[None][feat] MuseSpark GB300: port the decode kernel stack, switches become dispatch levers

tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py

@@ -116,6 +116,12 @@ class FastPath:
         "decode FMHA of steps with two or more requests as balanced page-split partials",
         "construct",
     ),
+    FastPath(
+        "multi_request_combine_gate",
+        "the multi-request decode FMHA combine storing the attention output's norm-gate",
+        "construct",
+        ("multi_request_split_kv",),
+    ),
     FastPath(
         "multi_request_prep_append",
         "every layer's K/V append folded into the decode prep of steps with two or more requests",
@@ -129,6 +135,12 @@ class FastPath:
     ),
     # Dense projections and prefetch.
     FastPath("dense_gemms", "dense projections on one named kernel per shape", "construct"),
+    FastPath(
+        "tep8_batch_splitk",
+        "TP8 decode projections of 33..160 rows on batched split-K kernels",
+        "construct",
+        ("dense_gemms",),
+    ),
     FastPath(
         "dense_gemm_prefetch",
         "L2 prefetch of the next dense projection's weight",
@@ -195,6 +207,13 @@ class FastPath:
         "construct",
         ("projection_gathers", "dense_gemms"),
     ),
+    FastPath(
+        "tep8_wide_direct_transports",
+        "pool and unpool slices published from their split-K GEMMs into 32-row multicast "
+        "transports, 16/24/32-row decode steps",
+        "construct",
+        ("wide_projection_gathers",),
+    ),
     FastPath(
         "projection_publish",
         "sharded pool/unpool GEMMs publishing into multicast transports their consumers poll, "
@@ -209,22 +228,58 @@ class FastPath:
         "construct",
         ("projection_gathers", "post_attention_hop"),
     ),
+    FastPath(
+        "tep8_router_shard",
+        "the TP8 MoE prologue's router shard and fused unpool gather on 16..32-row decode steps",
+        "construct",
+        ("moe_prologue",),
+    ),
+    FastPath(
+        "tep8_router_shard_wide",
+        "that router shard on the built-in wide bits instead of each row's one-request bits",
+        "construct",
+        ("tep8_router_shard",),
+    ),
     FastPath(
         "pool_into_route",
         "pool projection and router shard writing into the routed kernel's transport",
         "construct",
         ("moe_prologue", "dense_gemms"),
     ),
+    FastPath(
+        "tep8_wide_direct_pool",
+        "the pool GEMM publishing into the TP8 wide route transport, 16..32-row decode steps",
+        "load",
+        ("pool_into_route", "tep8_router_shard"),
+    ),
     FastPath(
         "finalize_transport",
         "expert finalize all-reduce on the fused hop's transport",
         "construct",
     ),
+    FastPath(
+        "finalize_entry_trigger",
+        "the fused finalize releasing its dependent at its own dependency wait",
+        "construct",
+        ("finalize_transport",),
+    ),
+    FastPath(
+        "tep8_finalize_transport_wide",
+        "the TP8 finalize transport laid out for calls of 17..40 tokens",
+        "construct",
+        ("finalize_transport",),
+    ),
     FastPath(
         "router_gemm",
         "router logits of 9..512-row calls as a BF16x3 split-K tensor-core GEMM",
         "construct",
     ),
+    FastPath(
+        "router_staged_gemm",
+        "router logits of 129..256-row calls on the staged BF16x3 GEMM",
+        "construct",
+        ("router_gemm",),
+    ),
     FastPath(
         "wide_moe_allreduce",
         "MoE output all-reduce of 129..1024-row decode steps on the MNNVL two-shot",
@@ -249,9 +304,52 @@ class FastPath:
         ("preattn_prep",),
     ),
     FastPath("replicated_experts", "ring-buddy expert replication, balanced", "load"),
