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
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.
… flush before control actions With TRTLLM_KV_DEFER_CLOSE=1, free_resources deferred the close of every non-dummy request while the overlap loop was armed, including preemption and admission rollbacks, whose callers count on the pages being free when the call returns. The executor also fired a drained control request before flushing, so the action could run while a finished request's cache was still open. - KVCacheManagerV2.deferring_closes(): free_resources defers a close only while armed and inside this scope; every other release closes at once. - PyExecutor._do_terminate_request releases inside the scope; preemption, recompute pauses and scheduler rollbacks keep closing at once. - PyExecutor._handle_control_request flushes the deferred closes before it fires the action. Off by default: with TRTLLM_KV_DEFER_CLOSE unset nothing changes. Tests: test_kv_cache_v2_deferred_close.py. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Feature-branch increment; nothing here blocks, but a few items should be on record before this lands in main via the branch merge.
Tracking. A 10k-line feature set across 15 commits carries [None]. Please attach the JIRA for the AgentX / modeling_v2 TEP work (or the per-path tickets) in the title.
Env vars are undocumented. TRTLLM_KV_DEFER_CLOSE, TRTLLM_SCHED_CTX_PAGE_BACKOFF_PASSES, TRTLLM_SCHED_ADMISSION_DIAG and TRTLLM_SCHED_ADMISSION_DIAG_PASSES (the last one is not in the PR description either) appear only in docstrings. Add them where the other TRTLLM_* knobs live under docs/source/ so a deployment can find them; at minimum list all four in the description.
Test coverage of the new target kernels. In-repo tests cover the draft-head transform, the three top-k ops, deferred close, the scheduler backoff and the tokenization cache. Nothing covers head_sampler.cu (772 lines; per-shard sampling that must agree across ranks), ce_gather, router_input, model_boundary, the mixed-step KV-split arming, the drafter hop/GEMM tables, or the four new CuTe DSL layout options. The description cites in-engine checksum arms, but those are not in the PR. For the bitwise paths a shape-pinned parity test against the built-in computation is cheap and would guard the cuBLAS/Triton/CuTe version bumps these claims depend on. Fine as a follow-up on this branch, but please file it.
Numerics disclosure. The sm107 TP4 prefill experts on the Rubin CuTe DSL GEMMs are the one path that is not bitwise with TRTLLM-gen. Since the target is always-on, please state in modeling.py's path table which paths are certified bitwise and which are not, so an accuracy regression on sm107 has a first suspect.
Interaction of the two shared opt-ins. With deferred close armed, the overlap loop schedules before it flushes, so pages of requests that finished in iteration i are invisible to the scheduler in iteration i+1. A page refusal caused only by those pages arms an N-pass backoff when TRTLLM_SCHED_CTX_PAGE_BACKOFF_PASSES is set. Both knobs are aimed at the same deployment; see the inline note on the scheduler.
The CuTe DSL unique_id and kernel_cache keys do include the new constructor options; checked.
| """A ``[rows, 256]`` FP32 stand-in for the step's logits, memoized to ``(selected, weights)``. | ||
|
|
||
| The stand-in is a zero-stride view of one element: it carries no values and is never | ||
| read, because every routing call on it returns the memo. |
There was a problem hiding this comment.
The placeholder is new_empty expanded over [rows, 256]: uninitialized memory whose only protection is that every consumer calls routing_method.apply with this exact object (the memo matches on identity and version). Any path that replaces or slices the tensor first misses the memo, and apply then runs its fallback top-k on garbage and returns plausible ids and weights with no error. Two such paths exist today in the shared MoE: the perfect-router substitution in interface.py (returns a different tensor) and the multi-chunk split (the SP branch refuses when rows > moe_max_num_tokens, which keeps chunking out, but nothing ties the two conditions together). Cheap guard: fill the single element with NaN so a memo miss produces NaN routing weights and fails loudly, and have shard_routing return None when experts._enable_perfect_router is set.
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
There was a problem hiding this comment.
The bitwise claim depends on cuBLAS choosing the same algorithm for a [shard, 8192] @ [8192, N] block as for the whole-step [rows, 8192] GEMM. cuBLAS heuristics key on M, and a split-K or different tile choice for the smaller M changes the fp32 accumulation order, so a cuBLAS bump can break this without any code change. Nothing in the repo checks it. Please add a GPU unit test that runs project vs one torch.matmul over the gathered rows at the deployed shapes (4096..max rows, TP4 shard), or do a one-time self-check in build and refuse (agreed over the group) when they differ.
| reader = threading.Thread( | ||
| target=lambda: result.append(self.flag_words(bound_s)), | ||
| name="musespark_open.ce_gather_flags", | ||
| daemon=True, |
There was a problem hiding this comment.
os.abort() from a daemon thread after a 120 s stall takes the whole rank down with a core dump, and the 120 s bound only exempts the very first forward. A sequence-parallel step later in the serve can legitimately exceed it: a first-seen shape triggering a Triton or flashinfer JIT on one rank, a peer paused by a debugger or profiler, or a peer held in a weight-update / sleep control action while this rank is mid-forward. A stall here is genuinely unrecoverable, but consider (a) making STALL_S tunable, (b) logging and raising on the next host call (set a flag forward checks) instead of aborting, so the executor's own hang detector and error reporting get the message, and (c) at least logger.error with flush before abort so the report is not lost in buffered output.
| """Advance the backoff by one scheduling pass.""" | ||
| if not has_generation: | ||
| return | ||
| if refused_this_pass: |
There was a problem hiding this comment.
With TRTLLM_KV_DEFER_CLOSE=1 the overlap loop runs schedule, then launch, then flush, so the pages of requests that finished in iteration i are still allocated when iteration i+1's scheduler runs. A first-chunk refusal that only those pages caused arms passes generation passes of hold here, even though the pages are free a few milliseconds later. Both knobs ship in this PR for the same deployment. Either count the deferred caches' pages as free in the admission check (the manager knows them), or flush deferred closes before scheduling when the backoff is armed, or document that the two are not meant to be combined and say which measurement showed it does not matter.
| kv_cache.discard_pending_stats() | ||
| # Dummy requests reuse their ids, which the stats-exclusion bookkeeping | ||
| # done after close() is keyed on, so they always close at once. | ||
| if self._defer_close_armed and self._defer_close_scope and not request.is_dummy: |
There was a problem hiding this comment.
free_resources defers only inside deferring_closes(), which _do_terminate_request enters, and dummy requests close at once because their ids recur. One more case worth a test: a request terminated while _early_freed_index_requests holds it (slot released early) now takes the deferred branch and only discards the id; the close that commits its last block still runs at flush. That is what you want, but the pairing of early-free with deferred close is not in test_kv_cache_v2_deferred_close.py (test_early_released_slot_is_not_released_twice checks the slot, not the block commit).
| "refinement or selector" | ||
| ) | ||
| else: | ||
| draft_model.__dict__[DEFERRED_ATTR] = (logits.data_ptr(), logits.numel()) |
There was a problem hiding this comment.
Recording (data_ptr, numel) and re-deriving it after the built-in reshape is correct today because reshape of a contiguous tensor is a view. If the built-in DFlash forward ever inserts a .contiguous(), .clone() or a slice between the transform call and _refine_block_logits, the step silently takes the built-in path with a one-time warning rather than failing. That is the intended degradation, but the warning is keyed per reason and emitted once per process, so a deployment that regresses to the slow path on every step sees one log line. Consider also counting built-in-path steps (as DecodeKvSplits does with split_calls/kept_calls) and exposing the count, so a perf regression is diagnosable after the fact.
| next_input, | ||
| next_norm.weight, | ||
| next_norm.variance_epsilon, | ||
| router_input, |
There was a problem hiding this comment.
join_with_router_input allocates router_input = torch.empty_like(residual) (an extra [shard, 8192] fp32 per layer per step) and hands it to the MoE through sp_router_input. _sp_router_logits consumes it only when handed[0] is hidden_states; if the layer passes a view or the MoE takes the super().forward branch, the tuple is dropped and the allocation was wasted, with no signal. Worth a logger.info_once on the first drop (as mixed.py does for its outcomes) so a silent miss of the hand-off is visible.
| class MixedStepKvSplitsAttention: | ||
| """Mixin re-classed in front of a full-attention layer's attention class. | ||
|
|
||
| A call on a step ``takes_generation`` accepts runs the class behind it with the model's |
There was a problem hiding this comment.
takes_generation reads num_contexts and num_generations with getattr(..., 0) defaults. On a metadata object that lacks either attribute (a different attention backend's metadata reaching this layer), the rule is silently never armed, and the info_once further down never fires because the armed() block is skipped. Prefer failing at install time by checking the metadata class the model actually uses, or at least a one-time warning when the attributes are missing, so a backend change does not quietly drop the path.
|
/bot run --disable-fail-fast |
|
PR_Github #81517 [ run ] triggered by Bot. Commit: |
|
PR_Github #81517 [ run ] completed with state |
c55183f
into
NVIDIA:feat/musespark-open-agentx
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