Repository navigation
Qwen3.6 MTP drafter on CUDA: draft a shared round's streams together - #1562
Merged
Merged
Conversation
The Qwen3.6 MTP drafter launches tf_q36m_attn once per stream, one ring each, so a shared round runs its streams' attention in series. tf_q36m_attn_rows reads each row's ring, slots and segment from a table and repeats the original's arithmetic, so the contract holds it byte-equal to one tf_q36m_attn launch per segment on mixed rings, with exact fp32 sums, an FP64 check and nvcc 13.0 launch records.
The CUDA MTP drafter drafted each stream apart: one-row passes, one head read and one sync per stream a step, so rounds grew with the streams. Shared rounds now absorb every stream's kept rows and draft each level for all streams in passes of up to 16 rows, with the same bits as alone; one stream and trees keep the serial path, and TF_MTP_BATCH=0 restores it for comparison. A chain's previous draft level lives beside its newest row's state, in a pool lane's run when the stream borrows one, and areaBytes counts it.
Reply hashes cannot show that rejected drafts match, so the MTP GPU check now drafts two sets of contexts on the same tap ring, serially and batched, and compares every held token, candidate, ring and state after absorb and after hold: mixed depths, picked rows, a gap, reused contexts and 17 streams.
The Qwen3.6 MTP drafter's attention folded one ring slot at a time, each slot waiting on its own load and a 64-bit modulo, so a warp took about 380 ns a slot and one stream's draft step spent 1.6 ms over a 4,200-slot ring. The three ring kernels now share a helper that loads eight slots a group with the next group in flight, steps the slot with a wrap, and folds the same terms in the same order, so every draft keeps its bits: the contract finds tf_q36m_attn_rows bit-equal to the old tf_q36m_attn, and on the checkpoint the trees equal their chains and batched drafts equal serial ones. The ring attention takes 0.24 ms, single-stream decode on a GB10 rises from 90.6 to 103.2 tok/s, and the nvcc 13.0 records follow the new register use, with 13.2 receipts pending.
The recipe's concurrent-request section described chains advancing a step at a time for all streams; it now describes the batched passes and TF_MTP_BATCH=0, and records one GB10's bench.py numbers on the MLX checkpoint with and without batching on the same binary. The CHANGELOG has the entry.
Owner
|
Thank you @laudney. This is in TensorFold 1.2.0, released today: https://github.com/ashhart/TensorFold/releases/tag/v1.2.0 It landed with your commits under your own name, so you're listed as a TensorFold contributor, and it's credited to you in the release notes and the changelog. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1560.
What
On CUDA, Qwen3.6-35B-A3B's MTP drafter drafted each stream of a shared round on its own: one MTP attention launch, one vocabulary-head read and one host synchronization per stream and draft level, so a round's time grew almost linearly with its streams. This PR:
tf_q36m_attn_rows(qwen36_mtp.cu): MTP attention for many streams' rows in one launch, each packed row over its own stream's ring, with the same arithmetic astf_q36m_attn.cuda_mtp_batch.zig): absorb passes over every stream's kept rows, then for each draft level passes over every stream still drafting, up to 16 rows a pass, with one head read and one synchronization a pass. Single-stream rounds and trees keep the serial path;TF_MTP_BATCH=0drafts one stream at a time for comparison. A chain's previous draft level lives beside its newest row's state, in a pool lane's run when the stream borrows one (areaBytescounts it).tf_q36m_attn,tf_q36m_attn_rowsandtf_q36m_attn_treefold ring slots in groups of eight with the next group's loads in flight and a wrapped slot index, in the same order.No draft changes: every stream drafts exactly what it drafts alone, and replies cannot change because the target verifies every draft.
Gate (on one DGX Spark: GB10, sm_121, CUDA 13.0)
tf-cuda-test mtp-rows, also registered inkernel-contract):tf_q36m_attn_rowsis byte-equal totf_q36m_attnlaunched once per stream at 5 served shapes (three rings, one wrapping, NaN in every slot outside a ring's range); NaN rule, five byte-identical repeats with guards, exact fp32 sums at 4,200 terms where fp16/bf16 sums miss by 124x/127x; FP64 at the served shapes. Launch records measured with nvcc 13.0 for sm 80/86/89/120/121 (later_pendingfor 13.2). The look-ahead loop stays byte-equal to the old one-slot loop: it was checked against the unchangedtf_q36m_attnbefore that kernel took it too.tf-qwen35-cuda <checkpoint> mtp <ids> 0onTensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP, revision f84b054): batched drafting equals serial byte for byte (held tokens, candidates, rings and states) for six streams with absorbs of 1 to 7 rows and depths 0 to 3, partial acceptance with a gap past a stream's end, contexts reused from the pool, and 17 streams (two passes a level); the existing chunked-absorb and tree-against-chain checks pass."draft": falseruns (6 of 6, greedy and two seeded temperatures); concurrent replies equal solo at 1 and 4 streams (8 of 8 cells).zig build testandzig build test-cuda-hostpass.Numbers
bench.pyof spark-sglang-docker, 4.2K-token prompt, 512 tokens, thinking on, the same binary withTF_MTP_BATCH=0and batched (both with the look-ahead attention):TF_MTP_BATCH=0The look-ahead attention alone takes a 4,200-slot ring from 1.6 ms to 0.24 ms a draft step (contract timing).
atc overlapagainstmain: no conflicts this PR adds; the registrations incontract_tests.zigandmain.zigsit apart from #1388's, #1426's and #1528's.