+    FastPath(
+        "tep8_ring_buddy_wide",
+        "the balanced tep8 stacked path on decode calls above its 16-row cap",
+        "load",
+        ("replicated_experts",),
+    ),
+    FastPath(
+        "tep8_wide_balance",
+        "the NVFP4 layers' balanced placement above the stacked cap, on the wide assign program",
+        "load",
+        ("replicated_experts",),
+    ),
+    FastPath(
+        "ring_buddy_stats",
+        "a placement census per call, reported on an interval by TP rank 0",
+        "load",
+        ("replicated_experts",),
+    ),
     # persistent_experts and bf16_routed_experts run over the ring-buddy stack with
     # replicated_experts, else over the rank's own expert block (the owner-only window).
     FastPath("persistent_experts", "NVFP4 experts as one persistent launch", "load"),
+    FastPath(
+        "tep8_persistent_batch_experts",
+        "stacked NVFP4 expert calls of 9..32 rows on the persistent kernel's batch builds",
+        "load",
+        ("persistent_experts", "replicated_experts"),
+    ),
+    FastPath(
+        "tep8_experts_in_kernel_placement",
+        "the batch builds placing the step's selection themselves, with no assign launch",
+        "load",
+        ("tep8_persistent_batch_experts",),
+    ),
+    FastPath(
+        "tep8_experts_interleaved_fetch",
+        "the batch builds' two fetch warps moving the FC2 operand's k-blocks alternately",
+        "load",
+        ("tep8_persistent_batch_experts",),
+    ),
+    FastPath(
+        "tep8_experts_early_start",
+        "the expert assignment triggering the persistent kernel's batch builds at its own "
+        "dependency wait",
+        "load",
+        ("tep8_persistent_batch_experts",),
+    ),
     FastPath(
         "half_window_experts",
         "the ring predecessor's upper expert half on each TP4 rank, balanced per step",
@@ -305,9 +403,60 @@ class FastPath:
         ("draft_head",),
     ),
     FastPath("drafter_fc_shard", "drafter fc projection column-sharded", "draft"),
+    FastPath(
+        "tep8_fc_shard_wide",
+        "the sharded drafter fc projection on TP8 calls of 9..128 rows",
+        "draft",
+        ("drafter_fc_shard",),
+    ),
     FastPath("fused_accept", "acceptance and forced acceptance in one launch", "draft"),
     FastPath("drafter_fixed_gemms", "drafter GEMMs on fixed low-M tactics", "draft"),
+    FastPath(
+        "drafter_conv_epilogue",
+        "the drafter coefficient GEMM of 9..32-row blocks with the input-side convolution in "
+        "its epilogue",
+        "draft",
+        ("drafter_block_conv", "drafter_fixed_gemms"),
+    ),
+    FastPath(
+        "drafter_finish_norm",
+        "each drafter hop's finish convolution, residual add and RMSNorm as one launch",
+        "draft",
+        ("drafter_block_decode",),
+    ),
+    FastPath(
+        "drafter_silu_epilogue",
+        "the drafter gate_up projection of 9..32-row blocks with silu(gate) * up in its epilogue",
+        "draft",
+        ("drafter_block_decode",),
+    ),
+    FastPath(
+        "swa_block_single_early_loads",
+        "the one-request sliding-window block's split and combine issuing their operand loads "
+        "at once",
+        "construct",
+        ("swa_block",),
+    ),
+    FastPath(
+        "tep8_attn_block_splitk",
+        "the TP8 attention block's QKV and o_proj of 9..32-row decode calls on batch split-K "
+        "cells, each staging its weight tile before its dependency wait",
+        "construct",
+        ("dense_gemms",),
+    ),
+    FastPath(
+        "drafter_multi_request_cells",
+        "the drafter GEMMs of 16, 24 and 32-row blocks on measured TP8 split-K cells",
+        "draft",
+        ("drafter_fixed_gemms",),
+    ),
     FastPath("context_kv_fork", "drafter context-K/V store on a side stream", "draft"),
