Repository navigation
[None][feat] Muse Spark Open modeling_v2 TEP target: GB300 and Vera Rubin fast paths, and the DFlash 2 drafter transform on the sharded draft head - #21143
Conversation
…ests' caches With TRTLLM_KV_DEFER_CLOSE=1 the overlap executor arms deferred closing on a target-model KVCacheManagerV2. free_resources then releases a finished request's index slot at once and queues its cache; the close, which commits the last block and releases the pages, runs right after the next forward step is launched, before the executor waits on an empty queue, and at shutdown. Dummy requests always close at once. Off by default (TRTLLM_KV_DEFER_CLOSE unset or 0): the unset path is unchanged; the index-slot release is factored into _detach_index_slot. Tests: test_kv_cache_v2_deferred_close.py. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ssion diagnostics TRTLLM_SCHED_CTX_PAGE_BACKOFF_PASSES=N: once a pass that schedules generation sees a first context chunk refused for KV pages, the rest of that pass and the next N generation passes try no more first chunks; one pass in N + 1 probes again. Chunk continuations and passes without generation are never held. The state depends only on the scheduling decisions, so ranks stay in lockstep. TRTLLM_SCHED_ADMISSION_DIAG=1 (window TRTLLM_SCHED_ADMISSION_DIAG_PASSES) logs the host time of context admission per phase. Both are off by default (backoff 0, diagnostics unset): no candidate is held back and no scheduler method is wrapped. Tests: test_kv_cache_v2_scheduler.py. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…put processor MuseSparkOpenOmniInputProcessor declares supports_tokenization_cache and accepts enable_tokenization_cache. A media-free prompt is encoded through the shared prefix-tokenization cache, whose spliced ids equal what the processor's own encode would produce. Prompts with media, and processors built without a tokenizer, keep the existing path. Off by default: enabled only by the enable_tokenization_cache LLM argument. Tests: test_musespark_open_tokenization_cache.py. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…parallel layers MuseSparkOpenDecoderLayer gains _sp_attn_input_gather. When a model-specific path sets it on a layer (with a QKV projection that consumes its result), the sequence-parallel context step gathers the attention input through it instead of the NCCL all-gather. Off by default: the class attribute is None, and the layer is unchanged unless a path sets it. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…muted store and 64-deep MMA - The Rubin gather + grouped GEMM + SwiGLU kernel (FC1) and its custom op take trtllm_gen_gated_layout: the weights are TRTLLM-gen's gated FC1 tensors, alpha holds the gate scales followed by the output scales, and the epilogue follows TRTLLM-gen's SwiGLU and NVFP4 quantization arithmetic. - The Rubin grouped GEMM finalize kernel (FC2) and its custom op take store_permuted (one output row per permuted row, alpha * acc, no token scale and no scatter), trtllm_gen_n_shuffle (B and its block scales in TRTLLM-gen's shuffled row order; output columns in logical order) and mma_k64 (the 64-deep SM100-family MMA, whose fp32 accumulators equal those of SM100-family NVFP4 GEMMs). - Kernel and tuning-config caches are keyed on the new options; the tiled MMA construction is factored into _make_tiled_mma with the same arguments. All four options are constructor arguments, False by default; with every option False the kernels and ops build and run as before. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ion, per-SM prep and split-KV launches - mixed_step_gen_rows (Rubin, sm107 TEP4): on the full-attention layers, a mixed step with 9..64 generation requests runs its context rows on the thop op (context phase only) and its generation rows on multi_request_split_kv's decode stand-in, as decode steps do. Sliding-window layers, pure context steps and mixed steps outside the band keep the thop attention. - decode_kv_splits on mixed steps (GB300, sm103 TEP4): a mixed step with 9..64 generation requests arms the decode KV-split rule for the generation phase's trtllm-gen decode FMHA call, as decode steps of 72..512 rows already do. The split stand-in counts the calls it split and the calls it kept. - multi_request_prep_append: on sm103 a packed launch geometry per row count (items per program, warps, vector lane loads and stores) and an L2 prefetch of the next projection's weight slice; sm107 keeps the grid of one program per (row, head). Both write the same bytes. - multi_request_split_kv: the ring depth of KV pages a CTA holds in flight is chosen per SM version and number of Q terms (on sm103 a deeper ring for a one-term Q; sm107 keeps its depth). Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ics and Rubin CuTe DSL prefill experts - decode_expert_tactics: the stacked NVFP4 expert call of tep4 decode steps runs a fixed TRTLLM-gen [tile_n, config] tactic per (expert window, rows) from a per-GPU table instead of the autotuner's power-of-two bucket choice. GB300 (sm103): the 128-token tile for 256..504 rows; every pinned tactic computes the expert rows bitwise as the autotuned ones. Rubin (sm107): the 32-token tile for 256..384 rows, bitwise as the autotuned tactics with the same FC2 MMA K. Unlisted row counts and windows, tuning passes and configs not valid for the call's shape keep the autotuner; a GPU without a table is refused. - prefill_cutedsl_experts (Rubin, sm107 at tp 4 only, selected by serves(compute_capability, tp_size)): the unfinalized expert calls of 4096..8192 NVFP4 rows (prefill and mixed steps) of every NVFP4 MoE layer on the rank's contiguous expert block run moe_sort, the Rubin CuTe DSL FC1 (TRTLLM-gen gated layout) and FC2 (permuted store, TRTLLM-gen N shuffle, 64-deep MMA) on the served TRTLLM-gen tensors, and return the TRTLLM-gen call's triple to the model's per-expert post-norm finalize. Kernels are compiled at install with fixed tactics per row bucket; the install is agreed over the TP group. Rounding differs from TRTLLM-gen; expert selection is unchanged. Every other GPU, tp size and call keeps the TRTLLM-gen expert call. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ngine gather on sequence-parallel context steps The sequence-parallel context path wires a per-SM set of parts (SHARD_PARTS_BY_SM) and refuses any other SM: - GB300 (sm103): the MoE routes each rank's own rows with the fused routing kernel and gathers the top-8 ids and FP32 weights instead of the router logits; the attention-side residual join also writes the router's FP32 input (router_input); and on steps of at least MIN_ROWS rows the copy engines gather the attention input into the peers' receive buffers while the QKV projection runs in row blocks behind per-peer flags (ce_gather). The gather is agreed on by every rank and falls back to the NCCL all-gather when its buffers cannot be built or a projection cannot run in row blocks; its host flag words use prefer_pinned(). - Rubin (sm107): the MoE re-class only; the router logits are gathered, every row of the step is routed, and the attention input is all-gathered. Routing ids and weights, the router input and the gathered rows are the built-in path's, bit for bit. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
- Input: a decode step of at most replicated_embedding.MAX_ROWS ids runs the replicated embedding gather, embed_norm, the FP32 residual copy and layer 0's pre-QKV norm as one CuTe DSL kernel (one 128-thread CTA per row, flashinfer's RMSNorm body), and hands its two outputs to layer 0 as the residual stream and the precomputed attention input. PDL: wait before the first global read, trigger at exit. - Output: while armed for a decode forward, the last layer's post-FFN join stores BF16 directly, so the cast before the final norm is a no-op. - install checks the model it acts on (replicated table, m_emb == 1, the pre-QKV ordering, flashinfer's RMSNorm geometry, the last layer not a DFlash capture layer) and is agreed over the TP group. inputs_embeds calls, other row counts and prefill keep the built-in chain. Bitwise with the built-in chain. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
- vocab_parallel_head (GB300, sm103 TEP4): decode steps of VP_MIN_ROWS (two K=7 verify blocks) up to the head's widest rows sample and argmax on each rank's own vocabulary shard; the ranks exchange per-thread states, greedy words and candidates over the tray's multicast and agree on the token in the head_sampler JIT CUDA kernels, instead of all-gathering the logits. Greedy rows are the argmax of the soft-capped row exactly; sampled rows keep the shipped CDF order; filtered rows run the shipped body on the row gathered from the ranks. - target_head.install(llm, mapping, gather, vocab_parallel=False) builds the exchange buffers (rebuilt by widen for the wider steps); 8-row steps keep the lm_head_inbox gather. - Split-sampler routing for the vocab-parallel pending rows. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…er GEMM tactics and per-GPU PDL chain widths - wide_drafter_hops: the fused drafter all-reduce hops (one-shot all-reduce, finish convolution and add-RMSNorm per site) also on block decodes above 8 rows, up to the device's served rows (GB300: 9..56, Rubin: 9..64; the blocks of several requests at K = 7), on a second, wide one-shot workspace. Forwards of at most 8 rows keep the first. - Drafter GEMMs: the fixed low-M split-K tables at TP4 on sm103 and sm107 list the blocks of 2..4 requests (16, 24 and 32 rows) next to the one-request block. - drafter_pdl_chain: the conv / merge / silu entry triggers act on every block width of 1..4 requests at which the drafter's tactic table covers every successor; the merge and silu widths are per GPU (sm103: 1..4 requests for both; sm107: merge on one request, silu on one and four). The block decode reads the widths the chain's install recorded. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…hs on GB300 and Rubin - Dispatch registers mixed_step_gen_rows, model_boundary, wide_drafter_hops and vocab_parallel_head with their dependencies. - GB300 (sm103) TEP4 composes vocab_parallel_head and wide_drafter_hops, and decode_kv_splits installs its mixed-step mixin. - Rubin (sm107) TEP4 composes mixed_step_gen_rows and wide_drafter_hops; load_weights passes the device's compute capability, so the CuTe DSL prefill experts install on sm107 at tp 4 only. - Both GPUs, TEP4 and TEP8, compose model_boundary (bitwise). - DecodeTarget runs the model_boundary input kernel and arms its BF16 last-layer join; adopt_model installs the mixed-step attention mixins after every other attention re-class; the drafter shell installs the wide hops and passes vocab_parallel to the target head. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…103 TEP4 and the wide projection table by device - half_window_experts joins the GB300 (sm103) TEP4 composition, as at Rubin (sm107) TEP4: the NVFP4 expert layers also hold their ring predecessor's upper expert half. Decode steps of at most 8 rows are placed by the prologue's routed kernel, those of 9..32 rows (two to four requests) by ring_buddy_assign's half flow, both balancing the selected experts over the four ranks. The window costs 45 GiB of KV cache per rank, so its capacity gate carries it only up to max_batch_size 8, and a plan carrying it leaves hot_replicas out (dispatch.YIELDS). Compositions above max_batch_size 8 are unchanged. - wide_projection_gathers takes its 9..32-row kernel table from wide_kernel_steps on the dense GEMM plan's compute capability: the sm103 table, and a refusal on sm107 (whose TEP4 composition runs those projections on wide_dense_gemms' own table) and on any other device. No served composition changes: only sm103 TEP4 serves the path. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ransform after the sharded draft top-k The Muse Spark Open DFlash 2 drafter declares a logit transform (output multiplier and softcap), which the DFlash worker applies to the draft logits right after the drafter's logits processor and which returns fp32. The vocab-sharded draft head takes only bf16 shard logits, so with the transform on every draft step fell back to the built-in selector: a full-vocab fp32 all-gather, torch.topk over the full vocabulary and the full-vocab fill, scatter and argmax, and the path-walk tail and the drafter PDL chain were lost. The transform never decreases, so each position's top-k is the same before and after it, except where it maps distinct bf16 logits to the same fp32 value (tanh saturates), which torch.topk then orders by vocab id. The draft head now keeps the bf16 shard and applies the transform after its top-k: - kernels/draft_head/logit_transform.py: per drafter, a canonical-key table (bf16 radix key -> key of its tie group's representative) and a values table (bf16 radix key -> the drafter's own transform of that value), built on the device by DFlashForCausalLM.apply_draft_logit_transform itself and checked exhaustively over the 65536 bf16 values (fp32 result, NaN and signed zeros kept, never decreasing, groups contiguous, canonical values comparing as the transformed ones, layout independence on a host-seeded probe of the runtime shape). - draft_shard_topk / draft_rank_topk take the canonical table and its identity band (only keys outside the band load it); draft_topk_merge takes the values table for the unary values. Without the tables every kernel is the previous one. Each op's reference and is_valid follow the new arguments, and the ops declare cells for the transform-on configurations. - draft_head.install puts a deferring transform on the drafter instance: it hands back the very shard tensor ShardDraftLogitsProcessor just returned and records it; ShardedDraftHeadDFlashWorker._refine_block_logits keeps those logits bf16 only when the step takes the sharded selector (which then owes the transform and must apply it before sampling), and applies the drafter's transform to everything else, so sampled drafts, d2t, other layouts and any other tensor keep the built-in computation bit for bit. Deferral also stops when a class composed above overrides the refinement or the selector, or when the drafter's attributes no longer match the tables. selector_draft_tokens takes the transform as a required keyword; the verify-fold tail walk passes it too. - The install builds and checks everything before changing anything, runs the tie-order probe also on logits whose transformed values tie (from the transform's own tie groups, and at many requests), and agrees its refusal and a digest of the tables over the TP group. wide_draft_topk's probe runs with the transform as well. No change outside the target: an identity transform installs nothing new and runs the same kernels as before. Tests (tests/unittest/_torch/modeling_v2/models/musespark_open/...): the tables against the drafter's method on every bf16 value and in the runtime layout; the ops on their cells; the merge and the end-to-end draft tokens against torch.topk and the built-in selector over the transformed logits (1 to 24 requests, TP4 and TP8, all three key layouts, tie-heavy and saturating logit families); the fast path engaged with the transform on (no vocabulary all-gather, no full-vocab torch.topk), under the fold tail and in CUDA-graph replay; every fallback; install atomicity, idempotence and the probe's teeth; and a source pin on DFlashWorker._forward_impl's statements between the logits processor and _refine_block_logits. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…es follow the drafter, warn on built-in steps The vocab-sharded draft head applies the drafter's logit transform after its top-k from tables built at install from the drafter's logit_multiplier and logit_softcap. Only the deferring transform checked that the drafter still declares those attributes; ShardedDraftHeadDFlashWorker._refine_block_logits and the selector's _draft_head_transform_for took bf16 logits as owing the installed tables. A drafter whose attributes changed to the identity (1.0 / None) after the install hands bf16 logits through, and the sharded selector then ranked and scored them with the stale tables. - _refine_block_logits keeps a step's logits untransformed only when the drafter still matches the tables, and otherwise applies the drafter's class method with its current attributes, the built-in computation. - A new _draft_head_takes decides the sharded route for both selectors (ShardedDraftHeadDFlashWorker and the verify-fold tail walk): the existing shard conditions plus, with a transform installed, tables that still match the drafter. A step it turns down takes the built-in path with the drafter's current transform. _draft_head_transform_for takes the drafter and raises DraftTransformStateError on tables that do not hold, and _draft_head_transformed applies the drafter's current transform. - Every way a step with a transform installed takes the built-in path warns once per reason through tensorrt_llm.logger.warning_once with the "[musespark_open] draft_head:" prefix: the deferring transform declining (logits processor replaced, another tensor, non-bf16 logits, changed attributes, a worker class that overrides the refinement or selector), and the refinement or selector turning a step down (full-vocab logits, layout, sampled drafts, d2t, changed attributes, logits the transform did not defer). The deferring transform records its reason on the drafter so a step warns once, not once per stage. - The install agrees the key exchange's refusal over the TP group too: a rank whose make_allreduce returns None and the ranks that built theirs all raise the same ValueError, before anything changes. Tests (test_modeling_v2_draft_head_transform.py): the drafter's multiplier and softcap changed after the install (to the identity, another softcap, a multiplier only, a softcap only), before the logits processor, after the deferring transform and after the refinement, at 1 and 3 requests and under the fold tail: the built-in transform of this rank's shard reaches the vocabulary all-gather bit for bit, the sharded selector never runs, the draft tokens equal the built-in ones of the new transform, and the tables serve again once the drafter declares their transform; the selector refuses tables that do not hold; each fallback reason warns exactly once over repeated steps, the fast path and a drafter without a transform never warn; the key-exchange refusal is agreed and leaves the worker and drafter as they were. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
JunyiXu-nv
left a comment
There was a problem hiding this comment.
Runtime part: with the new env vars unset, the KV cache manager, scheduler and executor paths are unchanged, but the opt-in TRTLLM_KV_DEFER_CLOSE=1 path also defers preemption releases (repro in the inline comment) and can leave a cache open across a drained control action.
| if self._defer_close_armed and not request.is_dummy: | ||
| if request.py_request_id in self._early_freed_index_requests: | ||
| self._early_freed_index_requests.discard(request.py_request_id) | ||
| else: | ||
| self._detach_index_slot(kv_cache) | ||
| self.index_mapper.remove_sequence(request.py_request_id) | ||
| self._deferred_closes = (*self._deferred_closes, (request.py_request_id, kv_cache)) | ||
| return |
There was a problem hiding this comment.
With TRTLLM_KV_DEFER_CLOSE=1 this also defers preemption: preempt_request -> _release_preempted -> free_resources returns before close(), so its "pages were released" result is no longer true. _try_recompute_pause_for_gen (scheduler_v2.py L2437-2449) retries the allocation right after each victim, still fails, and takes the next one. Running the head's own scheduler and manager code against a stub pool, one pool-full generation step recompute-paused 3 victims instead of 1 and was still not scheduled. Could _release_preempted (and the admission rollbacks) close at once, so only finished requests are deferred?
| idle = (total_num_live_requests == 0 and len(waiting_queue) == 0 | ||
| and not self.is_shutdown | ||
| and not self._has_pending_connector_transfers()) | ||
| if idle and self._kv_defer_close_armed: |
There was a problem hiding this comment.
This flush never runs while a control request is parked: the control_pending return above (L6105-6108) comes first. _handle_control_request then fires a drained update_weights/reset_prefix_cache/sleep once active_requests is empty, while the last finished request's cache is still open; it only closes after the action, at L5651. reset_prefix_cache drains so that no request holds or commits blocks across the reset, and a request cancelled mid-prefill still commits its partial block in close(). Could _handle_control_request flush the deferred closes before firing?
Description
This PR extends the Muse Spark Open modeling_v2 TEP target on GB300 (sm103) and Vera Rubin (sm107) with a further set of decode, prefill and drafter fast paths, and makes the vocab-sharded draft head serve the DFlash 2 drafter's declared logit transform. It follows #20770 (GB300) and #20844 (Vera Rubin). The changes are 15 commits:
_experimental/modeling_v2/models/musespark_open/...__tep);The modeling_v2 target. These paths are always on inside the target and selected by GPU architecture and tensor-parallel size; each path owns its row band.
mixed_step_gen_rowson sm107, the decode KV-split rule on sm103); a packed per-row-count launch geometry for the multi-request decode prep on sm103; the split-KV page ring depth per SM version.max_batch_size≤ 8).DFlash 2 drafter logit transform on the sharded draft head (
25525dc68c,fc6d7acc84). The drafter declares an output multiplier and softcap, which the DFlash worker applies in fp32. The vocab-sharded draft head took only bf16 shard logits, so with the transform on every draft step fell back to the built-in selector (full-vocab fp32 all-gather andtorch.topk). The transform never decreases, so the draft head now keeps the bf16 shard and applies it after its top-k, from tables built at install by the drafter's own method and checked exhaustively over the 65536 bf16 values, ties included. The draft tokens equal the built-in selector's over the transformed logits; target verification and the accept rule are unchanged. Any step that has to take the built-in path warns once per reason, and the install's refusals are agreed over the TP group.Shared paths: opt-in, off by default. With the variables or arguments unset, behaviour is unchanged.
KVCacheManagerV2:TRTLLM_KV_DEFER_CLOSE=1defers closing a finished request's cache until after the next forward step is launched (target-model manager only).KVCacheV2Scheduler:TRTLLM_SCHED_CTX_PAGE_BACKOFF_PASSES=Nbacks off first-chunk context admission for N generation passes after the KV pool refuses one;TRTLLM_SCHED_ADMISSION_DIAG=1logs the admission host time.MuseSparkOpenOmniInputProcessorsupportsenable_tokenization_cachefor media-free prompts.MuseSparkOpenDecoderLayer._sp_attn_input_gather: a hook a model-specific path can set to replace the sequence-parallel attention-input all-gather.Measured gains. AgentX traces on the 1M-token corpus, full protocol, forced draft acceptance (3.05; GB300 TEP4 c1 3.09). Each row compares this PR's tree with the previous target code with the same drafter transform, run in the same wave. x = p10 per-user output throughput; y = total throughput per GPU.
max_batch_size3max_num_tokens16k (two runs)max_batch_size3max_batch_size8On Vera Rubin each value is the mean of two runs of this PR. Every run of the sharded draft head logged no built-in-path warning.
Accuracy. Spec decoding on with natural acceptance; GSM8K (1319 questions behind a ≥ 22k-character preamble) and GPQA-Diamond, greedy; pooled paired McNemar against the previous target code. Every probe passes (p ≥ 0.05): accuracy deltas from −1.0 to +1.0 pp at Vera Rubin c72 / c32 and GB300 c72 / c10. Natural acceptance length rises with the correct drafter transform, e.g. 4.09 → 4.51 on GSM8K at Vera Rubin c72.
Validation. The tree was validated as built (wheel for sm103 and sm107) on
feat/musespark-open-agentxataf9c7ccd94. It merges and rebases onto the current branch tip without conflicts; the only files both sides change arekv_cache_manager_v2.pyandscheduler_v2.py, in disjoint regions. The unified DFlash KV cache (#20806) is not reached undermodeling_v2, where DFlashAUTOstill resolves toVANILLA. The per-launch decode KV split (#20934) defers to this target's decode stand-ins on the calls they intercept.Test Coverage
tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py: deferred close.tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py: context page backoff and admission diagnostics.tests/unittest/_torch/modeling/test_musespark_open_tokenization_cache.py: the omni processor's prefix-tokenization cache.tests/unittest/_torch/modeling_v2/models/musespark_open/.../kernels/draft_head/and.../ops/spec/: the transform tables on every bf16 value, the draft top-k ops on their cells, end-to-end draft tokens againsttorch.topkand the built-in selector (1 to 24 requests, TP4 and TP8, tie-heavy and saturating logits), the fast path under the fold tail and CUDA-graph replay, every fallback and its warning, install atomicity, and a source pin onDFlashWorker._forward_impl.accuracy/modeling_v2/test_musespark_open.py(on this branch): GSM8K and mean acceptance length undermodeling_v2="require".PR Checklist
🤖 Generated with Claude Code