+    FastPath(
+        "drafter_fused_prep",
+        "the drafter's page-table refresh, context K/V store and commit as fused launches",
+        "draft",
+        ("context_kv_fork",),
+    ),
     FastPath(
         "drafter_window_compact",
         "windowed drafter arena compaction in place, one read and write per window",

tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py

@@ -22,7 +22,14 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None:
     if block_conv or block_decode:
         from .kernels.drafter_block_decode import block_decode as _draft
 
-        _draft.install(shell.llm.draft_model, block_conv, block_decode)
+        _draft.install(
+            shell.llm.draft_model,
+            block_conv,
+            block_decode,
+            conv_epilogue=plan.enabled("drafter_conv_epilogue"),
+            finish_norm=plan.enabled("drafter_finish_norm"),
+            silu_epilogue=plan.enabled("drafter_silu_epilogue"),
+        )
     if plan.enabled("wide_drafter_window_attention"):
         # On the block decode drafter_block_decode just bound.
         from .kernels.wide_drafter_window_attention import attention as _wide_window
@@ -66,6 +73,8 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None:
             shell.llm.model,
             declined=plan.declined,
             paged_commit=plan.enabled("paged_context_commit"),
+            fused_prep=plan.enabled("drafter_fused_prep"),
+            fused_prep_min_requests=plan.param("drafter_fused_prep", "MIN_REQUESTS"),
         )
     verify_fold = plan.enabled("verify_fold")
     path_walk_tail = plan.enabled("path_walk_tail")

tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py

@@ -90,6 +90,7 @@
 from tensorrt_llm._utils import CUASSERT
 from tensorrt_llm.logger import logger
 
+from ...ops.spec import drafter_step_prep as _step_prep
 from ...ops.spec.context_kv_commit import MAX_REQUESTS, context_kv_commit
 from ...ops.spec.context_kv_commit_paged import context_kv_commit_paged
 from ...ops.spec.context_kv_commit_windowed import context_kv_commit_windowed
@@ -196,6 +197,9 @@ class Launched:
     point: str
     #: The captured rows the side chain projected (a view of the shared capture buffer).
     gen_hs: torch.Tensor
+    #: Whether the side chain and the commit run the fused bookkeeping
+    #: (``drafter_fused_prep``, ``ops/spec/drafter_step_prep.py``).
+    fused: bool = False
 
 
 class ContextKvFork:
@@ -207,7 +211,13 @@ class ContextKvFork:
     ``prepare_1st_drafter_inputs`` takes it, joins and commits.
     """
 
-    def __init__(self, worker: DFlashWorker, draft_model: nn.Module, point: str):
+    def __init__(
+        self,
+        worker: DFlashWorker,
+        draft_model: nn.Module,
+        point: str,
+        prep: Optional["_step_prep.DrafterStepPrep"] = None,
+    ):
         self.worker = worker
         self.draft_model = draft_model
         self.point = point
@@ -218,6 +228,8 @@ def __init__(self, worker: DFlashWorker, draft_model: nn.Module, point: str):
         self._armed: Optional[tuple] = None
         self._launched: Optional[Launched] = None
         self._fallback_logged = False
+        #: The fused decode-step bookkeeping (``drafter_fused_prep``), or None.
+        self.prep = prep
         #: ``paged_context_commit``: a paged context's commit runs ``context_kv_commit_paged``.
         self.paged_commit_on = False
         #: Reports a step a lever left to the path below it (``dispatch.Plan.declined``), or None.
@@ -264,10 +276,11 @@ def fire(self, point: str) -> None:
                 "the side chain cannot be forked from it"
             )
         capture = current_capture_id()
+        fused = self.prep is not None and self.prep.takes(self.worker, num_gens, k + 1)
         self.stream.wait_stream(main)
         with torch.cuda.stream(self.stream):
-            gen_hs = self.side_chain(spec_metadata, attn_metadata, num_gens, k)
-        self._launched = Launched(main, capture, num_gens, k, point, gen_hs)
+            gen_hs = self.side_chain(spec_metadata, attn_metadata, num_gens, k, fused)
+        self._launched = Launched(main, capture, num_gens, k, point, gen_hs, fused)
         self.launches[point] += 1
 
     # -- the worker side ---------------------------------------------------
@@ -314,7 +327,12 @@ def commit(self, launched: Launched, slots, num_accepted) -> Optional[torch.Tens
         base, last = arena_columns(worker)
         self.stream.wait_stream(main)
         with torch.cuda.stream(self.stream):
-            if worker._ctx_paged and self._commits_paged(worker, slots, num_accepted):
+            if worker._ctx_paged and launched.fused:
+                # drafter_fused_prep: the same context_kv_commit_paged op
+                # #21535's paged_context_commit branch below calls, reached through the
+                # fused bookkeeping's own gate (drafter_step_prep.takes).
+                self.prep.commit_context(worker, slots, num_accepted, rows, num_ctx)
+            elif worker._ctx_paged and self._commits_paged(worker, slots, num_accepted):
                 context_kv_commit_paged(
                     worker._ctx_len,
                     worker._ctx_base,
@@ -450,20 +468,32 @@ def plan(self, spec_metadata, attn_metadata) -> Optional[tuple[int, int]]:
             return None
         return num_gens, k
 
-    def side_chain(self, spec_metadata, attn_metadata, num_gens: int, k: int) -> torch.Tensor:
+    def side_chain(
+        self, spec_metadata, attn_metadata, num_gens: int, k: int, fused: bool = False
+    ) -> torch.Tensor:
         """The built-in store of the step's context K/V, unmasked; runs on the side stream.
 
         The worker's ``_store_step_context`` without its mask, after this step's page
         table decode over the KV cache manager's pages. Its first launch is a torch
         kernel, so the fork is a plain dependency on the main stream's last kernel.
+        With ``fused`` (``drafter_fused_prep``) the decode and the store
+        are ``drafter_step_prep``'s launches, the same bytes.
         """
         worker = self.worker
-        DFlashWorker._refresh_ctx_block_tables(worker, attn_metadata, num_gens)
+        if fused:
+            self.prep.refresh_block_tables(worker, attn_metadata, num_gens)
+        else:
+            DFlashWorker._refresh_ctx_block_tables(worker, attn_metadata, num_gens)
         slots = worker._batch_to_slot[:num_gens]
         # The built-in's get_hidden_states(num_tokens)[0 : rows]: the same view.
         gen_hs = spec_metadata.captured_hidden_states[
             : num_gens * (k + 1), : spec_metadata.hidden_size * spec_metadata.num_capture_layers
         ]
+        if fused:
+            self.prep.store_step_context(
+                worker, self.draft_model, gen_hs, slots, self.rows[:num_gens], k + 1
+            )
+            return gen_hs
         worker._store_step_context(
             self.draft_model, gen_hs, slots, self.rows[:num_gens], worker._ctx_len[slots]
         )
@@ -771,6 +801,8 @@ def install(
     late_join: bool = False,
     declined: Optional[Callable[..., None]] = None,
     paged_commit: bool = False,
+    fused_prep: bool = False,
+    fused_prep_min_requests: int = _step_prep.DEFAULT_MIN_REQUESTS,
 ) -> None:
     """Install the fork on the DFlash worker, the target core and the LM head (``POINT``).
 
@@ -803,7 +835,12 @@ def install(
     fork: Optional[ContextKvFork] = worker.__dict__.get("_kv_fork")
     fresh = fork is None
     if fresh:
-        fork = ContextKvFork(worker, drafter, point)
+        fork = ContextKvFork(
+            worker,
+            drafter,
+            point,
+            _step_prep.from_plan(fused_prep, fused_prep_min_requests),
+        )
         worker._kv_fork = fork
     fork.draft_model = drafter
     fork.declined = declined