From 1eb0935c2f7af6c8eee453b315591c2e4ac627f1 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian Date: Mon, 9 Nov 2026 17:47:42 -0800 Subject: [PATCH 1/9] [None][feat] MuseSpark GB300: port the decode kernel stack, switches become dispatch levers Ports the MuseSpark GB300 decode kernel stack onto feat/musespark-open-agentx and replaces the environment switches it carried with dispatch levers: staged router GEMM, ring-buddy cap and wide balance, router shard, drafter fast path, wide direct transports and pool, persistent experts, experts early start and interleaved fetch, in-kernel placement, finalize entry trigger, the SWA cluster block, the widened TP8 finalize transport, batch split-K, and the multi-instance context FMHA. EVERY LEVER IS DEFAULT OFF. This is the port; it changes no behaviour until a lever is knobbed on. Proved rather than asserted: a verifier rebuilds FAST_PATHS, DEPENDS_ON and both compositions by AST and resolves them with an EMPTY environment -- the resolved path sets are identical to the base at all ten deployments (44 paths at TEP8, 54/53 at TEP4), and each lever is then shown to turn on by its own knob. THREE DEFECTS FIXED, each found by running the code after static checks passed: 1. weights._ring_buddy_row_cap called max_rows_for with two arguments against a one-argument signature, on the unconditional replicated_experts install path. Every TP8 deployment would have raised TypeError at load_weights. 2. moe_prologue declared _wide_direct: bool as a class attribute, shadowing projection_gathers' transport object (default None) through the MRO. Every 8-row decode step raised AttributeError("'bool' object has no attribute 'is_armed'") at warmup, inside attention_output_into_hop, which is served by default at TP8. 3. pool_into_route called _wide_router_shard with four arguments against five. Latent, behind a default-off lever. THIRD-PARTY LICENSING. multi_instance_context_kernel.py derives from the CUTLASS CuTe DSL Blackwell FMHA example (v4.8.0) and declares that derivation with the full BSD-3-Clause notice in-file, matching 51 existing files in this repository and the sm107 sibling added in 736cfdc7c00. CUTLASS is already attributed in ATTRIBUTIONS-Python.md and both per-arch CPP files, so no attribution file changes. MEASURED. With every lever off this branch is indistinguishable from its base on GB300 TEP8 decode (11.67 vs 11.93 ms at bs3, 12.26 vs 12.28 at bs4), which is a correctness result -- the port does not disturb the default path -- and says nothing about what the levers are worth. Non-bitwise levers still need their acceptance, drift and step-time gates before any of them ships on by default. Signed-off-by: Iman Tabrizian --- .../dispatch.py | 110 ++ .../drafter.py | 2 + .../kernels/context_kv_fork/fork.py | 51 +- .../kernels/dense_gemms/dispatcher.py | 270 ++++ .../kernels/drafter_hops/hops.py | 4 + .../kernels/finalize_transport/transport.py | 12 +- .../kernels/finalize_two_shot/install.py | 100 +- .../kernels/moe_prologue/prologue.py | 196 ++- .../multi_request_split_kv/attention.py | 87 +- .../kernels/pool_into_route/route.py | 192 ++- .../kernels/projection_gathers/gathers.py | 15 + .../kernels/projection_publish/publish.py | 22 +- .../kernels/replicated_experts/ring_buddy.py | 621 +++++++- .../kernels/router_gemm/router.py | 21 +- .../wide_direct_transports/__init__.py | 3 + .../wide_direct_transports/transports.py | 482 ++++++ .../kernels/wide_prologue/__init__.py | 3 + .../kernels/wide_prologue/prologue.py | 309 ++++ .../layers.py | 31 + .../lever_settings.py | 173 +++ .../levers.py | 195 ++- .../modeling.py | 40 +- .../ops/allreduce/drafter_hop.py | 36 + .../ops/attention/multi_request_split_kv.py | 203 ++- .../ops/attention/swa_block.py | 139 +- .../ops/attention/swa_cluster_block.py | 1071 ++++++++++++++ .../ops/comm/pool_direct_route.py | 24 +- .../ops/comm/pool_gather_route.py | 5 +- .../ops/comm/projection_publish.py | 11 +- .../ops/comm/unpool_direct_boundary.py | 20 +- .../ops/gemm/batch_splitk_gemm.py | 805 ++++++++++ .../ops/moe/_persistent_experts/kernel.py | 677 +++++++-- .../ops/moe/finalize_allreduce_transport.py | 37 +- .../ops/moe/finalize_two_shot.py | 68 +- .../ops/moe/persistent_experts.py | 245 ++- .../ops/moe/ring_buddy_assign.py | 261 +++- .../ops/moe/router_gemm.py | 96 +- .../ops/moe/router_gemm_staged.py | 428 ++++++ .../ops/moe/router_shard.py | 525 +++++++ .../ops/spec/drafter_step_prep.py | 504 +++++++ .../spec_worker.py | 46 +- .../weights.py | 138 +- .../backends/fmha/flashinfer_trtllm_gen.py | 64 +- .../backends/fmha/multi_instance_context.py | 817 ++++++++++ .../fmha/multi_instance_context_kernel.py | 1312 +++++++++++++++++ 45 files changed, 10180 insertions(+), 291 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/transports.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_cluster_block.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/batch_splitk_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_gemm_staged.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_step_prep.py create mode 100644 tensorrt_llm/_torch/attention/backends/fmha/multi_instance_context.py create mode 100644 tensorrt_llm/_torch/attention/backends/fmha/multi_instance_context_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py index 09b1a146311f..716fd90febef 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py +++ b/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,21 @@ 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("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", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py index 618ea48cc5c8..2d7201adb531 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py @@ -66,6 +66,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") diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py index 5eb226f2ae24..52574478b569 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py +++ b/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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py index 02fc11282317..2cc4c7ad452d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py @@ -27,6 +27,10 @@ that acts on larger decode steps arms the plan instead (``DenseGemmPlan.batch``, set for one decode forward by ``batch_dense_gemms``): while it is set, a re-classed projection's ``apply_linear`` hands the call to it with the built-in GEMM as a callable. +With ``tep8_batch_splitk`` (``batch_splitk``, sm103 TEP8 only) and no such +path armed, ``apply_linear`` runs a 33..160-row call of a site ``SM103_TP8_BATCH_KERNELS`` +names on ``ops/gemm/batch_splitk_gemm`` instead of the built-in cuBLAS call; the pooler and +unpooler slices of those rows run it through ``projection_publish`` (``batch_splitk_apply``). With ``dense_gemm_prefetch`` a projection that ran also forks a side stream that asks the L2 to fetch the next projection's weight (QKV -> the same layer's attention output -> @@ -146,6 +150,144 @@ } +#: ``batch_splitk``: the switch (read once, at ``install``). Off by default; on sm103 at TP8 with +#: the sharded pooler and unpooler, ``1`` runs the 33..160-row decode calls (5..20 requests at +#: K = 7) of the sites ``SM103_TP8_BATCH_KERNELS`` names on ``ops/gemm/batch_splitk_gemm`` (a +#: cluster split-K kernel whose kernel-N tile holds up to 256 rows, reducing in DSMEM) instead +#: of the built-in cuBLAS call (a split-K ``nvjet`` kernel plus ``cublasLt::splitKreduce`` for +#: most of these shapes and rows). Anywhere else it changes nothing but a warning. +#: The row band of ``batch_splitk``: above the 9..32-row wide band, up to the 20-request step. +BATCH_MIN_ROWS = WIDE_MAX_ROWS + 1 +BATCH_MAX_ROWS = 160 +#: ``tep8_batch_splitk``: whether ``install`` builds that table. Off keeps the built-in calls. +_BATCH_SPLITK = False + + +def configure(*, batch_splitk: bool = False) -> None: + """Set whether ``install`` arms the batched split-K table. + + ``lever_settings.apply`` sets it from the model's plan before any install; the default + keeps the built-in cuBLAS calls at every row count above the wide band. + + Args: + batch_splitk: the ``tep8_batch_splitk`` lever (``_BATCH_SPLITK``). + """ + global _BATCH_SPLITK + _BATCH_SPLITK = bool(batch_splitk) + + +#: ``batch_splitk`` on sm103 at TP8 (per-rank QKV ``[2304, 8192]``, pooler slice ``[512, 8192]``, +#: unpooler slice ``[1024, 4096]``), as ``(last row count, kernel)`` steps from 33 rows; rows past +#: a site's last step keep the built-in call. ``bsplitk:,,,``: +#: ``mma_n`` public rows per CTA (a call takes ``ceil(rows / mma_n)`` row tiles). Every cell is +#: split_k 4, the split_k of each site's 8-row cell in ``DEFAULT_KERNELS`` and of the 9..32-row +#: kernels: K is summed in the same four slices in the same order, so every row gets the bits it +#: gets in a one-request step (``install`` checks the split_k). Chosen in 62-layer chains of cold +#: weights against the served cuBLAS call on one GB300 (microbench ``mb-h58-bsk-2``, each cell +#: the fastest split_k-4 tactic launched without PDL; checked through the overlay in +#: ``mb-h58-bsk-verify``): QKV one row tile (144 CTAs, the most clusters of 4 that fit) up to 128 +#: rows, -0.1..+2.2 us against cuBLAS without PDL and +1.1..+3.4 us with it (from 136 rows +#: cuBLAS's two-CTA tiles tie or win); the pooler slice four row tiles (128 CTAs) and the +#: unpooler slice two (128 CTAs), +0.2..+2.6 us without PDL at every row count. The attention +#: output projection is not named: its split_k 1 cells (cuBLAS's bits) are slower than cuBLAS +#: without PDL at every row count. +SM103_TP8_BATCH_KERNELS: Mapping[str, Tuple[Tuple[int, str], ...]] = { + "qkv": ( + (40, "bsplitk:64,40,4,8"), + (48, "bsplitk:64,48,4,7"), + (56, "bsplitk:64,56,4,7"), + (64, "bsplitk:64,64,4,6"), + (72, "bsplitk:64,72,4,6"), + (80, "bsplitk:64,80,4,5"), + (88, "bsplitk:64,88,4,5"), + (96, "bsplitk:64,96,4,5"), + (104, "bsplitk:64,104,4,4"), + (112, "bsplitk:64,112,4,4"), + (120, "bsplitk:64,120,4,4"), + (128, "bsplitk:64,128,4,4"), + ), + "pool": ( + (48, "bsplitk:64,24,4,10"), + (64, "bsplitk:64,16,4,11"), + (96, "bsplitk:64,24,4,10"), + (128, "bsplitk:64,32,4,9"), + (160, "bsplitk:64,40,4,8"), + ), + "unpool": ( + (48, "bsplitk:64,24,4,8"), + (64, "bsplitk:64,32,4,8"), + (80, "bsplitk:64,40,4,8"), + (96, "bsplitk:64,48,4,7"), + (112, "bsplitk:64,56,4,7"), + (128, "bsplitk:64,64,4,6"), + (144, "bsplitk:64,72,4,6"), + (160, "bsplitk:64,80,4,5"), + ), +} + +#: The ``batch_splitk`` table of each ``(compute capability, TP size)`` it is measured for. +BATCH_KERNELS_BY_SM_TP: Mapping[ + Tuple[Tuple[int, int], int], Mapping[str, Tuple[Tuple[int, str], ...]] +] = {((10, 3), 8): SM103_TP8_BATCH_KERNELS} + + +def batch_kernels( + table: Mapping[str, Tuple[Tuple[int, str], ...]] = SM103_TP8_BATCH_KERNELS, +) -> Dict[str, Tuple[Optional[str], ...]]: + """The kernel of each site of a ``batch_splitk`` table per row count ``BATCH_MIN_ROWS.. + BATCH_MAX_ROWS`` (index ``rows - BATCH_MIN_ROWS``); ``None`` past the site's last step.""" + out = {} + for site, steps in table.items(): + kernels: List[Optional[str]] = [] + first = BATCH_MIN_ROWS + for last, kernel in steps: + assert first <= last <= BATCH_MAX_ROWS, steps + kernels += [kernel] * (last - first + 1) + first = last + 1 + kernels += [None] * (BATCH_MAX_ROWS - first + 1) + out[site] = tuple(kernels) + return out + + +def batch_kernel_table( + capability: Tuple[int, int], + tp_size: int, + sharded: bool, + kernels: Mapping[str, Tuple[str, ...]], +) -> Dict[str, Tuple[Optional[str], ...]]: + """``batch_splitk``'s kernel per row count 33..160 of each site its table names on ``capability`` + at ``tp_size``, checked against ``kernels`` (the plan's 1..8-row table): every cell must be a + split-K tactic with the ``split_k`` of its site's 8-row cell. Reads only the tables. + + Raises: + ValueError: if ``BATCH_KERNELS_BY_SM_TP`` has no table for ``(capability, tp_size)``, the + projections are not ``sharded``, or a cell does not sum the K slices of its site's + 8-row cell. + """ + from ...ops.gemm import batch_splitk_gemm as _bsk + + key = (tuple(capability), int(tp_size)) + steps = BATCH_KERNELS_BY_SM_TP.get(key) + if steps is None or not sharded: + raise ValueError( + f"[musespark_open] batch_splitk is measured for SM {(10, 3)} at TP8 with the sharded " + f"pooler and unpooler; got SM {key[0]} at TP{key[1]}, " + f"{'sharded' if sharded else 'replicated'} projections" + ) + table = batch_kernels(steps) + for role, cells in table.items(): + family, fields = parse_kernel(kernels[role][MAX_ROWS - 1]) + for kernel in cells: + if kernel is None: + continue + if family != SPLITK or _bsk.parse(kernel).split_k != fields[2]: + raise ValueError( + f"[musespark_open] batch_splitk: {role} kernel {kernel} does not sum the K " + f"slices of the {MAX_ROWS}-row kernel {kernels[role][MAX_ROWS - 1]}" + ) + return table + + def _by_rows(steps: Tuple[Tuple[int, str], ...]) -> List[str]: kernels, first = [], 1 for last, kernel in steps: @@ -374,6 +516,10 @@ def __init__( #: ``BatchDenseGemms``), whose ``projection`` every re-classed projection's #: ``apply_linear`` calls; ``None`` otherwise. self.batch = None + #: ``batch_splitk``: the kernel of a site's 33..160-row calls + #: (``batch_splitk[site][rows - BATCH_MIN_ROWS]``, ``None`` past its last step); empty + #: when off (``set_batch_splitk_kernels``). + self.batch_splitk: Dict[str, Tuple[Optional[str], ...]] = {} self.prefetch = prefetch self.enable_pdl = enable_pdl #: The compute capability of the device the plan serves (``install`` reads it once); @@ -410,6 +556,28 @@ def kernel(self, role: str, rows: int) -> Optional[str]: """``kernel_for``: the kernel of a ``role`` call of ``rows`` rows, or ``None``.""" return self.kernel_for(role, rows) + def batch_kernel_for(self, role: str, rows: int) -> Optional[str]: + """``batch_splitk``'s kernel of a 33..160-row ``role`` call, or ``None`` (the built-in call).""" + cells = self.batch_splitk.get(role) + if cells is None or not BATCH_MIN_ROWS <= rows <= BATCH_MAX_ROWS: + return None + return cells[rows - BATCH_MIN_ROWS] + + def set_batch_splitk_kernels(self, kernels: Mapping[str, Tuple[Optional[str], ...]]) -> None: + """Give sites a ``batch_splitk`` kernel per row count 33..160 (``batch_kernel_table``). + + Raises: + ValueError: for an unknown site or a table that does not cover 33..160 rows. + """ + rows = BATCH_MAX_ROWS - BATCH_MIN_ROWS + 1 + for role, table in kernels.items(): + if role not in SITES or len(table) != rows: + raise ValueError( + f"[musespark_open] dense_gemms: a batch_splitk table names {rows} kernels per " + f"site of {SITES}, got {role!r} with {len(table)}" + ) + self.batch_splitk = {role: tuple(table) for role, table in kernels.items()} + def set_wide_kernels(self, kernels: Mapping[str, Tuple[str, ...]]) -> None: """Give sites a kernel per row count 9..32 (``kernels[site][rows - WIDE_MIN_ROWS]``). @@ -687,6 +855,58 @@ def _apply_substitute(self, input_tensor: torch.Tensor, weight: torch.Tensor) -> self.site.plan.fork_prefetch(self.site.next_weight()) return product.view(*input_tensor.shape[:-1], int(weight.shape[0])) + def batch_splitk_apply( + self, + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + out: Optional[torch.Tensor] = None, + ) -> Optional[torch.Tensor]: + """A 33..160-row call on the plan's ``batch_splitk`` kernel for its rows, or ``None`` (the caller + runs the built-in call): BF16, no bias, the site's shape, contiguous, 32-byte aligned, not + compiled or differentiated. ``out`` (``projection_publish``'s transport slot) receives the + product; the kernel stores each element of it once, after the reduction. Joins the plan's + forks the site's dispatcher would have joined before its GEMM.""" + plan = self.site.plan + if not plan.batch_splitk or bias is not None: + return None + if torch.is_grad_enabled() or torch.compiler.is_compiling(): + return None + if not ( + input_tensor.is_cuda + and input_tensor.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and weight.dim() == 2 + and input_tensor.dim() >= 1 + and (int(weight.shape[1]), int(weight.shape[0])) == self.site.shape + and int(input_tensor.shape[-1]) == int(weight.shape[1]) + and input_tensor.is_contiguous() + and weight.is_contiguous() + and input_tensor.data_ptr() % 32 == 0 + and weight.data_ptr() % 32 == 0 + ): + return None + k = int(weight.shape[1]) + rows = input_tensor.numel() // k + kernel = plan.batch_kernel_for(self.site.role, rows) + if kernel is None: + return None + from ...ops.gemm import batch_splitk_gemm as _bsk + + plan.join() + if plan.named_forks: + plan.join_role(self.site.role) + x = input_tensor.view(-1, k) + if out is None: + out = torch.empty((rows, int(weight.shape[0])), dtype=x.dtype, device=x.device) + _bsk.run(x, weight, out, _bsk.parse(kernel), plan.enable_pdl) + logger.info_once( + f"[musespark_open] batch_splitk: a {rows}-row {self.site.role} call (layer " + f"{self.site.layer_idx}) ran {kernel}", + key=f"musespark_open.batch_splitk.{self.site.role}", + ) + return out.view(*input_tensor.shape[:-1], int(weight.shape[0])) + def _apply_publish( self, module: torch.nn.Module, @@ -775,6 +995,10 @@ def apply_linear(self, input, bias, lora_params=None, layer_idx=None): batch = dispatcher.site.plan.batch lora = getattr(self, "lora", None) is not None and bool(lora_params) if batch is None or lora: + if dispatcher.site.plan.batch_splitk and not lora and _builtin_is_cublas(self): + out = dispatcher.batch_splitk_apply(input, self.weight, bias) + if out is not None: + return out return Linear.apply_linear(self, input, bias, lora_params, layer_idx) return batch.projection( dispatcher.site, @@ -812,6 +1036,18 @@ def adopt( return linear +def _builtin_is_cublas(linear: Linear) -> bool: + """Whether the built-in call of ``linear`` above 32 rows is ``F.linear`` on its BF16 weight + (``UnquantizedLinearMethod.apply``): no custom cuBLAS GEMM, no CuTe DSL BF16 GEMM.""" + from tensorrt_llm._torch.modules.linear import UnquantizedLinearMethod + + return ( + isinstance(getattr(linear, "quant_method", None), UnquantizedLinearMethod) + and not getattr(linear, "use_custom_cublas_mm", False) + and not getattr(linear, "use_cute_dsl_bf16_gemm", False) + ) + + def _projection_weight(moe, linear: Linear, sharded: bool) -> Callable[[], torch.Tensor]: """The weight a pooler / unpooler GEMM is handed: the rank's row slice when sharded.""" if not sharded: @@ -865,8 +1101,20 @@ def install( "(MUSESPARK_OPEN_SHARDED_PROJECTIONS=1); without them the pooler GEMM reads the " "whole replicated weight" ) + batch_table: Mapping[str, Tuple[Optional[str], ...]] = {} + if _BATCH_SPLITK: + if (tuple(capability), tp_size) in BATCH_KERNELS_BY_SM_TP and sharded: + batch_table = batch_kernel_table(capability, tp_size, sharded, kernels) + else: + logger.warning( + "[musespark_open] tep8_batch_splitk is measured for SM (10, 3) at " + f"TP8 with the sharded pooler and unpooler; SM {tuple(capability)} at TP{tp_size} " + f"({'sharded' if sharded else 'replicated'}) keeps the built-in calls" + ) plan = DenseGemmPlan(kernels, prefetch=prefetch, enable_pdl=LAUNCH_PDL, capability=capability) plan.tp_size = tp_size + if batch_table: + plan.set_batch_splitk_kernels(batch_table) layers = list(model.layers) planned = [] @@ -899,6 +1147,18 @@ def install( raise ValueError( f"{where}: kernel {kernel} cannot serve {rows} x {shape}: {error}" ) + if index == 0 and role in batch_table: + from ...ops.gemm import batch_splitk_gemm as _bsk + + for rows, kernel in enumerate(batch_table[role], start=BATCH_MIN_ROWS): + if kernel is None: + continue + try: + _bsk.validate(_bsk.parse(kernel), rows, shape[1], shape[0]) + except ValueError as error: + raise ValueError( + f"{where}: kernel {kernel} cannot serve {rows} x {shape}: {error}" + ) from error planned.append((index, role, modules[role], shape, weights)) next_of = {"qkv": "o_proj", "o_proj": "pool"} @@ -920,4 +1180,14 @@ def install( f"{MAX_ROWS} rows by {describe(kernels)}; " f"L2 prefetch of the next weight {'on' if prefetch else 'off'}" ) + if batch_table: + table = BATCH_KERNELS_BY_SM_TP[(tuple(capability), tp_size)] + steps = "; ".join( + f"{site}=" + ", ".join(f"{k}@..{last}" for last, k in table[site]) for site in table + ) + logger.info( + f"[musespark_open] tep8_batch_splitk: TP8 calls of {BATCH_MIN_ROWS}.." + f"{BATCH_MAX_ROWS} rows by {steps} (split_k 4, each row the bits of its {MAX_ROWS}-row " + "kernel); o_proj and the rows past each site's last step keep the built-in cuBLAS call" + ) return plan diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py index d6970d7f8eb1..5db2d365a3c0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py @@ -395,6 +395,10 @@ def install_wide(draft_model: Optional[nn.Module], mapping=None) -> None: norms += [layer.input_layernorm for layer in layers[1:]] + [draft_model.model.norm] eps = sorted({float(norm.variance_epsilon) for norm in norms}) rows = _hop.wide_served_rows() + if narrow.ranks == _hop.MAX_RANKS: + # TP8 has no WIDE_SERVED_ROWS_BY_SM entry: the lever's TP8_MAX_ROWS names the bound + # (drafter_hop.configure, from lever_settings). + rows = _hop.tp8_wide_rows() call.wide = _hop.DrafterHopWorkspace( layers[0].self_attn.o_proj.mapping, sites=narrow.sites, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_transport/transport.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_transport/transport.py index e8eeb233955f..caa2841f0ff0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_transport/transport.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_transport/transport.py @@ -36,6 +36,16 @@ ATTRIBUTE = "_finalize_ar_workspace" +#: ``finalize_entry_trigger`` (off by default) builds the transport with the ``entry`` trigger +#: (``finalize_allreduce_transport.ENTRY_TRIGGER``), which lets the dependent grid launch right +#: after this kernel's own ``griddepcontrol.wait`` instead of after its output store. The +#: dependent (the unpool slice GEMM) then issues its weight loads while this kernel waits for its +#: peers; it still reads the finalize's output only after its own wait, which returns once this +#: grid has completed. Same statements, same bits; only the dependent's launch point moves. Off: +#: the built-in place. The lever reaches this module through +#: ``finalize_allreduce_transport.configure``, which rebuilds ``_fin.DEFAULT``. + + def refusal(model: MuseSparkOpenModel, wide: bool = False) -> Optional[str]: """Why ``install`` cannot replace this model's fused finalize, or None. Changes nothing.""" legacy = getattr(model, ATTRIBUTE, None) @@ -86,11 +96,11 @@ def install(model: MuseSparkOpenModel, wide: bool = False): RuntimeError: if the workspace could not be built on some rank (on every rank). """ + transport = _fin.DEFAULT reason = agreed_refusal(model, wide) if reason is not None: raise ValueError(reason) legacy = getattr(model, ATTRIBUTE) - transport = _fin.DEFAULT workspace = _fin.FinalizeTransportWorkspace(legacy, legacy._mapping, transport) for layer in model.layers: setattr(layer.feed_forward, ATTRIBUTE, workspace) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_two_shot/install.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_two_shot/install.py index 469fb19f4a2b..7d0feaf26957 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_two_shot/install.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_two_shot/install.py @@ -28,17 +28,30 @@ ``MUSESPARK_OPEN_FUSED_FINALIZE_AR_CHUNKED`` set. The switch is host state, a constant per CUDA-graph capture, and every rank arms it alike (the choice depends on the step's kind and row count only). + +At TP8 (sm103, two trays of four in one NVLink domain: ``dispatch.TWO_TRAYS_OF_4``) the +workspace's receive memory is the trays' fabric, and the band is ``ROWS_BY_SM_TP[(103, 8)]``: +there the built-in fused finalize is a one-shot over seven peers, ingress-bound from about +three requests, and the unfused combine + cast + MNNVL two-shot above 128 rows is three +launches. Measured per call in 62-layer chains on two GB300 trays (``mb-tep8-finalize-1``), +built-in -> this kernel: 8.9 -> 8.4 us at 24 rows, 15.0 -> 8.9 at 80, 19.9 -> 9.8 at 128 +(the one-shot); 13.4 -> 10.0 at 136..152 (the unfused path). At 160 rows (more tokens than +SMs) the two-register-bound build ties the unfused path, so the band ends at 152. Bitwise the +built-in fused kernel at 24..128 rows on every rank in that run; above 128 rows the served +unfused path rounds the per-expert norm's sum of squares in another order (within one BF16 +place of this kernel), so arming 129..152 rows changes bits there. """ from __future__ import annotations import contextlib -from typing import Iterator, List, Optional +from typing import Iterator, List, Optional, Tuple import torch from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenModel +from ... import dispatch as _dispatch from ...ops.moe import finalize_two_shot as _two #: The attribute the built-in MoE reads its fused finalize workspace from, and the @@ -52,19 +65,73 @@ MIN_ROWS_BY_SM = {103: 72} MIN_ROWS = MIN_ROWS_BY_SM[103] MAX_ROWS = _two.MAX_ROWS +#: The band, per (SM, TP size): ``(fewest rows, most rows)`` of an armed decode forward. TP4 +#: keeps ``MIN_ROWS_BY_SM`` up to ``MAX_ROWS``. TP8 (sm103 only, two trays): from three +#: requests (24 rows at K = 7; the kernel takes 17 and up, ``_two.MIN_ROWS_BY_RANKS``) to 19 +#: (152 rows, the device's SM count of tokens: see the module docstring). +ROWS_BY_SM_TP = { + (103, 4): (MIN_ROWS_BY_SM[103], MAX_ROWS), + (103, 8): (24, 152), +} +#: ``finalize_two_shot__TP8_MAX_ROWS``: a lower last row count of the TP8 band; None keeps +#: ``ROWS_BY_SM_TP[(103, 8)]``'s. ``128`` keeps the band to the rows whose served kernel is +#: the built-in fused finalize, where this kernel returns its bits, so no step's values change. +_TP8_MAX_ROWS: Optional[int] = None + + +def configure(*, tep8_max_rows: Optional[int] = None) -> None: + """Set the last row of the TP8 band ``band`` reports. + + ``lever_settings.apply`` sets it from the model's plan before any install; None, the + default, keeps the band ``ROWS_BY_SM_TP`` declares. + + Args: + tep8_max_rows: the last armed row count at TP8 (``_TP8_MAX_ROWS``), inside + ``ROWS_BY_SM_TP[(103, 8)]``. + + Raises: + ValueError: for a bound outside that band. + """ + global _TP8_MAX_ROWS + if tep8_max_rows is None: + _TP8_MAX_ROWS = None + return + first, last = ROWS_BY_SM_TP[(103, 8)] + rows = int(tep8_max_rows) + if not first <= rows <= last: + raise ValueError( + f"[musespark_open] finalize_two_shot TP8 row bound {rows}: expected {first}..{last}" + ) + _TP8_MAX_ROWS = rows + + +def band(sm: int, tp_size: int) -> Tuple[int, int]: + """``(fewest rows, most rows)`` armed for ``(sm, tp_size)``; at TP8 ``configure`` may lower + the last.""" + first, last = ROWS_BY_SM_TP[(sm, tp_size)] + if tp_size == 8 and _TP8_MAX_ROWS is not None: + last = min(last, _TP8_MAX_ROWS) + return first, last class FinalizeTwoShotSwitch: """What ``DecodeTarget`` arms: the workspace's ``active`` flag and the layers' chunk bound.""" - def __init__(self, workspace, moes: List[torch.nn.Module], min_rows: int = MIN_ROWS) -> None: + def __init__( + self, + workspace, + moes: List[torch.nn.Module], + min_rows: int = MIN_ROWS, + max_rows: int = MAX_ROWS, + ) -> None: self.workspace = workspace self.moes = moes self.min_rows = int(min_rows) + self.max_rows = int(max_rows) def takes(self, rows: int) -> bool: """Whether a decode forward of ``rows`` rows is armed.""" - return self.min_rows <= rows <= MAX_ROWS + return self.min_rows <= rows <= self.max_rows @contextlib.contextmanager def armed(self) -> Iterator[None]: @@ -89,8 +156,20 @@ def refusal(model: MuseSparkOpenModel) -> Optional[str]: if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): return "it is written for sm_103" mapping = model.layers[0].mapping - if mapping.tp_size not in _two.RANK_COUNTS or mapping.is_multi_node(): - return f"it serves TP4 on one node, got TP{mapping.tp_size}" + one_tray = mapping.tp_size == 4 and not mapping.is_multi_node() + two_trays = mapping.tp_size == 8 and _dispatch.TWO_TRAYS_OF_4.holds(mapping) + if mapping.tp_size not in _two.RANK_COUNTS or not (one_tray or two_trays): + return ( + f"it serves TP4 on one node or TP8 over {_dispatch.TWO_TRAYS_OF_4.description}, got " + f"TP{mapping.tp_size} ({_dispatch.describe_topology(mapping)})" + ) + major, minor = torch.cuda.get_device_capability() + if (10 * major + minor, mapping.tp_size) not in ROWS_BY_SM_TP: + return f"it has no band for TP{mapping.tp_size} on sm_{major}{minor}" + try: + band(10 * major + minor, mapping.tp_size) + except ValueError as error: + return str(error) if mapping.enable_attention_dp or mapping.moe_tp_size != 1: return "the MoE output is not all-reduced over the TP group (attention DP or MoE TP)" workspace = getattr(model.layers[0].feed_forward, ATTRIBUTE, None) @@ -147,6 +226,11 @@ def install(model: MuseSparkOpenModel) -> FinalizeTwoShotSwitch: moes.append(layer.feed_forward) bound_unarmed_calls(wrapped, moes) major, minor = torch.cuda.get_device_capability(workspace.device) - min_rows = MIN_ROWS_BY_SM[10 * major + minor] - _two.log_once_enabled(len(model.layers), workspace.max_rows, min_rows) - return FinalizeTwoShotSwitch(workspace, moes, min_rows) + min_rows, max_rows = band(10 * major + minor, mapping.tp_size) + # A finalize transport widened past 16 tokens (TP8: + # ``finalize_allreduce_transport.TP8_MAX_TOKENS_ENV``) is faster than this kernel up to + # its last token count, so the band starts above it. + min_rows = max(min_rows, int(getattr(wrapped, "max_tokens", 0)) + 1) + max_rows = min(max_rows, workspace.max_rows) + _two.log_once_enabled(len(model.layers), max_rows, min_rows) + return FinalizeTwoShotSwitch(workspace, moes, min_rows, max_rows) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/moe_prologue/prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/moe_prologue/prologue.py index bf09594a1104..66c37c941f78 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/moe_prologue/prologue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/moe_prologue/prologue.py @@ -90,13 +90,61 @@ from ...ops.comm.projection_inbox import MAX_ROWS, ProjectionInbox from ...ops.moe import router_shard as _shard from ..context_kv_fork.fork import dedicated_stream -from ..projection_gathers.gathers import MuseSparkOpenTepGatherMoE, _require_aligned +from ..projection_gathers.gathers import MuseSparkOpenTepGatherMoE, _require_aligned, fused_rows #: The built-in code the routed prologue reproduces: the router kernel and its #: launch choice (``_router_projection_kernel``, ``MuseSparkOpenMoE.router_logits``), #: the routing kernel (``_full_routing_kernel``) and the memo the experts and the #: finalize read (``MuseSparkOpenRoutingMethod.apply``); its ``reference`` runs them. +#: ``_report_tep8``'s (rows, reason) pairs already logged in this process. +_TEP8_REPORTED = set() + + +def _report_tep8( + world: int, + rows: int, + layer: int, + declined: Optional[str], + router_rows: bool = False, + direct: bool = False, +) -> None: + """Log, once per process and row count (and reason), what a TP8 wide prologue call did. + + TP8 only (``wide_prologue.install_tep8``, ``tep8_router_shard``): the first + served call at a row count names the router shard and the routed pool kernel; a declined + call names why it keeps the built-in prepare (the BF16-expert layers decline by design). + TP4 logs nothing, as before. + """ + if world != 8 or (rows, declined, direct) in _TEP8_REPORTED: + return + _TEP8_REPORTED.add((rows, declined, direct)) + if declined is None and direct: + router = "router_shard_rows (one-request bits)" if router_rows else "router_shard_wide" + logger.info( + f"[musespark_open] tep8 router shard: a captured {rows}-row decode step runs {router} into " + "the wide direct pool transport's logit block, the pool slice GEMM into its slot, and " + "pool_direct_route's wide routing (tep8_wide_direct_pool; first at layer " + f"{layer}; logged once per row count)" + ) + elif declined is None: + router = "router_shard_rows (one-request bits)" if router_rows else "router_shard_wide" + logger.info( + f"[musespark_open] tep8 router shard: a {rows}-row decode step runs {router} on the " + "router fork and the logit exchange + routing in pool_gather_route on the TP8 wide " + f"inbox (first at layer {layer}; logged once per row count)" + ) + elif declined.startswith("no fused"): + logger.info( + f"[musespark_open] tep8 router shard: a {rows}-row call of layer {layer} keeps the " + f"built-in prepare ({declined}; expected on the BF16-expert layers)" + ) + else: + logger.warning( + f"[musespark_open] tep8 router shard: a {rows}-row call of layer {layer} keeps the " + f"built-in prepare ({declined})" + ) + def _placement(moe, world: Optional[int] = None): """``(kind, static_from, state)`` of the replicated experts' placement on this layer, if installed and active. @@ -131,8 +179,35 @@ class MuseSparkOpenTepPrologueMoE(MuseSparkOpenTepGatherMoE): #: ``pool_direct.DirectPoolSite`` (its producers publish into the routed kernel's #: transport); None otherwise. _pool_direct_armed = None + #: ``wide_prologue``: the TP4 wide inbox (``WIDE_MAX_ROWS`` rows) of the 16..32-row + #: prologue, or None (those steps keep the built-in prepare). At TP8 the inbox + #: ``wide_prologue.install_tep8`` builds (``tep8_router_shard``). + _wide_inbox: Optional[ProjectionInbox] = None + #: ``wide_prologue``: whether the 16..32-row router shard gives each row its + #: one-request bits (``router_shard_rows``) instead of the built-in tensor-core + #: router's (``router_shard_wide``); the ``wide_router_per_row`` lever at TP4, + #: ``tep8_router_shard`` without ``tep8_router_shard_wide`` under ``install_tep8``. + _wide_router_rows: bool = False + #: ``wide_prologue.install_tep8`` with ``TRTLLM_MUSESPARK_TEP8_ROUTER_SHARD_EARLY``: on the + #: router fork's side stream, ``router_shard_rows`` is a programmatic launch (its weight tile + #: before its dependency wait; ``"entry"`` also releases the routed pool kernel at entry) + #: instead of a plain one. None keeps the plain launch. + _wide_router_early: Optional[str] = None + #: ``wide_prologue``: whether ``pool_into_route`` and ``unpool_into_boundary`` also + #: build their 16..32-row direct transports; the ``wide_direct_transports`` lever. + #: + #: NOT ``_wide_direct``: this mixin is re-classed in front of the MoE, and + #: projection_gathers.MuseSparkOpenTepGatherMoE uses ``_wide_direct`` for the wide direct + #: transport OBJECT (default None), read as ``.is_armed``. A bool here shadows that None + #: through the MRO, so ``wide_direct is not None`` passed and every 8-row decode step + #: raised AttributeError("'bool' object has no attribute 'is_armed'") at warmup. + _wide_direct_built: bool = False def prepare_moe_inputs(self, hidden_states: torch.Tensor): + if self._wide_inbox is not None: + wide = self._wide_prepare(hidden_states) + if wide is not None: + return wide warps = self._prologue_router_warps(hidden_states) if warps is None: self._pool_direct_refuse_declined(hidden_states) @@ -214,6 +289,125 @@ def prepare_moe_inputs(self, hidden_states: torch.Tensor): self._prologue_record_routing(out, place, state, world) return out.logits, Fp4QuantizedTensor(out.fp4, out.scale, is_sf_swizzled=False) + def _wide_router_shard( + self, hidden_states: torch.Tensor, rank: int, warps: int, out: torch.Tensor, world: int + ) -> None: + """This rank's router logit columns of a 16..32-row call into ``out`` (``[rows, L]`` FP32). + + By default ``router_shard_wide``: the built-in 16..32-row router's bits (its TF32x3 + tensor-core split-K and reduce). With ``_wide_router_rows``, ``router_shard_rows`` at + ``warps``: each row's logits are the built-in small-batch router's for that row in a + one-request call, one launch and no reduce, but other bits than the built-in + 16..32-row launch's. With ``_wide_router_early`` and on the router fork's stream, the + launch is programmatic (``router_shard_rows(enable_pdl=True)``): the same bits, the + weight tile streamed while the producer of ``hidden_states`` finishes. + """ + if self._wide_router_rows: + early = self._wide_router_early + side = self._router_stream + if early is not None and ( + side is None or torch.cuda.current_stream().cuda_stream != side.cuda_stream + ): + early = None + _shard.router_shard_rows( + hidden_states, + self.router.weight, + rank, + self.metap_multiplier, + warps, + out, + ranks=world, + enable_pdl=early is not None, + early_trigger=early == "entry", + ) + else: + _shard.router_shard_wide( + hidden_states, self.router.weight, rank, self.metap_multiplier, out, ranks=world + ) + + def _wide_prepare(self, hidden_states: torch.Tensor): + """``wide_prologue``: the prologue of a 16..32-row decode step, or None (nothing launched). + + Where the built-in decode would run its tensor-core router and the fused routing on + a pool-deferring call of 16..32 rows: this rank's router columns + (``_wide_router_shard``) on the router fork's side stream, the pool slice GEMM on the + current one, and ``pool_gather_route`` on the + wide inbox -- the pool gather + quantization, the logit exchange and the routing + (no placement: above 8 rows the experts are the built-in ones and read only the + routing memo). The pool slice is this rank's columns of the built-in replicated + pool projection, by a GEMM of its own shape. + """ + rows = hidden_states.shape[0] if hidden_states.dim() == 2 else 0 + if not _shard.WIDE_MIN_ROWS <= rows <= _shard.WIDE_MAX_ROWS: + return None + if self._prologue_router_conditions(hidden_states, _shard.WIDE_MAX_ROWS) is None: + # pool_into_route's wide site (tep8_wide_direct_pool) armed this layer: + # its peers poll the transport for its record, so a decline raises instead. + self._pool_direct_refuse_declined(hidden_states) + return None + inbox = self._wide_inbox + world = inbox.world + # The pool slice reaches the routed kernel only where the pool projection defers its + # gather at this row count: always under wide_prologue's TP4 install (_defer_max_rows + # 32); at TP8 (install_tep8) while wide_projection_gathers is armed (fused_rows 32). + if rows > max(self._defer_max_rows, fused_rows(self)): + self._pool_direct_refuse_declined(hidden_states) + _report_tep8(world, rows, self.layer_idx, "the pool slice is not deferred") + return None + columns = _route.local_columns(world) + probe = hidden_states.view(-1)[: rows * columns].view(rows, columns) + args = self._pool_quant_args(probe, max_rows=inbox.max_rows) + if args is None: + self._pool_direct_refuse_declined(hidden_states) + _report_tep8(world, rows, self.layer_idx, "no fused NVFP4 pre-expert quantization") + return None + gain, global_scale = args + # One request's warp count (the built-in launch's at 8 rows), for router_shard_rows. + warps = _shard.builtin_warps( + _shard.MAX_ROWS, getattr(self, "_prepare_stream", None) is not None + ) + if self._pool_direct_armed is not None: + # pool_into_route's wide site: the same chain with the producers publishing. + return self._pool_direct_armed.prepare(self, hidden_states, warps, gain, global_scale) + routing = self.experts.routing_method + side = self._router_stream + local = _shard.local_experts(world) + with self.nvtx("router"): + main = torch.cuda.current_stream() + logits_local = torch.empty( + (rows, local), device=hidden_states.device, dtype=torch.float32 + ) + if side is None: + self._wide_router_shard(hidden_states, inbox.rank, warps, logits_local, world) + else: + side.wait_stream(main) + with torch.cuda.stream(side): + self._wide_router_shard(hidden_states, inbox.rank, warps, logits_local, world) + with self.nvtx("pool"): + projected = self._replicated_projection( + self.pooler, hidden_states, self._pool_gather_workspace + ) + if side is not None: + main.wait_stream(side) + bias = routing.e_score_correction_bias + _require_aligned(pooled=projected, gain=gain, bias=bias, logits=logits_local) + out = _route.pool_gather_route( + projected, + logits_local, + inbox, + gain, + global_scale, + self.pre_expert_norm.variance_epsilon, + bias, + routing.renorm_eps, + _route.PLACE_NONE, + enable_pdl=self._prologue_enable_pdl, + return_gathered=False, + ) + self._prologue_record_routing(out, _route.PLACE_NONE, None, world) + _report_tep8(world, rows, self.layer_idx, None, self._wide_router_rows) + return out.logits, Fp4QuantizedTensor(out.fp4, out.scale, is_sf_swizzled=False) + def _prologue_record_routing(self, out, place: int, state, world: int = 8) -> None: """What the experts and the finalize read, keyed on the logits object: the routing memo, and the replicated experts' placement.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_split_kv/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_split_kv/attention.py index 76a304e5487e..36a1cf8b040e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_split_kv/attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_split_kv/attention.py @@ -154,6 +154,13 @@ class MultiRequestCall: window: int launched: bool = False fallback: Optional[str] = None + #: ``COMBINE_GATE``: launch the split only; ``pending`` launches the combine (the layer's + #: output gate does, with the gate fused when the built-in would run its Triton norm-gate). + defer: bool = False + pending: Optional[object] = None + out_ptr: int = 0 + out_shape: Tuple[int, ...] = () + out_stride: int = 0 def decline(self, kwargs: dict) -> Optional[str]: """Why this decode FMHA call is not the plain case, or ``None``. Host-side reads only.""" @@ -250,17 +257,24 @@ def decline(self, kwargs: dict) -> Optional[str]: def launch(self, kwargs: dict) -> None: query = kwargs["query"] rows = query.shape[0] - _mr.multi_request_split_kv( + out = kwargs["out"].view(rows, -1) + result = _mr.multi_request_split_kv( query.view(rows, -1), kwargs["kv_cache"][0], kwargs["block_tables"], kwargs["seq_lens"], kwargs["bmm1_scale"], kwargs["bmm2_scale"], - kwargs["out"].view(rows, -1), + out, window=self.window, enable_pdl=bool(kwargs["enable_pdl"]), + defer_combine=self.defer, ) + if self.defer: + self.pending = result + self.out_ptr = out.data_ptr() + self.out_shape = tuple(out.shape) + self.out_stride = out.stride(0) self.launched = True @@ -323,6 +337,10 @@ class MultiRequestSplitKvAttention: _mr_kv_heads: int = 4 _mr_min_requests: int = 2 _mr_census: Optional[MultiRequestCensus] = None + #: ``COMBINE_GATE``: this call's output gate (``forward_impl``) and the deferred combine + #: (``_attn_impl``), which ``apply_output_gate`` launches. + _mr_gate: Optional[torch.Tensor] = None + _mr_pending: Optional[MultiRequestCall] = None def forward(self, position_ids, hidden_states, attn_metadata, *args, **kwargs): census = self._mr_census @@ -333,12 +351,30 @@ def forward(self, position_ids, hidden_states, attn_metadata, *args, **kwargs): census.end() return result + def forward_impl(self, *args, output_gate=None, **kwargs): + self._mr_gate = output_gate + try: + return super().forward_impl(*args, output_gate=output_gate, **kwargs) + finally: + self._mr_gate = None + call, self._mr_pending = self._mr_pending, None + if call is not None: + # No output gate consumed the deferred combine: the attention output still + # has to be written, before anything reads it. + call.pending() + def _attn_impl(self, q, k, v, attn_metadata, *args, **kwargs): if not takes_step(attn_metadata, self._mr_min_requests): return super()._attn_impl(q, k, v, attn_metadata, *args, **kwargs) - call = MultiRequestCall(kv_heads=self._mr_kv_heads, window=self._mr_window) + call = MultiRequestCall( + kv_heads=self._mr_kv_heads, + window=self._mr_window, + defer=_mr.COMBINE_GATE and self._mr_gate is not None, + ) with intercepting(call): result = super()._attn_impl(q, k, v, attn_metadata, *args, **kwargs) + if call.launched and call.pending is not None: + self._mr_pending = call census = self._mr_census if census is not None: if call.launched: @@ -347,6 +383,51 @@ def _attn_impl(self, q, k, v, attn_metadata, *args, **kwargs): census.record(self.layer_idx, BUILTIN, call.fallback or "the FMHA was not called") return result + def apply_output_gate(self, attention_output, gate): + call, self._mr_pending = self._mr_pending, None + if call is None: + return super().apply_output_gate(attention_output, gate) + if self._fits_fused_gate(call, attention_output, gate): + trigger = getattr(self, "_norm_gate_trigger", None) or "exit" + logger.info_once( + "[musespark_open] multi_request_combine_gate: the " + f"combine applies the output norm-gate ({attention_output.shape[0]} rows, " + f"trigger {trigger})", + key=f"musespark_open.multi_request_split_kv.combine_gate_{attention_output.shape[0]}", + ) + return call.pending(gate, self.o_head_norm.variance_epsilon, trigger) + call.pending() + return super().apply_output_gate(attention_output, gate) + + def _fits_fused_gate(self, call: MultiRequestCall, attention_output, gate) -> bool: + """Exactly when the built-in ``apply_output_gate`` (or the PDL norm-gate copy) would run + its Triton kernel on this call's attention output (at most 128 rows), so the fused gate's + bits are that kernel's.""" + rows = call.out_shape[0] if call.out_shape else 0 + return ( + isinstance(attention_output, torch.Tensor) + and type(attention_output) is torch.Tensor + and attention_output.is_cuda + and attention_output.dtype == torch.bfloat16 + and attention_output.ndim == 2 + and attention_output.data_ptr() == call.out_ptr + and tuple(attention_output.shape) == call.out_shape + and attention_output.stride(0) == call.out_stride + and attention_output.stride(-1) == 1 + and 0 < rows <= _mr.COMBINE_GATE_MAX_ROWS + and not torch.compiler.is_compiling() + and self.head_dim == _mr.HEAD_DIM + and attention_output.shape[1] % self.head_dim == 0 + and isinstance(gate, torch.Tensor) + and gate.is_cuda + and gate.dtype == torch.bfloat16 + and gate.ndim in (2, 3) + and gate.stride(-1) == 1 + and gate.shape[0] == rows + and gate.numel() == attention_output.numel() + and (gate.ndim == 2 or gate.shape[-1] == self.head_dim) + ) + @functools.lru_cache(maxsize=None) def _mr_class(base: type) -> type: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/pool_into_route/route.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/pool_into_route/route.py index fba05328747c..263fa0d5909a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/pool_into_route/route.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/pool_into_route/route.py @@ -51,6 +51,19 @@ ``moe_prologue`` then runs ``DirectPoolSite.prepare`` in place of the tagged chain. Everything is disarmed on any exit. +Steps of two to four requests at TP8 (``tep8_wide_direct_pool``, default off; +the TP8 side of upstream b5ca973's TP4 wide direct pool sites, #21377): when the TEP8 router +shard (``tep8_router_shard``) gave every routed layer its TP8 wide route inbox, +``install_sites`` also builds a second transport on the wide layout +(``PoolDirectWorkspace(max_rows=WIDE_MAX_ROWS)``, routing only) with the same layers and buffer +sequence, and a 16..32-row forward captured into a CUDA graph arms those sites. The router +shard's wide path (``_wide_prepare``) then runs ``DirectPoolSite.prepare``: its router shard +(``_wide_router_shard``) writes this rank's ``[rows, 32]`` logits into the record's logit block +on the router fork's side stream, the pool GEMM (``wide_projection_gathers``' split-K cell) +writes the slot, and the routed kernel only polls them and runs the wide tagged kernel's +statements (``pool_gather_route``'s on the 32-row route inbox). Same bits as that path; the +placement stays ``ring_buddy_assign``'s, from the routing memo. + An armed layer whose call does not take the prologue's routed path raises: its peers would poll for a record nobody publishes. Eager steps and other row counts keep the tagged chain; the transports never share memory. @@ -69,6 +82,7 @@ from ... import dispatch as _tp_agreement from ...ops.comm import pool_direct_route as _pdr from ...ops.comm import unpool_direct_boundary as _direct +from ...ops.comm.projection_inbox import WIDE_MAX_ROWS from ...ops.moe import router_shard as _shard from ..dense_gemms.dispatcher import ( FLASHINFER_FAMILIES, @@ -79,8 +93,14 @@ ) from ..projection_gathers.gathers import _require_aligned -#: The one row count the transport is armed for: the K = 7 verify step. +#: The one row count the 8-row transport is armed for: the K = 7 verify step. ROWS = 8 +#: Rows of the narrowest wide call: two requests' K = 7 verify step. Pre-#21754 this was +#: ``router_shard.WIDE_MIN_ROWS``; 736cfdc7 moved the wide router (and that constant) out of +#: the sm103 tree into sm107's, so the wide sites carry the bound themselves. Same value. +WIDE_MIN_ROWS = 16 +#: The lever of the TP8 wide sites (16..32-row steps beside the TEP8 router shard): +#: ``tep8_wide_direct_pool``. Off, nothing is built. def _publish_copy(slot: torch.Tensor, product: torch.Tensor) -> None: @@ -104,16 +124,24 @@ def __init__( self.buffer = int(buffer) self.dispatcher = dispatcher - def prepare(self, moe, hidden_states: torch.Tensor, warps: int, gain, global_scale): + def prepare(self, moe, hidden_states: torch.Tensor, warps, gain, global_scale): """The prologue MoE's ``prepare_moe_inputs`` result for an armed call, through the transport. ``moe`` is the layer's prologue MoE with the router fork's side stream. ``warps``, ``gain`` and ``global_scale`` are what its routed path computed for this - call (its 8-row path). + call (its 8-row path, or its wide path on a wide site, where ``warps`` is unused: + the router shard picks its own). """ from ..moe_prologue.prologue import _placement rows = hidden_states.shape[0] + if self.workspace.wide: + if not WIDE_MIN_ROWS <= rows <= self.workspace.max_rows: + raise RuntimeError( + f"[musespark_open] pool_into_route: a layer armed for wide calls got {rows} rows; " + "its peers poll the transport for its records" + ) + return self._prepare_wide(moe, hidden_states, warps, gain, global_scale) if rows != ROWS: raise RuntimeError( f"[musespark_open] pool_into_route: a layer armed for {ROWS}-row calls got {rows} rows; its peers " @@ -174,29 +202,91 @@ def prepare(self, moe, hidden_states: torch.Tensor, warps: int, gain, global_sca moe._prologue_record_routing(out, place, state) return out.logits, Fp4QuantizedTensor(out.fp4, out.scale, is_sf_swizzled=False) + def _projected_into_slot(self, moe, hidden_states: torch.Tensor, slot: torch.Tensor) -> None: + self.dispatcher.publish = slot + try: + projected = moe._replicated_projection( + moe.pooler, hidden_states, moe._pool_gather_workspace + ) + finally: + self.dispatcher.publish = None + if projected.data_ptr() != slot.data_ptr(): + raise RuntimeError( + "[musespark_open] pool_into_route: the pool GEMM of an armed layer did not write the transport " + "slot (it bypassed the dense_gemms pool dispatcher); its peers poll for that slot" + ) + + def _prepare_wide(self, moe, hidden_states: torch.Tensor, warps, gain, global_scale): + """``prepare`` of a 16..32-row call on the wide layout: ``_wide_prepare``'s chain with the + producers publishing (upstream b5ca973's ``_prepare_wide``, routing only at TP8).""" + from ..moe_prologue.prologue import _report_tep8 + + rows = hidden_states.shape[0] + world = self.workspace.world + place, static_from, state = _wide_placement(moe, world, rows) + workspace, buffer = self.workspace, self.buffer + routing = moe.experts.routing_method + side = moe._router_stream + main = torch.cuda.current_stream() + with moe.nvtx("router"): + block = workspace.block(buffer, rows) + side.wait_stream(main) + with torch.cuda.stream(side): + moe._wide_router_shard(hidden_states, workspace.rank, warps, block, world) + with moe.nvtx("pool"): + self._projected_into_slot(moe, hidden_states, workspace.slot(buffer, rows)) + main.wait_stream(side) + bias = routing.e_score_correction_bias + _require_aligned(gain=gain, bias=bias) + out = _pdr.pool_direct_route( + workspace, + buffer, + rows, + gain, + global_scale, + moe.pre_expert_norm.variance_epsilon, + bias, + routing.renorm_eps, + place, + static_from, + enable_pdl=moe._prologue_enable_pdl, + return_gathered=False, + ) + moe._prologue_record_routing(out, place, state, world) + _report_tep8(world, rows, moe.layer_idx, None, moe._wide_router_rows, direct=True) + return out.logits, Fp4QuantizedTensor(out.fp4, out.scale, is_sf_swizzled=False) + class DirectPool: """What ``DecodeTarget`` holds under ``pool_into_route``: the workspaces and, once loaded, their sites. - ``sites`` is ``(moe, site)`` per layer that runs the routed path, in call order. + ``sites`` is ``(moe, site)`` per layer that runs the routed path, in call order; + ``wide_sites`` the same layers on the TP8 wide workspace (``wide_workspace``, built by + ``install_sites`` under ``tep8_wide_direct_pool`` beside the TEP8 router shard), else empty. """ def __init__( self, workspace: _pdr.PoolDirectWorkspace, + wide_workspace: Optional[_pdr.PoolDirectWorkspace] = None, ) -> None: self.workspace = workspace + self.wide_workspace = wide_workspace self.sites: Tuple[Tuple[torch.nn.Module, DirectPoolSite], ...] = () + self.wide_sites: Tuple[Tuple[torch.nn.Module, DirectPoolSite], ...] = () def takes(self, rows: int) -> bool: """Whether a decode forward of ``rows`` rows runs on a direct transport. - Inside a CUDA graph capture: exactly 8 rows. Every rank captures the same graphs with - the same row counts and installs the same sites, so either every rank arms or none does. + Inside a CUDA graph capture: exactly 8 rows, or 16..32 rows when the wide sites are + installed. Every rank captures the same graphs with the same row counts and installs + the same sites, so either every rank arms or none does. """ if not torch.cuda.is_current_stream_capturing(): return False - return rows == ROWS + if rows == ROWS: + return True + return bool(self.wide_sites) and WIDE_MIN_ROWS <= rows <= WIDE_MAX_ROWS @contextlib.contextmanager def armed(self, rows: int, skip: Optional[Callable[[int], bool]] = None): @@ -207,7 +297,7 @@ def armed(self, rows: int, skip: Optional[Callable[[int], bool]] = None): rank passes the same predicate, so every rank arms the same sites. Without it every site of the transport is armed. """ - sites = self.sites + sites = self.sites if rows <= ROWS else self.wide_sites if skip is not None: sites = tuple((moe, site) for moe, site in sites if not skip(int(moe.layer_idx))) for moe, site in sites: @@ -238,11 +328,12 @@ def variants(self, layers: Sequence) -> list: keys.add(_pdr.variant(place, static_from, enable_pdl, True, early)) return sorted(keys) - def install_sites(self, model, mapping) -> int: + def install_sites(self, model, mapping, wide: bool = False) -> int: """Find, compile and claim the sites once the weights are final (collective). - Returns the number of sites. Raises ``ValueError`` on every TP rank alike when - fewer than two layers run the routed path, when another fast path already + ``wide`` is the ``tep8_wide_direct_pool`` lever: with it the TP8 wide sites are + built too. Returns the number of sites. Raises ``ValueError`` on every TP rank alike + when fewer than two layers run the routed path, when another fast path already publishes from a pool GEMM, or when a variant fails to compile on some rank. """ layers = list(model.layers) @@ -309,8 +400,87 @@ def install_sites(self, model, mapping) -> int: f"{len(variants)} kernel variants); eager steps and other row " "counts keep the tagged chain" ) + if wide: + self._install_wide_sites(routed, buffers, mapping) return len(sites) + def _install_wide_sites(self, routed: Sequence, buffers: Sequence[int], mapping) -> None: + """``tep8_wide_direct_pool``: the TP8 wide workspace and its sites on the routed layers (collective). + + Needs the TEP8 router shard's wide route inbox (``_wide_inbox``, TP8) on every routed + layer: its ``_wide_prepare`` is the path the wide sites publish for. A refusal is + agreed over the TP group and only warns (16..32-row steps keep that path's tagged + route inbox); a variant that does not compile raises on every rank. + """ + reason = None + if self.workspace.world != 8: + reason = f"TP{self.workspace.world}: the wide sites are built at TP8 only" + else: + for layer in routed: + inbox = getattr(layer.feed_forward, "_wide_inbox", None) + if inbox is None or getattr(inbox, "world", None) != 8: + reason = ( + f"layer {layer.layer_idx} has no TP8 wide route inbox (set " + "tep8_router_shard, whose 16..32-row path these sites publish for)" + ) + break + reason = _tp_agreement.agreed_refusal(mapping, reason) + if reason is not None: + logger.warning( + "[musespark_open] tep8_wide_direct_pool is on but the wide sites are not built: " + f"{reason}" + ) + return + # The group the 8-row workspace was built for (``install``: the layers' mapping). + wide = _pdr.shared_workspace(routed[0].mapping, max_rows=WIDE_MAX_ROWS) + self.wide_workspace = wide + wide_variants = self.wide_variants(routed) + failed: Optional[Exception] = None + try: + wide.build_variants(wide_variants) + except Exception as error: # noqa: BLE001 - every failure must reach the agreement + failed = error + reason = _tp_agreement.agreed_refusal( + mapping, + None + if failed is None + else f"[musespark_open] pool_into_route could not build its wide routed kernel: {failed!r}", + ) + if reason is not None: + raise ValueError(reason) from failed + wide_sites = [] + for (moe, site), buffer in zip(self.sites, buffers): + wide_sites.append((moe, DirectPoolSite(wide, buffer, site.dispatcher))) + self.wide_sites = tuple(wide_sites) + wide.start_guard_watch() + logger.info( + f"[musespark_open] tep8_wide_direct_pool: TP8 decode steps of " + f"{WIDE_MIN_ROWS}..{WIDE_MAX_ROWS} rows (captured) publish the pool slice and the " + f"router shard's logits into the wide routed kernel's direct transport on {len(wide_sites)} " + f"layers (routing only, {len(wide_variants)} kernel variant(s); buffers as the 8-row sites')" + ) + + def wide_variants(self, layers: Sequence) -> list: + """The wide routed kernel's variants the layers' wide calls launch (``_wide_placement``).""" + keys = set() + for layer in layers: + moe = layer.feed_forward + for rows in (WIDE_MIN_ROWS, WIDE_MAX_ROWS): + place, static_from, _ = _wide_placement(moe, self.wide_workspace.world, rows) + keys.add(_pdr.variant(place, static_from, bool(moe._prologue_enable_pdl), False)) + return sorted(keys) + + +def _wide_placement(moe, world: int, rows: int): + """``(kind, static_from, state)`` of a wide call: the half window's placement on a layer + whose placement serves ``rows`` rows (TP4), else routing only (``ring_buddy_assign`` places).""" + from ..moe_prologue.prologue import _placement + + place, static_from, state = _placement(moe, world) + if place == _pdr.PLACE_HALF and world == 4 and state is not None and rows <= state.max_rows: + return place, static_from, state + return _pdr.PLACE_NONE, _pdr.LOCAL_EXPERTS, None + def _pool_dispatcher(moe) -> Optional[DenseGemmDispatcher]: pooler = getattr(moe, "pooler", None) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_gathers/gathers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_gathers/gathers.py index c15c1d29af12..23a943cef957 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_gathers/gathers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_gathers/gathers.py @@ -130,11 +130,21 @@ class MuseSparkOpenTepGatherMoE(MuseSparkOpenMoE): _published = None _published_pool_buffer: Optional[int] = None _published_unpool_buffer: Optional[int] = None + #: ``wide_direct_transports`` (sm103 TEP8, ``tep8_wide_direct_transports``): + #: its arming (``WideDirectTransports``) and this layer's pool buffer; None without it. + _wide_direct = None + _wide_direct_pool_buffer: Optional[int] = None def _replicated_projection(self, module, hidden_states, workspace): defer = (self._defer_unpool and module is self.unpooler) or ( self._defer_pool and module is self.pooler ) + wide_direct = self._wide_direct + if defer and module is self.pooler and wide_direct is not None and wide_direct.is_armed: + # 16/24/32 rows: the pool slice GEMM publishes into the wide pool transport. + slot = wide_direct.project_pool(self, module, hidden_states) + if slot is not None: + return slot published = self._published if defer and published is not None and published.is_armed: # 33..512 rows: the slice GEMM publishes into the site's transport. @@ -168,6 +178,11 @@ def _replicated_projection(self, module, hidden_states, workspace): return projected def prepare_expert_input(self, pooled: torch.Tensor): + wide_direct = self._wide_direct + if wide_direct is not None and wide_direct.is_armed: + quantized = wide_direct.consume_pool(self, pooled) + if quantized is not None: + return quantized published = self._published if published is not None and published.is_armed: quantized = published.consume_pool(self, pooled) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_publish/publish.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_publish/publish.py index c0b58d7e186f..faec2688e963 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_publish/publish.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_publish/publish.py @@ -66,13 +66,17 @@ MIN_ROWS, MAX_ROWS = _publish.MIN_ROWS, 448 -def _pool_args(moe) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: +def _pool_args( + moe, max_rows: int = MAX_ROWS +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: """(gain, global scale) exactly when the built-in ``prepare_expert_input`` would quantize a - BF16 ``[33..448, 4096]`` pool output with flashinfer's fused ``rmsnorm_fp4quant``; else None. + BF16 ``[1..max_rows, 4096]`` pool output with flashinfer's fused ``rmsnorm_fp4quant``; else None. The built-in's conditions that do not depend on the call's tensor (which is a fresh, - contiguous BF16 CUDA tensor of 33..448 rows on the armed path). Reads nothing that - differs between ranks. + contiguous BF16 CUDA tensor of at most ``max_rows`` rows on the armed path). Reads nothing + that differs between ranks. ``max_rows`` defaults to this path's own bound (448), so the + 33..448-row publish asks exactly what it asked before; ``wide_direct_transports`` passes + its own 32 (an expert op admitting 32 tokens need not admit 448). """ from tensorrt_llm._torch.modules.fused_ops.rmsnorm_fp4_quant import rmsnorm_fp4_quant_available from tensorrt_llm._torch.moe.fused_moe.moe_op_backend import TRTLLMOpBackend @@ -84,7 +88,7 @@ def _pool_args(moe) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: or not isinstance(getattr(backend, "op_backend", None), TRTLLMOpBackend) or not getattr(backend, "has_nvfp4", False) or getattr(backend, "fc31_act_scale", None) is not None - or MAX_ROWS > getattr(moe.experts, "moe_max_num_tokens", MAX_ROWS) + or max_rows > getattr(moe.experts, "moe_max_num_tokens", max_rows) or moe.expert_hidden_size != _pool.HIDDEN or not rmsnorm_fp4_quant_available() ): @@ -352,7 +356,13 @@ def project(self, moe, module, hidden_states: torch.Tensor) -> Optional[torch.Te weight = module.weight[begin : begin + columns] if site == "pool" and getattr(_publish, "TEP8_POOL_PROGRAMMATIC", False): self._programmatic_pool_fork(moe, hidden_states, weight, slot) - _publish.publish_projection(hidden_states, weight, slot) + # ``batch_splitk`` (dense_gemms, tep8_batch_splitk): its kernel stores + # each element of the slot once, after its reduction, as the cuBLAS producer does; every + # rank takes the same branch (the switch and the plan's table are the same on every rank). + dispatcher = module.__dict__.get("_low_m_dispatcher") + batch = getattr(dispatcher, "batch_splitk_apply", None) + if batch is None or batch(hidden_states, weight, None, out=slot) is None: + _publish.publish_projection(hidden_states, weight, slot) self._unconsumed += 1 self._report(site, rows) return slot diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py index 80f748436e22..0bbfee1005d3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py @@ -15,7 +15,9 @@ parameters are re-bound as views of the stack's first half, so the own block is stored once. -Each decode step (``M <= 16`` rows), the routing method's call additionally +Each decode step of at most ``max_rows`` rows (``MAX_ROWS``, 16, or what +``tep8_ring_buddy_wide`` sets on the tep8 ring, ``max_rows_for``), +the routing method's call additionally runs ``ops/moe/ring_buddy_assign`` on the main stream. In the built-in MoE input schedule it overlaps the pool/quantize work of the preparation stream; where the main stream is the longer branch instead (the router after the pool @@ -36,6 +38,14 @@ Prefill rows and any call the stacked path does not cover take the built-in expert call, which now reads the own-block views. +``tep8_wide_balance``'s ``MAX_ROWS`` (``wide_balance_rows``, tep8; default +off) raises the NVFP4 layers' cap to ``rows`` (up to 256): above the cap, a decode step +otherwise runs owner-only and every layer waits in the MoE all-reduce for the rank that +streams the most experts. Those layers' placements of more than 64 rows run on +``ring_buddy_assign``'s wide program (the same placement in about 2 us, where the +``[SLOTS, RANKS]`` program takes 3.5 us at 1024 slots and 5.8 us at 2048). Their stacked calls +keep the runner's default tactic, as below the cap. The two BF16 layers keep the cap. + This module is a load-time installer, not an op: it patches the built-in modules once the weights are final, so it sits beside ``modeling.py`` and is imported only when ``replicated_experts`` is planned. Its one kernel is the @@ -89,8 +99,10 @@ from __future__ import annotations +import threading +import time import types -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Tuple import torch @@ -104,15 +116,200 @@ BLOCK, EXPERTS, HALF_SPAN, + MAX_SLOTS, + PLACEMENT_STATS_FIELDS, RANKS, TOP_K, + WIDE_MIN_SLOTS, + accumulate_placement_stats, block_of, ring_buddy_assign, + wide_balance_rows, ) #: Decode rows the stacked path serves; larger calls keep the built-in path. MAX_ROWS = 16 +#: ``tep8_ring_buddy_wide``'s ``MAX_ROWS`` replaces ``MAX_ROWS`` on the tep8 ring (32-expert +#: blocks, the ring whose balanced placement spreads a step's hot experts over the ranks): the +#: stacked path then balances the decode calls of up to this many rows, ``8 * B`` for a verify +#: step of ``B`` requests at K = 7 (128 for 16 requests, 160 for 20), up to ``MAX_SLOTS // +#: TOP_K`` (256: the rows of one ``ring_buddy_assign`` program; ``install`` refuses more). Off +#: (the default), ``MAX_ROWS`` stands. +_TEP8_MAX_ROWS: Optional[int] = None + +#: ``tep8_experts_in_kernel_placement`` (sm103 tep8, the persistent kernel's batch builds): the +#: NVFP4 layers' stacked calls of 9..``persistent_rows`` rows skip ``ring_buddy_assign`` and +#: hand the step's selection to the batch build, which computes the same balanced placement for +#: its rank (``persistent_experts``' ``place_rank``). Read at ``install``. Not with the +#: placement statistics (``ring_buddy_stats``), which read the program's census: where they are +#: on, ``_State.current`` keeps the launch (its ``_stats is None``) and this does nothing. +_IN_KERNEL_PLACEMENT = False + +#: ``ring_buddy_stats``' ``INTERVAL_S``: whole seconds between reports, or None (off). +_STATS_INTERVAL_S: Optional[int] = None + + +def configure( + *, + in_kernel_placement: bool = False, + tep8_max_rows: Optional[int] = None, + wide_balance_rows: Optional[int] = None, + stats_interval_s: Optional[int] = None, +) -> None: + """Set the settings this module reads at ``install`` and on its decode calls. + + ``lever_settings.apply`` sets them from the model's plan before any install; the defaults + are the stock path (the 16-row cap, no in-kernel placement, no wide band, no census). + + Args: + in_kernel_placement: ``tep8_experts_in_kernel_placement`` (``_IN_KERNEL_PLACEMENT``). + tep8_max_rows: ``tep8_ring_buddy_wide``'s ``MAX_ROWS``; None keeps ``MAX_ROWS``. + wide_balance_rows: ``tep8_wide_balance``'s ``MAX_ROWS``, handed to + ``ring_buddy_assign.configure``; None keeps the band off. + stats_interval_s: ``ring_buddy_stats``' ``INTERVAL_S``; None counts nothing. + + Raises: + ValueError: for a row cap outside ``1..MAX_SLOTS // TOP_K`` or an interval below 1 s. + """ + global _IN_KERNEL_PLACEMENT, _TEP8_MAX_ROWS, _STATS_INTERVAL_S + if tep8_max_rows is not None: + rows = int(tep8_max_rows) + if not 0 < rows <= MAX_SLOTS // TOP_K: + raise ValueError( + f"[musespark_open] tep8_ring_buddy_wide MAX_ROWS {rows}: the row cap is placed " + f"by one ring_buddy_assign program, so it is at most {MAX_SLOTS // TOP_K} rows " + f"({MAX_SLOTS} slots)" + ) + tep8_max_rows = rows + if stats_interval_s is not None: + seconds = int(stats_interval_s) + if seconds < 1: + raise ValueError( + f"[musespark_open] ring_buddy_stats INTERVAL_S {seconds}: expected whole " + "seconds between reports, at least 1" + ) + stats_interval_s = seconds + _IN_KERNEL_PLACEMENT = bool(in_kernel_placement) + _TEP8_MAX_ROWS = tep8_max_rows + _STATS_INTERVAL_S = stats_interval_s + _configure_wide_balance(wide_balance_rows) + + +def _configure_wide_balance(rows: Optional[int]) -> None: + """Hand ``tep8_wide_balance``'s ``MAX_ROWS`` to the assign module, which owns the band. + + One ``configure`` reaches this module (``lever_settings``), and the band belongs to + ``ring_buddy_assign``; this keeps both in step without a second entry in the settings + table. ``_wide_rows_here`` reads it back. + """ + from ...ops.moe.ring_buddy_assign import configure as _assign_configure + + _assign_configure(wide_balance_rows=rows) + + +#: A placement's stats slot when the expert call places it +#: (``tep8_experts_in_kernel_placement``): its ids are then the step's selection and its +#: weights the routing's FP32 ones. +_PLACED_IN_KERNEL = "placed in the expert kernel" + + +def _in_kernel_placement() -> bool: + """Whether ``tep8_experts_in_kernel_placement`` is on (``configure``).""" + return _IN_KERNEL_PLACEMENT + + +#: ``(the built-in cap, the cap tep8_ring_buddy_wide set)`` once ``max_rows_for`` applied the +#: lever in this process, else None. +_override: Optional[Tuple[int, int]] = None +_override_logged = False + + +def max_rows_for(ranks: Optional[int]) -> int: + """The stacked path's row cap on a ring of ``ranks``. + + ``MAX_ROWS``, which ``tep8_ring_buddy_wide``'s ``MAX_ROWS`` (``configure``) replaces on the + tep8 ring when it is on; on any other ring the lever is ignored, with a warning. A + half-window layer keeps ``HALF_MAX_ROWS`` either way. + """ + global _override + rows = _TEP8_MAX_ROWS + if rows is None: + return MAX_ROWS + if ranks != 8: + logger.warning( + f"[musespark_open] tep8_ring_buddy_wide MAX_ROWS={rows} has no effect on a ring of " + f"{ranks} ranks: it replaces the row cap of the tep8 ring only; the stacked path " + f"keeps {MAX_ROWS} rows" + ) + return MAX_ROWS + _override = (MAX_ROWS, rows) + logger.info( + f"[musespark_open] replicated_experts: tep8_ring_buddy_wide MAX_ROWS={rows} sets the " + f"balanced placement's row cap at tep{ranks} (the built-in {MAX_ROWS}): decode calls of " + f"up to {rows} rows take the stacked path, each placed by one ring_buddy_assign program" + ) + return rows + + +def _log_override_served(rows: int, rank: int) -> None: + """Once per process: the first placement of a call above the built-in cap (``_override``).""" + global _override_logged + if _override_logged: + return + _override_logged = True + stock, cap = _override + logger.info( + f"[musespark_open] replicated_experts: the balanced placement serves a {rows}-row decode " + f"call on TP rank {rank}, above the built-in {stock} rows (tep8_ring_buddy_wide " + f"MAX_ROWS={cap}; logged " + "once per process)" + ) + + +_wide_logged = False + + +def _log_wide_served(rows: int, rank: int, base: int) -> None: + """Once per process: the first placement of a call in the band ``tep8_wide_balance`` adds.""" + global _wide_logged + if _wide_logged: + return + _wide_logged = True + logger.info( + f"[musespark_open] replicated_experts: the balanced placement serves a {rows}-row " + f"({rows * TOP_K}-slot) decode call of an NVFP4 layer on TP rank {rank} on the wide " + f"ring_buddy_assign program, above the {base}-row cap (tep8_wide_balance; logged once " + "per process)" + ) + + +def _wide_layer(kind: str, windows: int, half: bool, balance: bool, ranks: Optional[int]) -> bool: + """Whether ``tep8_wide_balance`` extends a layer: an NVFP4 layer's balanced two-window + placement on the tep8 ring (the BF16 layers, the owner-only window and the half window + keep their caps).""" + return kind == "nvfp4" and windows >= 2 and not half and bool(balance) and ranks == 8 + + +def _wide_rows_here(ranks: Optional[int]) -> Optional[int]: + """``tep8_wide_balance``'s cap where it applies (the tep8 ring of this sm103 target), else + None, with a warning when it is set elsewhere. + + Raises: + ValueError: for a malformed value (``wide_balance_rows``), on every rank alike. + """ + rows = wide_balance_rows() + if rows is None: + return None + if ranks != 8: + logger.warning( + f"[musespark_open] tep8_wide_balance MAX_ROWS={rows} has no effect on a ring of " + f"{ranks} ranks: it extends the tep8 balanced placement only" + ) + return None + return rows + + #: Decode rows the stacked path serves on a half-window layer (``half_window_experts``): #: up to four requests' verify steps, the slots of one ``ring_buddy_assign`` program. HALF_MAX_ROWS = 32 @@ -256,6 +453,7 @@ def __init__( windows: int = 2, ranks: int = RANKS, half: bool = False, + max_rows: Optional[int] = None, ): self.kind = kind self.stacks = stacks @@ -279,7 +477,8 @@ def __init__( # The half window (``half_window_experts``): the owner-only placement over the # stack's own block, and the prologue's half-window placement over all of it. self.half = half - self.max_rows = HALF_MAX_ROWS if half else MAX_ROWS + # ``install``'s ``max_rows_for`` (``tep8_ring_buddy_wide``), else the built-in cap. + self.max_rows = HALF_MAX_ROWS if half else (MAX_ROWS if max_rows is None else max_rows) if half: self.default_window = self.block self.by_window = { @@ -298,6 +497,24 @@ def __init__( # The per-rank loads (before, after) of the last placement; kept across # memo clears for the shadow validation mode. self.last_stats = None + # ``tep8_wide_balance``: the placements run ``ring_buddy_assign(wide=True)``, and calls + # above ``wide_base`` rows (the cap without the switch) are the band it adds. + self.wide = False + self.wide_base = self.max_rows + # Early start (``tep8_experts_early_start``, ``_patch_experts``): the + # stacked calls of ``early_start_rows`` (first, last) rows launch the placement with + # its early trigger, and the persistent kernel's early-start variant behind it. + # Only a placement the expert forward itself launches (``in_forward``; its ``x`` was + # written before) does; ``early_ids`` is that placement's ``local_ids`` until the same + # forward's expert call takes it. + self.early_start_rows: Optional[Tuple[int, int]] = None + self.in_forward = False + self.early_ids = None + # In-kernel placement (``tep8_experts_in_kernel_placement``): a placement the forward + # itself makes for a + # stacked call of ``place_rows`` (first, last) rows hands the selection to the batch + # build (its stats slot is ``_PLACED_IN_KERNEL``). + self.place_rows: Optional[Tuple[int, int]] = None def current(self, router_logits: torch.Tensor): """The placement computed from ``router_logits`` as it is now, or None.""" @@ -405,6 +622,28 @@ def apply(self, router_logits, input_ids=None): and weights.is_contiguous() and not torch.compiler.is_compiling() ): + if ( + state.in_forward + and state.place_rows is not None + and state.place_rows[0] <= router_logits.shape[0] <= state.place_rows[1] + and _stats is None + ): + # In-kernel placement: no placement launch; this forward's expert call hands + # the selection to the batch build, which places it for this rank. + state.early_ids = None + state.placement = ( + router_logits, + _version(router_logits), + selected, + weights, + _PLACED_IN_KERNEL, + ) + return selected, weights + early = ( + state.in_forward + and state.early_start_rows is not None + and state.early_start_rows[0] <= router_logits.shape[0] <= state.early_start_rows[1] + ) local_ids, weights_bf16, stats = ring_buddy_assign( selected, weights, @@ -413,7 +652,10 @@ def apply(self, router_logits, input_ids=None): state.static_from, ranks=state.ranks, half=state.half, + **({"wide": True} if state.wide else {}), + **({"early_trigger": True} if early else {}), ) + state.early_ids = local_ids if early else None state.placement = ( router_logits, _version(router_logits), @@ -424,6 +666,13 @@ def apply(self, router_logits, input_ids=None): if state.half: state.placement += (HALF_SPAN,) state.last_stats = stats + rows = router_logits.shape[0] + if _override is not None and rows > _override[0]: + _log_override_served(rows, state.rank) + if state.wide and rows > state.wide_base: + _log_wide_served(rows, state.rank, state.wide_base) + if _stats is not None: + _stats.add(stats, rows) return selected, weights def clear_routing_memo(self): @@ -434,7 +683,16 @@ def clear_routing_memo(self): routing.clear_routing_memo = types.MethodType(clear_routing_memo, routing) -def _nvfp4_call(backend, routing, persistent: bool = False, stacks=None, persistent_rows=8): +def _nvfp4_call( + backend, + routing, + persistent: bool = False, + stacks=None, + persistent_rows=8, + persistent_batch: bool = False, + persistent_early_start: bool = False, + persistent_place_rank: Optional[int] = None, +): """The stacked call for the native NVFP4 leaf, as its ``run_moe`` makes the built-in one. ``persistent`` (``persistent_experts``): calls of up to ``persistent_experts``'s @@ -450,6 +708,27 @@ def _nvfp4_call(backend, routing, persistent: bool = False, stacks=None, persist The call's window is the leading dim of the stacks it is given; the ``local_expert_offset=0`` TRTLLM-gen call runs over exactly that window. + + ``persistent_rows`` (``wide_persistent_experts``: ``persistent_experts.WIDE_MAX_ROWS``): + calls of up to that many rows on a window of at most 64 experts go to the persistent + kernel's wide build. (Restored here: #21754 / 736cfdc7 moved the sm107 wide lever out of + this tree and dropped the sm103 wiring with it; without the gate below, no stacked call + above ``persistent_experts.MAX_ROWS`` can reach the persistent kernel at all.) + + ``persistent_batch`` (sm103 tep8, ``tep8_persistent_batch_experts``): + those 9..``persistent_rows``-row calls go to the persistent kernel's batch builds + (``persistent_experts.BATCH_NTOKS``) instead of the wide build. + + ``persistent_early_start`` (with ``persistent_batch``, + ``tep8_experts_early_start``): the batch builds' early-start variants are + compiled too, and a batch call made with ``early_start=True`` (the forward launched its + placement with the early trigger) runs that variant. + + ``persistent_place_rank`` (with ``persistent_batch``, + ``tep8_experts_in_kernel_placement``): the batch builds' + in-kernel placement variants for this rank are compiled too, and a call made with + ``place_rank`` (its ids are the step's selection, its weights the routing's FP32 ones) runs + them; a call the batch builds do not serve is placed by ``ring_buddy_assign`` here first. """ from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( # noqa: E501 fp4_block_scale_moe_runner, @@ -476,14 +755,25 @@ def _nvfp4_call(backend, routing, persistent: bool = False, stacks=None, persist # One compiled kernel per window the layer's placements address. for window in sorted(persistent_windows): window_stacks = stacks[window] + widen = {} + if persistent_rows > _persistent.MAX_ROWS and window <= 64: + if persistent_batch: + widen = {"batch_rows": persistent_rows} + if persistent_early_start: + widen["early_start"] = True + if persistent_place_rank is not None and window == 64: + widen["place_rank"] = persistent_place_rank + else: + widen = {"wide": True} _persistent.persistent_experts.prepare( window_stacks["w3_w1_weight"], window_stacks["w3_w1_weight_scale"].view(torch.float8_e4m3fn), window_stacks["w2_weight"], window_stacks["w2_weight_scale"].view(torch.float8_e4m3fn), + **widen, ) - def call(stacks, x, local_ids, weights_bf16): + def call(stacks, x, local_ids, weights_bf16, early_start=False, place_rank=None): if not isinstance(x, Fp4QuantizedTensor) or x.is_sf_swizzled: return None x_sf = x.scaling_factor.reshape(-1).view(torch.float8_e4m3fn) @@ -491,8 +781,19 @@ def call(stacks, x, local_ids, weights_bf16): served = ( persistent_op is not None and window in persistent_windows - and x.fp4_tensor.shape[0] <= persistent_op.MAX_ROWS + and x.fp4_tensor.shape[0] + <= (persistent_rows if window <= 64 else persistent_op.MAX_ROWS) ) + if place_rank is not None and not ( + served + and persistent_batch + and persistent_place_rank is not None + and window == 64 + and x.fp4_tensor.shape[0] > persistent_op.MAX_ROWS + ): + # Only the batch builds place in the kernel: place the selection here first. + local_ids, weights_bf16, _ = ring_buddy_assign(local_ids, weights_bf16, place_rank) + place_rank = None if window > two_windows and not served: # The three-window placement runs only where the persistent kernel serves the # call (the prologue's steps of at most 8 rows); every rank shares this shape. @@ -501,6 +802,16 @@ def call(stacks, x, local_ids, weights_bf16): f"of {x.fp4_tensor.shape[0]} rows the persistent expert kernel does not serve" ) if served: + batch = {} + if persistent_batch and x.fp4_tensor.shape[0] > persistent_op.MAX_ROWS: + batch = {"batch": True} + _log_batch_once(x.fp4_tensor.shape[0], window) + if place_rank is not None: + batch["place_rank"] = place_rank + _log_placed_once(x.fp4_tensor.shape[0], window) + elif early_start and persistent_early_start: + batch["early_start"] = True + _log_early_start_once(x.fp4_tensor.shape[0], window) return persistent_op.persistent_experts( x.fp4_tensor, x_sf, @@ -513,6 +824,7 @@ def call(stacks, x, local_ids, weights_bf16): stacks["fc31_scale_c"], stacks["fc31_alpha"], stacks["fc2_alpha"], + **batch, ) tactic = None if not AutoTuner.get().is_tuning_mode: @@ -596,6 +908,75 @@ def call(stacks, x, local_ids, weights_bf16): return call +_BATCH_LOGGED: List[bool] = [] +_EARLY_START_LOGGED: List[bool] = [] +_PLACED_LOGGED: List[bool] = [] +_PLACEMENT_INSTALL_LOGGED: List[bool] = [] + + +def _log_placed_once(rows: int, window: int) -> None: + """One INFO line per process: the first stacked call the in-kernel placement serves.""" + if _PLACED_LOGGED: + return + _PLACED_LOGGED.append(True) + logger.info( + f"[musespark_open] persistent_experts: a {rows}-row stacked NVFP4 call on the " + f"{window}-expert window runs the batch build's in-kernel placement, with no " + "ring_buddy_assign launch (tep8_experts_in_kernel_placement; logged once per process)" + ) + + +def _log_placement_install_once(placed: bool) -> None: + """One line per process at install: the in-kernel placement is on (INFO), or the switch is + set where it does nothing (WARNING).""" + if _PLACEMENT_INSTALL_LOGGED: + return + _PLACEMENT_INSTALL_LOGGED.append(True) + if placed: + logger.info( + "[musespark_open] replicated_experts: tep8_experts_in_kernel_placement: the NVFP4 " + "layers' stacked " + "calls the persistent kernel's batch builds serve skip ring_buddy_assign; the batch " + "builds compute the balanced tep8 placement from the selection themselves" + ) + else: + logger.warning( + "[musespark_open] replicated_experts: tep8_experts_in_kernel_placement has no effect " + "here: it " + "needs sm103 tep8 with the persistent kernel's batch builds on the balanced " + "two-window ring-buddy stack (tep8_persistent_batch_experts)" + ) + + +def _log_early_start_once(rows: int, window: int) -> None: + """One INFO line per process: the first stacked call the early-start variant serves.""" + if _EARLY_START_LOGGED: + return + _EARLY_START_LOGGED.append(True) + logger.info( + f"[musespark_open] persistent_experts: a {rows}-row stacked NVFP4 call on the " + f"{window}-expert window runs the batch build's early-start variant behind " + "ring_buddy_assign's early trigger (tep8_experts_early_start; logged " + "once per process)" + ) + + +def _log_batch_once(rows: int, window: int) -> None: + """One INFO line per process: the first stacked call a persistent batch build serves. + + Python runs at eager warmup and graph capture only, so this marks the graphs captured + on the batch path. + """ + if _BATCH_LOGGED: + return + _BATCH_LOGGED.append(True) + logger.info( + f"[musespark_open] persistent_experts: a {rows}-row stacked NVFP4 call on the " + f"{window}-expert window runs the persistent kernel's batch build " + "(tep8_persistent_batch_experts; logged once per process)" + ) + + def _bf16_call(backend, routing, unpacked: bool = False): """The stacked call for the FlashInfer BF16 leaf, as its ``run_moe`` makes the built-in one. @@ -678,19 +1059,52 @@ def _patch_experts( persistent: bool = False, bf16_unpacked: bool = False, persistent_rows: int = 8, + persistent_batch: bool = False, + persistent_early_start: bool = False, ) -> None: experts = moe.experts backend = _backend(moe) builtin_forward = experts.forward routing = experts.routing_method if state.kind == "nvfp4" and persistent: + # Early start: only where the batch builds serve the layer's 9..persistent_rows-row + # calls (the two-window stack of an NVFP4 leaf the persistent kernel computes). + early_start = bool( + persistent_early_start + and persistent_batch + and persistent_rows > _persistent_max_rows() + and not state.half + and state.default_window <= 64 + and persistent_eligible(backend) + ) + # In-kernel placement: where the batch builds serve the balanced two-window tep8 stack. + placed = bool( + _in_kernel_placement() + and persistent_batch + and persistent_rows > _persistent_max_rows() + and not state.half + and state.default_window == 2 * state.block + and state.ranks == RANKS + and state.balance + and state.static_from == state.block + and persistent_eligible(backend) + ) + if _in_kernel_placement(): + _log_placement_install_once(placed) stacked_call = _nvfp4_call( backend, routing, persistent=True, stacks=state.by_window, persistent_rows=persistent_rows, + persistent_batch=persistent_batch, + persistent_early_start=early_start, + persistent_place_rank=state.rank if placed else None, ) + if early_start: + state.early_start_rows = (_persistent_max_rows() + 1, persistent_rows) + if placed: + state.place_rows = (_persistent_max_rows() + 1, persistent_rows) elif state.kind == "nvfp4": stacked_call = _nvfp4_call(backend, routing) elif bf16_unpacked: @@ -703,13 +1117,31 @@ def _patch_experts( def forward(self, x, router_logits, *args, do_finalize=True, **kwargs): if not state.active or do_finalize or args or not 0 < x.shape[0] <= state.max_rows: return builtin_forward(x, router_logits, *args, do_finalize=do_finalize, **kwargs) - routing.apply(router_logits) + # Early start pairs a placement this call launches (``x`` was written before it) + # with this call's expert launch only. + state.early_ids = None + state.in_forward = True + try: + routing.apply(router_logits) + finally: + state.in_forward = False placement = state.current(router_logits) if placement is None: + state.early_ids = None return builtin_forward(x, router_logits, *args, do_finalize=do_finalize, **kwargs) local_ids, weights_bf16 = placement[2], placement[3] window = placement[5] if len(placement) > 5 else default_window - outputs = stacked_call(by_window[window], x, local_ids, weights_bf16) + early = state.early_ids is not None and state.early_ids is local_ids + placed = placement[4] is _PLACED_IN_KERNEL + state.early_ids = None + if placed: + outputs = stacked_call( + by_window[window], x, local_ids, weights_bf16, place_rank=state.rank + ) + elif early: + outputs = stacked_call(by_window[window], x, local_ids, weights_bf16, early_start=True) + else: + outputs = stacked_call(by_window[window], x, local_ids, weights_bf16) if outputs is None: return builtin_forward(x, router_logits, *args, do_finalize=do_finalize, **kwargs) return outputs @@ -717,6 +1149,13 @@ def forward(self, x, router_logits, *args, do_finalize=True, **kwargs): experts.forward = types.MethodType(forward, experts) +def _persistent_max_rows() -> int: + """The persistent kernel's 8-row build bound (``persistent_experts.MAX_ROWS``).""" + from ...ops.moe import persistent_experts as _persistent + + return _persistent.MAX_ROWS + + def _act_type(backend) -> int: from tensorrt_llm._torch.moe.fused_moe.trtllm_gen.kernel_inputs import to_trtllm_gen_act_type @@ -769,6 +1208,8 @@ def install( persistent_rows: int = 8, windows: int = 2, half: bool = False, + persistent_batch: bool = False, + persistent_early_start: bool = False, ) -> List[int]: """Replicate and patch every eligible MoE layer of ``text_model``; returns their indices. @@ -789,7 +1230,12 @@ def install( ``bf16_unpacked`` are served without replication. ``persistent=True`` (``persistent_experts``) serves the NVFP4 layers' stacked calls of up to 8 rows with the - persistent expert kernel (``ops/moe/persistent_experts.py``). + persistent expert kernel (``ops/moe/persistent_experts.py``); + ``persistent_rows`` above 8 also its calls of up to that many rows on a window of at + most 64 experts, on the wide build or, with ``persistent_batch``, on the batch builds; + with ``persistent_early_start`` as well, those batch calls launch their placement with + its early trigger and run the batch builds' early-start variants + (``tep8_experts_early_start``). ``bf16_unpacked=True`` (``bf16_routed_experts``) gives the two BF16 layers' stacked calls unpacked routing (``ops/moe/bf16_routed_experts.py``): bitwise the packed call, without its three packing launches. @@ -823,9 +1269,19 @@ def install( refusal = _install_refusal( text_model, windows, persistent, balance, bf16_unpacked, ranks, static_from, half ) + max_rows = MAX_ROWS + wide_rows = None + if refusal is None: + try: + max_rows = max_rows_for(ranks) + wide_rows = _wide_rows_here(ranks) + except ValueError as error: + refusal = str(error) refusal = _tp_agreement.agreed_refusal(mapping, refusal) if refusal is not None: raise ValueError(refusal) + if wide_rows is not None and wide_rows <= max_rows: + wide_rows = None # the cap already covers the band # A mapping outside the served rings skips every layer (``eligible``); the # block only names the static placement in the log line then. block = block_of(ranks) if ranks is not None else BLOCK @@ -833,6 +1289,7 @@ def install( static_from = block failed: Optional[str] = None installed, refreshed = [], [] + widened_rows: List[int] = [] widened = [] halved = [] skipped = {} @@ -850,6 +1307,8 @@ def install( _refresh(moe, state, mapping) state.balance = balance state.static_from = static_from + state.max_rows = HALF_MAX_ROWS if state.half else max_rows + _apply_wide(state, wide_rows, widened_rows) refreshed.append(index) continue reason = eligible(moe, mapping) @@ -867,13 +1326,30 @@ def install( stacks = _replicate(moe, kind, mapping, layer_windows, layer_half) added += sum(t[block:].numel() * t.element_size() for t in stacks.values()) state = _State( - kind, stacks, mapping.tp_rank, balance, static_from, layer_windows, ranks, layer_half + kind, + stacks, + mapping.tp_rank, + balance, + static_from, + layer_windows, + ranks, + layer_half, + max_rows=max_rows, ) + _apply_wide(state, wide_rows, widened_rows) moe._replicated_experts = state if failed is None: _patch_routing(moe, state) try: - _patch_experts(moe, state, persistent, bf16_unpacked, persistent_rows) + _patch_experts( + moe, + state, + persistent, + bf16_unpacked, + persistent_rows, + persistent_batch, + persistent_early_start, + ) except Exception as error: # noqa: BLE001 - every failure must reach the agreement failed = f"layer {index}: {error!r}" installed.append(index) @@ -913,7 +1389,9 @@ def install( f"{'(persistent NVFP4 expert kernel) ' if persistent else ''}" f"{'(unpacked BF16 routing) ' if bf16_unpacked else ''}" f"{f'(three ring windows on {len(widened)} layers) ' if windows == 3 else ''}" - f"({placement} placement, " + f"({placement} placement, ring-buddy max_rows={max_rows}" + f"{f' ({HALF_MAX_ROWS} on the half-window layers)' if halved else ''}" + f"{f' ({wide_rows} on {len(widened_rows)} NVFP4 layers)' if widened_rows else ''}, " f"{dict(sorted(kinds.items()))}), refreshed in place on {len(refreshed)}; " f"per rank: replicated experts +{added / 2**30:.2f} GiB, torch allocated " f"+{allocated_added / 2**30:.2f} GiB, device +{device_added / 2**30:.2f} GiB, of " @@ -921,9 +1399,33 @@ def install( f"communicator buffers); built-in path kept for layers {sorted(skipped)}: " f"{sorted(set(skipped.values()))}" ) + if widened_rows: + logger.info( + f"[musespark_open] replicated_experts: tep8_wide_balance MAX_ROWS={wide_rows}: the " + f"balanced placement of {len(widened_rows)} NVFP4 layers serves decode calls of up to " + f"{wide_rows} rows ({wide_rows * TOP_K} slots; the cap was {max_rows}), placed by " + f"ring_buddy_assign's wide program above {WIDE_MIN_SLOTS // TOP_K} rows; the other " + f"layers keep {max_rows}" + ) + if balance: + _start_stats(mapping, ranks, max_rows, len(installed) + len(refreshed)) return sorted(installed + refreshed) +def _apply_wide(state: "_State", wide_rows: Optional[int], widened: List[int]) -> None: + """``tep8_wide_balance`` on one installed layer: its cap and the wide program, where + ``_wide_layer`` admits it; any other layer, or the switch unset, is left as it is.""" + state.wide = False + state.wide_base = state.max_rows + if wide_rows is None or not _wide_layer( + state.kind, state.windows, state.half, state.balance, state.ranks + ): + return + state.wide = True + state.max_rows = max(state.max_rows, wide_rows) + widened.append(state.max_rows) + + def _install_refusal( text_model, windows: int, @@ -986,3 +1488,98 @@ def _install_refusal( except ImportError as error: return f"bf16_routed_experts calls flashinfer.fused_moe.trtllm_bf16_routed_moe: {error}" return None + + +#: ``ring_buddy_stats``' ``INTERVAL_S`` seconds: every placement this process computes also +#: adds its census to a pinned host accumulator +#: (``ring_buddy_assign.accumulate_placement_stats``: one more one-program launch per +#: placement, diagnostics only), and TP rank 0 logs the totals at that interval. Off (the +#: default), nothing is counted. +_STATS_TAG = "ring_buddy_stats" + +_stats: Optional["_PlacementStats"] = None + + +class _PlacementStats: + """Placement census totals in pinned host memory, in two bands of + ``PLACEMENT_STATS_FIELDS``: calls of at most ``split`` rows (band 0) and above (band 1). + + The device adds to it in stream order (graph replays included); the host reads it + without a CUDA call, so the reader cannot disturb a graph capture on another thread. + """ + + def __init__(self, interval: float, split: int, ranks: int, rank: int): + self.interval = interval + self.split = split + self.ranks = ranks + self.rank = rank + self.layers = 0 + self.failed = False + self.acc = torch.zeros(2 * PLACEMENT_STATS_FIELDS, dtype=torch.int64, pin_memory=True) + + def add(self, stats: torch.Tensor, rows: int) -> None: + if self.failed: + return + band = int(rows > self.split) * PLACEMENT_STATS_FIELDS + try: + accumulate_placement_stats( + stats, self.acc[band : band + PLACEMENT_STATS_FIELDS], rows, ranks=self.ranks + ) + except Exception as error: # noqa: BLE001 - diagnostics must not stop serving + self.failed = True + logger.warning( + f"[musespark_open] {_STATS_TAG}: placement stats off on TP rank {self.rank}: " + f"{error!r}" + ) + + def summary(self) -> str: + values = self.acc.tolist() + parts = [] + for band, name in ((1, f"above {self.split} rows"), (0, f"up to {self.split} rows")): + calls, rows, owner, placed, distinct = values[ + band * PLACEMENT_STATS_FIELDS : band * PLACEMENT_STATS_FIELDS + 5 + ] + if not calls: + parts.append(f"{name}: none") + continue + mean = max(distinct, 1) / self.ranks + parts.append( + f"{name}: {calls} placements (about {calls / max(self.layers, 1):.0f} decode steps " + f"over {self.layers} MoE layers), {rows / calls:.1f} rows each; the busiest rank's " + f"distinct experts {owner / calls:.2f} owner-only -> {placed / calls:.2f} placed, " + f"the mean rank's {mean / calls:.2f} (busiest/mean {owner / mean:.3f} -> " + f"{placed / mean:.3f})" + ) + head = f"[musespark_open] replicated_experts placement stats (TP rank {self.rank}): " + return head + "; ".join(parts) + + def _report(self) -> None: + last = None + while True: + time.sleep(self.interval) + line = self.summary() + if line != last: + logger.info(line) + last = line + + def start(self) -> None: + if self.rank == 0: + threading.Thread(target=self._report, name="ring-buddy-stats", daemon=True).start() + + +def _start_stats(mapping: Mapping, ranks: Optional[int], max_rows: int, layers: int) -> None: + """``install``'s ``ring_buddy_stats``: the process's accumulator and its reporter (once).""" + global _stats + interval = _STATS_INTERVAL_S + if interval is None or ranks is None or layers == 0: + return + if _stats is None: + split = _override[0] if _override is not None else max_rows + _stats = _PlacementStats(float(interval), split, ranks, mapping.tp_rank) + _stats.start() + logger.info( + f"[musespark_open] {_STATS_TAG} INTERVAL_S={interval}: placement census of {layers} " + f"MoE layers, bands up to and above {split} rows, reported every {interval} s by TP " + "rank 0" + ) + _stats.layers = layers diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/router_gemm/router.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/router_gemm/router.py index 8b5bf06dc4f8..286240eea41b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/router_gemm/router.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/router_gemm/router.py @@ -67,7 +67,11 @@ def report(self, hidden_states: torch.Tensor, served: bool) -> None: return self._reported.add((rows, served)) if served: - logger.info(f"[musespark_open] router_gemm: a {rows}-row decode step runs the op") + build = _op.staged() if _op.staged_takes(rows) else None + logger.info( + f"[musespark_open] router_gemm: a {rows}-row decode step runs the op" + + (f" (router_staged_gemm {build}: the staged GEMM)" if build is not None else "") + ) else: logger.warning( f"[musespark_open] router_gemm: an armed {rows}-row router call keeps the built-in " @@ -137,7 +141,9 @@ def install( model: torch.nn.Module, min_rows: int = _op.MIN_ROWS, staged_epilogue: Optional[str] = None ) -> RouterGemmCall: """Re-class every layer's MoE with ``RouterGemmMoE`` in front; returns the call to arm - on decode steps of ``min_rows..MAX_ROWS`` rows. ``staged_epilogue`` is ignored. + on decode steps of ``min_rows..MAX_ROWS`` rows. ``staged_epilogue`` is ignored (this + target selects the staged GEMM per process from the ``router_staged_gemm`` lever + (``ops/moe/router_gemm.configure``), over its own row band). Every layer is checked before any is re-classed; a layer already carrying the mixin is re-pointed at the returned call. @@ -171,4 +177,15 @@ def install( f"(2 x 128-row tiles)" + f") in {len(layers)} MoE layers; other steps keep the built-in router" ) + build = _op.staged() + if build is not None: + # Every rank logs it: the lever is resolved per process (router_staged_gemm). + epilogue, prewait = _op.STAGED_BUILDS[build] + logger.info( + f"[musespark_open] router_staged_gemm {build}: calls of " + f"{max(min_rows, _op.STAGED_MIN_ROWS)}..{_op.STAGED_MAX_ROWS} rows run the staged GEMM " + f"({epilogue} epilogue through shared memory" + + (", weight tiles loaded before griddepcontrol.wait" if prewait else "") + + "; same partials bit for bit)" + ) return call diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/__init__.py new file mode 100644 index 000000000000..0c4008051fb5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""wide_direct_transports: TP8 pool / unpool slice GEMMs of 16..32-row steps publishing into their consumers.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/transports.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/transports.py new file mode 100644 index 000000000000..8794cd93ab37 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/transports.py @@ -0,0 +1,482 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``wide_direct_transports``: on sm103 TEP8 decode steps of two to four requests (16, 24 and 32 +rows at K = 7), the pool and unpool slice GEMMs publish straight into their consumers' transports. + +Off unless ``tep8_wide_direct_transports`` is on (its ``SITES``: both sites, or a +list of ``pool`` / ``unpool``); unset, nothing here is imported and every step runs as before. + +At 9..32 rows ``wide_projection_gathers`` computes each rank's pool (512 columns) and unpool (1024 +columns) slices with flashinfer's split-K into fresh tensors; their consumers -- the pre-expert +RMSNorm + NVFP4 quantization (``pool_gather_fp4quant``) and the post-FFN boundary +(``unpool_gather_boundary``) -- wait on the GEMM, reload the slice, publish it to the seven peers +as tagged 64-bit words through a 32-row inbox, and then poll the peers' words. Here, on the steps of +``ROWS`` rows, the same split-K cell's output tensor is this rank's slot of a multicast transport, +so its epilogue stores are the publication (the 8-row ``unpool_into_boundary`` / +``pool_into_route`` mechanism, and ``projection_publish``'s at 33..512 rows, on the 9..32-row +band neither covers): + +* unpool: a 32-row ``UnpoolDirectWorkspace`` (``unpool_direct_boundary``'s wide layout, which + b5ca973 / #21377 landed in this tree TP4 only; ``check_layout`` here admits TP8 as well), + consumed by ``UnpoolDirectBoundary`` -- the gather boundary's + geometry and arithmetic at 9..32 rows (a cluster of four CTAs, two peers' slices per CTA at TP8), + polling its own copy -- through ``unpool_into_boundary``'s re-classed join, early-release trigger + and DFlash capture fold included. Its sites and buffers are ``unpool_into_boundary``'s. +* pool: a 32-row ``projection_publish.PublishTransport`` of the 512-column slices, consumed by + ``pool_publish_fp4quant`` (flashinfer's RMSNorm + NVFP4 quantization reading its row from the + transport) on the MoE preparation stream, launched with programmatic dependent launch as the + gather consumer it replaces. The publishing layers are those whose pre-expert input is the fused + NVFP4 quantization (``projection_publish``'s choice), fixed at the first eager 16/24/32-row forward + (the engine's CUDA-graph warm-up), once the weights are final, and agreed over the TP group. The + pool site yields to ``tep8_router_shard`` (a TP8 ``_wide_inbox`` on the prologue + MoE): that path's ``_wide_prepare`` owns the 16..32-row pool prepare and feeds the pool slice to + ``pool_gather_route``, which must not be handed a write-only multicast slot. The unpool site + composes with it. + +Both sites arm only on forwards captured into a CUDA graph (as the 8-row ``unpool_into_boundary`` / +``pool_into_route`` and upstream's 16..32-row lever do); eager forwards keep the 32-row inboxes, and +the first eager one at a ``ROWS`` row count only chooses the pool layers. + +Producers: the split-K cells of ``wide_projection_gathers.TP8_WIDE_KERNELS`` (flashinfer's: the +owner CTA of each cluster stores each output element once, from its epilogue, after its dependency +wait). The sites' producer policies are ``unpool_into_boundary``'s and ``pool_into_route``'s (same +families). Slot tails: a 24-row call runs the ``mma_n`` 32 cell into the first 24 rows of a 32-row +slot; its stores past row 24 would be read as arrived by a later 32-row call on that buffer, so the +tail is predicated by the kernel (``exp/mb-tep8-wide-direct-*``: every half past the call's rows +still holds the sentinel after the call). + +Numerics: bitwise. The producer is the served cell with another output pointer, and both consumers +are certified bit for bit against the gather consumers they replace (same thread-to-column maps and +arithmetic, statement for statement). + +Every rank takes the same path at every call: arming reads the step's row count only, the layer +choice is agreed over the TP group, and an armed call that does not fit raises on every rank alike +(its GEMM already published). +""" + +from __future__ import annotations + +import contextlib +from typing import FrozenSet, Iterator, List, Optional, Sequence, Tuple + +import torch + +from tensorrt_llm._torch.distributed.communicator import Distributed +from tensorrt_llm._torch.utils import Fp4QuantizedTensor +from tensorrt_llm.logger import logger + +from ... import dispatch as _tp_agreement +from ...ops.comm import pool_publish_fp4quant as _pool +from ...ops.comm import projection_publish as _publish +from ...ops.comm import unpool_direct_boundary as _direct +from ...ops.comm.projection_inbox import WIDE_MAX_ROWS +from ...ops.gemm.dense_gemm import parse_kernel +from ..dense_gemms.dispatcher import FLASHINFER_FAMILIES, WIDE_MIN_ROWS, claim_publish + +#: The sites the path can publish; ``tep8_wide_direct_transports__SITES`` names a subset. +SITES = ("pool", "unpool") +#: The decode row counts the path arms: two, three and four requests at K = 7 (the 9..32-row +#: steps of the TEP8 generation workers' CUDA graphs); other counts keep the 32-row inboxes. +ROWS = (16, 24, 32) +#: The TP size and device the path serves: sm103 TEP8 over two trays. +TP_SIZE = 8 +CAPABILITY = (10, 3) + + +def requested_sites(sites: Optional[Sequence[str]] = None) -> FrozenSet[str]: + """The sites ``sites`` names, as a set; every site when it is None. + + Raises: + ValueError: for a name outside ``SITES``. + """ + if sites is None: + return frozenset(SITES) + names = frozenset(str(part).strip().lower() for part in sites if str(part).strip()) + unknown = sorted(names - set(SITES)) + if unknown: + raise ValueError( + f"[musespark_open] wide_direct_transports: {unknown}: expected a subset of " + f"{', '.join(SITES)}" + ) + return names + + +def _cells(plan, site: str) -> List[str]: + """The kernel the dense GEMM plan's armed 9..32-row table runs at each of ``ROWS`` for ``site``.""" + table = getattr(plan, "wide_kernels", {}).get(site) + if table is None or len(table) != WIDE_MAX_ROWS - WIDE_MIN_ROWS + 1: + return [] + return [table[rows - WIDE_MIN_ROWS] for rows in ROWS] + + +def refusal(model, plan, wide_projections, unpool_direct, sites: FrozenSet[str], capability) -> Optional[str]: + """Why the path cannot run here, or None. Reads construction-time state only (rank-agnostic).""" + layers = list(model.layers) + mapping = layers[0].mapping + if int(mapping.tp_size) != TP_SIZE: + return f"TP{mapping.tp_size}: the path serves TP{TP_SIZE} (sm103 TEP8) only" + if capability is not None and tuple(capability) != CAPABILITY: + return f"SM {tuple(capability)}: the path serves sm103 only" + if plan is None or wide_projections is None: + return "wide_projection_gathers (and its dense GEMM plan) is not installed" + for site in sorted(sites): + cells = _cells(plan, site) + if len(cells) != len(ROWS): + return f"the dense GEMM plan has no 9..32-row {site} table" + for rows, cell in zip(ROWS, cells): + if parse_kernel(cell)[0] not in FLASHINFER_FAMILIES: + return ( + f"the {rows}-row {site} cell {cell} is not a flashinfer kernel, which alone stores " + "into a transport slot" + ) + if "unpool" in sites and (unpool_direct is None or len(getattr(unpool_direct, "sites", ())) < 2): + return "unpool_into_boundary (whose joins, sites and buffers the unpool site rides) is not installed" + return None + + +def pool_owner_conflict(layers) -> Optional[str]: + """Why another fast path owns the 16..32-row pool prepare of ``layers``, or None. + + ``tep8_router_shard`` (``wide_prologue.install_tep8``) gives the prologue MoE a + TP8 ``_wide_inbox``: its ``_wide_prepare`` then runs the pool slice GEMM itself and hands the + product to ``pool_gather_route`` as its local slice, so that GEMM must not publish into this + path's write-only multicast slot. Reads construction-time state only. + """ + for layer in layers: + moe = layer.feed_forward + if getattr(moe, "_wide_inbox", None) is not None: + return ( + f"layer {getattr(layer, 'layer_idx', '?')}'s 16..32-row pool prepare belongs to the TEP8 " + "router shard (tep8_router_shard, a wide route inbox), which routes the " + "pool slice itself" + ) + return None + + +class WideDirectTransports: + """What ``DecodeTarget`` holds under ``tep8_wide_direct_transports``. + + ``unpool`` is the 32-row direct unpool workspace and ``unpool_sites`` its ``(dispatcher, join, + buffer)`` per fused layer (``unpool_into_boundary``'s); ``pool`` the 32-row pool transport and, + once built, ``pool_layers`` the publishing layers' indices (None: not built yet). + """ + + def __init__( + self, + model, + plan, + wide_projections, + sites: FrozenSet[str], + unpool: Optional[_direct.UnpoolDirectWorkspace], + unpool_sites: Sequence[Tuple[object, object, int]], + pool: Optional[_publish.PublishTransport], + ) -> None: + self._layers = list(model.layers) + self.mapping = self._layers[0].mapping + self.plan = plan + self.wide_projections = wide_projections + self.sites = frozenset(sites) + self.unpool = unpool + self.unpool_sites = tuple(unpool_sites) + self.pool = pool + self.pool_layers: Optional[Tuple[int, ...]] = None + self._pool_refused: Optional[str] = None + self.enable_pdl = bool(getattr(plan, "enable_pdl", True)) + self._armed = False + self._rows = 0 + self._unconsumed = 0 + self._reported = set() + for layer in self._layers: + layer.feed_forward._wide_direct = self + + @property + def is_armed(self) -> bool: + return self._armed + + def _pool_ready(self) -> bool: + return self.pool is not None and self.pool_layers is not None and self._pool_refused is None + + def takes(self, rows: int) -> bool: + """Whether a decode forward of ``rows`` rows publishes on this path. + + Only a forward of ``ROWS`` rows captured into a CUDA graph. An eager one (the warm-up, or a + step without a graph) is never armed; the first eager one at a ``ROWS`` row count chooses + the pool layers (``_build_pool``, collective: ``DecodeTarget`` asks at the same point of the + same forward on every rank). A capture before that arms the unpool site alone. Every rank + warms up and captures the same graphs, so every rank answers alike. + """ + if rows not in ROWS or torch.compiler.is_compiling(): + return False + if not torch.cuda.is_current_stream_capturing(): + if "pool" in self.sites and self.pool is not None and self.pool_layers is None: + self._build_pool() + return False + return ("unpool" in self.sites and self.unpool is not None) or self._pool_ready() + + def _build_pool(self) -> None: + """Choose the publishing pool layers (collective; at the first eager 16/24/32-row forward).""" + from ..pool_into_route.route import PUBLISH_POLICY as POOL_POLICY + from ..projection_gathers.gathers import MuseSparkOpenTepGatherMoE + from ..projection_publish.publish import _pool_args + + chosen: List[int] = [] + for index, layer in enumerate(self._layers): + moe = layer.feed_forward + if isinstance(moe, MuseSparkOpenTepGatherMoE) and _pool_args(moe, WIDE_MAX_ROWS) is not None: + chosen.append(index) + local = tuple(chosen) + everyone = Distributed.get(self.mapping).tp_allgather(local, small_payload=True) + reason = None + if any(tuple(other) != local for other in everyone): + reason = "the TP ranks chose different publishing pool layers" + elif len(chosen) < 2: + reason = f"fewer than two publishing pool layers ({len(chosen)})" + else: + reason = pool_owner_conflict(self._layers[index] for index in chosen) + if reason is None: + for index in chosen: + dispatcher = getattr(self._layers[index].feed_forward.pooler, "_low_m_gemm_dispatcher", None) + held = getattr(dispatcher, "publish_policy", None) + if dispatcher is None or not hasattr(dispatcher, "publish"): + reason = f"layer {index}'s pool projection is not on the dense GEMM dispatcher" + break + if held is not None and held.owner != POOL_POLICY.owner: + reason = f"layer {index}'s pool GEMM already publishes for {held.owner}" + break + reason = _tp_agreement.agreed_refusal(self.mapping, reason) + if reason is not None: + self._pool_refused = reason + self.pool_layers = () + logger.warning( + f"[musespark_open] wide_direct_transports: the pool site is not armed ({reason}); " + f"{'/'.join(map(str, ROWS))}-row steps keep the 32-row pool inbox" + ) + return + for index, buffer in zip(chosen, _publish.buffer_sequence(len(chosen))): + moe = self._layers[index].feed_forward + moe._wide_direct_pool_buffer = buffer + dispatcher = moe.pooler._low_m_gemm_dispatcher + if dispatcher.publish_policy is None: + claim_publish(dispatcher, POOL_POLICY) + self.pool_layers = tuple(chosen) + self.pool.start_guard_watch("wide pool") + logger.info( + f"[musespark_open] wide_direct_transports: TP{TP_SIZE} decode steps of " + f"{'/'.join(map(str, ROWS))} rows publish the pool slice from {len(chosen)} layers into the " + f"32-row pool transport (cells {'|'.join(dict.fromkeys(_cells(self.plan, 'pool')))}; " + f"consumer pool_publish_fp4quant, PDL {'on' if self.enable_pdl else 'off'})" + ) + + @contextlib.contextmanager + def armed(self, rows: int) -> Iterator[None]: + """Arm the requested sites for one decode forward of ``rows`` rows; disarm on any exit. + + Raises: + RuntimeError: if the forward returned with a published pool slot no consumer took. + """ + pool_on = "pool" in self.sites and self._pool_ready() + unpool_on = "unpool" in self.sites and self.unpool is not None + self._rows = int(rows) + self._unconsumed = 0 + if unpool_on: + for dispatcher, join, buffer in self.unpool_sites: + dispatcher.publish = self.unpool.slot(buffer, rows) + join._unpool_direct_armed = (self.unpool, buffer) + self._report("unpool", rows) + self._armed = pool_on + try: + yield + if self._unconsumed: + raise RuntimeError( + f"[musespark_open] wide_direct_transports: a {rows}-row decode forward returned " + f"with {self._unconsumed} published pool slot(s) no consumer took" + ) + finally: + self._armed = False + self._rows = 0 + self._unconsumed = 0 + if unpool_on: + for dispatcher, join, _ in self.unpool_sites: + dispatcher.publish = None + join._unpool_direct_armed = None + + def _report(self, site: str, rows: int) -> None: + if (site, rows) in self._reported: + return + self._reported.add((site, rows)) + cell = _cells(self.plan, site)[ROWS.index(rows)] if rows in ROWS else "?" + logger.info( + f"[musespark_open] wide_direct_transports: a {rows}-row decode step publishes the {site} " + f"slice from its GEMM ({cell}) into the 32-row direct transport" + ) + + def project_pool(self, moe, module, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: + """The pool slot this layer's slice GEMM just published into, or None when the layer does + not publish.""" + buffer = getattr(moe, "_wide_direct_pool_buffer", None) + if buffer is None or module is not moe.pooler: + return None + rows = hidden_states.shape[0] if hidden_states.dim() == 2 else -1 + if rows != self._rows or hidden_states.dtype != torch.bfloat16 or not hidden_states.is_contiguous(): + raise RuntimeError( + f"[musespark_open] wide_direct_transports: armed for {self._rows} rows, the pool " + f"projection got {tuple(hidden_states.shape)} {hidden_states.dtype}" + ) + from tensorrt_llm._torch.modules.low_m_gemm import apply_low_m_gemm + + slot = self.pool.slot(buffer, rows) + columns = self.pool.columns + begin = self.mapping.tp_rank * columns + weight = module.weight[begin : begin + columns] + dispatcher = module._low_m_gemm_dispatcher + dispatcher.publish = slot + try: + projected = apply_low_m_gemm(module, hidden_states, weight, None, force_active=True) + finally: + dispatcher.publish = None + if projected is None or projected.data_ptr() != slot.data_ptr(): + raise RuntimeError( + "[musespark_open] wide_direct_transports: the pool GEMM of an armed layer did not write " + "its transport slot (it bypassed the dense GEMM dispatcher); its peers poll for it" + ) + self._unconsumed += 1 + self._report("pool", rows) + return slot + + def consume_pool(self, moe, pooled: torch.Tensor) -> Optional[Fp4QuantizedTensor]: + """The pre-expert NVFP4 input from the published pool slot ``pooled``, or None when ``pooled`` + is not this layer's slot.""" + buffer = getattr(moe, "_wide_direct_pool_buffer", None) + if buffer is None or not self.pool.is_slot(pooled, buffer, self._rows): + return None + from ..projection_publish.publish import _pool_args + + args = _pool_args(moe, WIDE_MAX_ROWS) + if args is None: + raise RuntimeError( + f"[musespark_open] wide_direct_transports: layer {moe.layer_idx}'s pool slice was " + "published but its pre-expert quantization no longer fits" + ) + gain, global_scale = args + fp4, scale = _pool.pool_publish_rmsnorm_fp4quant( + self.pool, + buffer, + self._rows, + gain, + global_scale, + moe.pre_expert_norm.variance_epsilon, + enable_pdl=self.enable_pdl, + ) + self._unconsumed -= 1 + return Fp4QuantizedTensor(fp4, scale, is_sf_swizzled=False) + + +def _unpool_variants(layers, unpool_sites, captured: FrozenSet[int]) -> list: + """The direct boundary variants the fused layers launch at ``ROWS`` (``unpool_into_boundary``'s + settings per layer, at the wide row counts).""" + index_of = {id(layer.add_ffn): index for index, layer in enumerate(layers)} + keys = set() + for _, join, _ in unpool_sites: + index = index_of[id(join)] + layer = layers[index] + for rows in ROWS: + for capture in (False, True) if index in captured else (False,): + keys.add( + _direct.variant( + rows, + True, + layer.post_ffn_norm.variance_epsilon, + layers[index + 1].pre_qkv_proj_norm.variance_epsilon, + join._preqkv_norm_fp32, + join._boundary_trigger, + capture, + ) + ) + return sorted(keys) + + +def install( + model, + plan, + wide_projections, + unpool_direct, + captured: FrozenSet[int], + capability=None, + sites: Optional[Sequence[str]] = None, +) -> Optional[WideDirectTransports]: + """Build ``sites``' 32-row transports (collective) and hand the arming to the layers. + + ``sites`` is the ``tep8_wide_direct_transports__SITES`` subset the plan resolved; None is + every site. The refusal is agreed over the TP group before anything collective; a refused + path logs a warning and returns None (the steps keep the 32-row inboxes). The unpool + variants are compiled here (agreed), the pool layers chosen at the first eager armed forward. + """ + sites = requested_sites(sites) + if not sites: + return None + layers = list(model.layers) + mapping = layers[0].mapping + reason = _tp_agreement.agreed_refusal( + mapping, refusal(model, plan, wide_projections, unpool_direct, sites, capability) + ) + if reason is not None: + logger.warning( + f"[musespark_open] tep8_wide_direct_transports is on but not armed: {reason}" + ) + return None + unpool = None + unpool_sites: Tuple = () + if "unpool" in sites: + unpool_sites = tuple(unpool_direct.sites) + variants = _unpool_variants(layers, unpool_sites, frozenset(captured)) + unpool = _direct.shared_workspace(mapping, max_rows=WIDE_MAX_ROWS) + failed: Optional[Exception] = None + try: + unpool.build_variants(variants) + except Exception as error: # noqa: BLE001 - every failure must reach the agreement + failed = error + reason = _tp_agreement.agreed_refusal( + mapping, + None if failed is None else f"the wide boundary variants did not build: {failed!r}", + ) + if reason is not None: + raise ValueError(f"[musespark_open] wide_direct_transports: {reason}") from failed + unpool.start_guard_watch() + pool = None + if "pool" in sites: + # Order-independent backstop: _build_pool checks the same again at the first eager step. + conflict = _tp_agreement.agreed_refusal(mapping, pool_owner_conflict(layers)) + if conflict is not None: + logger.warning( + f"[musespark_open] wide_direct_transports: the pool site is not armed ({conflict}); " + f"{'/'.join(map(str, ROWS))}-row steps keep that path's pool prepare" + ) + sites = sites - {"pool"} + if not sites: + return None + if "pool" in sites: + from tensorrt_llm._torch.distributed.ops import _get_mnnvl_workspace_comm + + pool = _publish.PublishTransport( + mapping, + _publish.local_columns(_publish.POOL_HIDDEN, TP_SIZE), + WIDE_MAX_ROWS, + comm=_get_mnnvl_workspace_comm(mapping), + ) + wide = WideDirectTransports(model, plan, wide_projections, sites, unpool, unpool_sites, pool) + logger.info( + f"[musespark_open] tep8_wide_direct_transports ({', '.join(sorted(sites))}): TP{TP_SIZE} " + f"decode steps of {'/'.join(map(str, ROWS))} rows publish the " + f"{' and '.join(sorted(sites))} slice(s) from their split-K GEMMs into 32-row multicast " + "transports" + + ( + f"; unpool on {len(unpool_sites)} layers (cells " + f"{'|'.join(dict.fromkeys(_cells(plan, 'unpool')))}, buffers " + f"{unpool_sites[-1][2] if unpool_sites else '-'} last)" + if unpool is not None + else "" + ) + + ("; pool layers chosen at the first eager 16/24/32-row step" if pool is not None else "") + + "; captured graphs only (eager forwards and other row counts keep the 32-row inboxes)" + ) + return wide diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/__init__.py new file mode 100644 index 000000000000..3c7a442f3092 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""wide_prologue: the TP4 MoE prologue and projection gathers of 16..32-row decode steps.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py new file mode 100644 index 000000000000..c541d665db18 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py @@ -0,0 +1,309 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fast path ``wide_prologue``: the TP4 sharded MoE prologue on decode steps of two to four requests. + +At TP4 a decode step of one request (8 rows at K=7) runs the sharded MoE: the +pool and unpool projections as per-rank column slices whose gathers are fused into +their consumers (``projection_gathers``), and the MoE prologue as one chain +(``moe_prologue``: router shard, pool slice GEMM, ``pool_gather_route``). Without this +fast path a step of two to four requests (16..32 rows) runs the built-in replicated path: +every rank reads the whole 8 MiB router weight with the tensor-core router kernel and then +routes, and multiplies the whole 64 MiB pool and unpool weights, because the +built-in sharded projections, their fused gathers and the projection inboxes all +stop at 8 rows (``ops/comm/projection_inbox.MAX_ROWS``). + +This fast path gives those steps the same structure on wide inboxes +(``ProjectionInbox(max_rows=WIDE_MAX_ROWS)``, separate allocations whose slot layout +holds 32 rows; the 8-row inboxes and kernels are unchanged): + +* the prologue MoE (``moe_prologue.MuseSparkOpenTepPrologueMoE._wide_prepare``) runs + ``ops/moe/router_shard.router_shard_wide`` (this rank's 64 logit columns, bit for + bit the built-in tensor-core router's) on the router fork's side stream, the pool + slice GEMM, and ``pool_gather_route`` on the wide route inbox (routing only: above + 8 rows the experts are the built-in ones and read the routing memo); +* the projection-gather MoE hands its pool and unpool slices on ungathered up to + ``WIDE_MAX_ROWS`` rows (``_defer_max_rows``), and the post-FFN join takes 9..32-row + calls through ``ops/comm/unpool_gather_boundary`` on the wide unpool inbox + (``_wide_gather_inbox``). + +Values: the router logits, the routing and the boundary are the built-in path's bit +for bit given the same projections; the pool and unpool slices are each rank's +columns of the built-in projection by a GEMM of the slice's shape, so they may +differ from the replicated GEMM's in FP32 summation order (each output one BF16 +rounding of its full-K FP32 dot product either way). + +Two levers, given to ``install`` by the plan (the same on every rank): + +* ``wide_router_per_row``: the router shard of these steps is + ``router_shard_rows``, which gives each row its one-request logits (those the built-in + router gives the row when its request decodes alone) in one launch without the split-K + reduce. Those are other bits than the built-in 16..32-row launch's, so the routing of + a wide step may differ from the built-in path's. +* ``wide_direct_transports``: ``pool_into_route`` and ``unpool_into_boundary`` also build their + wide direct transports (the pool and unpool GEMMs and the router shard publish into + the routed kernel's and the boundary's transports on 16..32-row CUDA-graph steps, and on + a half-window layer the routed kernel places). Same bits as the tagged wide path given + the same router logits. + +Every rank takes the same path at every call: the choice reads shapes, dtypes, +contiguity and construction-time state only. + +TEP8 (sm103, two trays), default off: ``install_tep8``, which ``layers.adopt`` calls after +``wide_projection_gathers`` under ``tep8_router_shard`` (with ``tep8_router_shard_wide``), +gives the 16..32-row decode steps the same prologue on a TP8 wide route inbox +(``ProjectionInbox(row_words=row_words(8), max_rows=WIDE_MAX_ROWS)``, 2 MiB of fabric memory +per rank): ``_wide_prepare`` runs this rank's 32 router logit columns on the router fork's side +stream (``router_shard_rows``: each row's one-request bits, the ones its request gets in a +one-request step; with ``wide``, ``router_shard_wide``: the built-in 16..32-row tensor-core +router's bits, two launches and slower at EP8), the pool slice GEMM that +``wide_projection_gathers`` already defers at those rows (its split-K table), and +``pool_gather_route`` (routing only) in place of ``pool_gather_fp4quant``. That replaces the +served chain of those steps -- ``router_gemm``'s BF16x3 split-K GEMM and sum, the built-in +``_full_routing_kernel`` -- with one side-stream launch (two with ``router_shard_wide``'s +reduce) and the routing warps of the pool gather kernel; ``ring_buddy_assign`` then follows +the routed kernel instead of running beside the pool chain. The experts are reached as before: +``replicated_experts``' routing hook places the memo's selection with ``ring_buddy_assign``. +The pool site's unpool side stays ``wide_projection_gathers``'; ``router_gemm`` stays installed +for every other row count and for any declined call. Values: the routing and the quantization +are the built-in kernels' bits on the gathered logits and rows; the logits are not the served +BF16x3 router's, so top-8 near-ties can flip (N-class). +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch + +from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenModel +from tensorrt_llm.logger import logger + +from ... import dispatch as _tp_agreement +from ...ops.comm import pool_gather_route as _route +from ...ops.comm import unpool_gather_boundary as _unpool +from ...ops.comm.projection_inbox import WIDE_MAX_ROWS, ProjectionInbox +from ...ops.moe import router_shard as _shard +from ..moe_prologue.prologue import MuseSparkOpenTepPrologueMoE +from ..projection_gathers.gathers import MuseSparkOpenTepGatedResidual + +#: The TP size the wide inboxes are laid out for: one tray of four. +WORLD = 4 + + +def refusal(model: MuseSparkOpenModel): + """Why ``install`` cannot run on this rank, or None. Reads construction-time state only.""" + layers = list(model.layers) + if not layers: + return "wide_prologue: the model has no layers" + world = layers[0].mapping.tp_size + if world != WORLD: + return f"wide_prologue is laid out for TP{WORLD}, got TP{world}" + for layer in layers: + moe = layer.feed_forward + if type(moe) is not MuseSparkOpenTepPrologueMoE or moe._prologue_inbox is None: + return ( + f"wide_prologue extends the one-chain MoE prologue; layer {layer.layer_idx} has " + f"{type(moe).__name__} without its inbox" + ) + join = layer.add_ffn + if not isinstance(join, MuseSparkOpenTepGatedResidual) or join._gather_inbox is None: + return ( + f"wide_prologue extends the fused unpool gather; layer {layer.layer_idx}'s post-FFN " + "join has no unpool inbox" + ) + if not (_route.available() and _unpool.available()): + return "wide_prologue needs CuTe DSL, flashinfer's RMSNormFP4QuantKernel and the cluster boundary" + return None + + +def install(model: MuseSparkOpenModel, router_rows: bool = False, direct: bool = False) -> int: + """Build the wide route and unpool inboxes and hand them to every layer. + + Called after ``moe_prologue.install`` and ``install_router_fork``. The refusal is + agreed over the TP group before the inboxes (fabric memory, collective) are built. + ``router_rows`` (``wide_router_per_row``) and ``direct`` (``wide_direct_transports``) + are the plan's levers. Returns the number of layers installed. + + Raises: + ValueError: on every rank alike, naming the refusal. + """ + layers = list(model.layers) + mapping = layers[0].mapping + reason = _tp_agreement.agreed_refusal(mapping, refusal(model)) + if reason is not None: + raise ValueError(f"[musespark_open] {reason}") + route_inbox = ProjectionInbox( + mapping, row_words=_route.row_words(WORLD), max_rows=WIDE_MAX_ROWS + ) + unpool_inbox = ProjectionInbox( + mapping, row_words=_unpool.local_columns(WORLD) // 2, max_rows=WIDE_MAX_ROWS + ) + router_rows, direct = bool(router_rows), bool(direct) + for layer in layers: + moe = layer.feed_forward + moe._wide_inbox = route_inbox + moe._defer_max_rows = WIDE_MAX_ROWS + moe._wide_router_rows = router_rows + # NOT ``_wide_direct``: that attribute is the wide direct transport OBJECT, set by + # kernels/wide_direct_transports/transports.install and read in projection_gathers as + # ``.is_armed`` / ``.project_pool`` / ``.consume_pool``. layers.py installs the + # transports first and this prologue second, so assigning the bool here overwrote the + # object and every 8-row decode step raised + # AttributeError("'bool' object has no attribute 'is_armed'") + # at warmup. Nothing reads this flag; it is kept under its own name. + moe._wide_direct_enabled = direct + layer.add_ffn._wide_gather_inbox = unpool_inbox + router = ( + "router_shard_rows, each row's one-request bits" + if router_rows + else "router_shard_wide, the built-in bits" + ) + logger.info( + f"[musespark_open] wide_prologue: TP{WORLD} sharded MoE prologue (router shard of " + f"{_shard.local_experts(WORLD)} experts ({router}), routing in the pool kernel) and fused " + f"unpool gather on decode steps of {_shard.WIDE_MIN_ROWS}..{_shard.WIDE_MAX_ROWS} rows " + f"on all {len(layers)} layers; wide direct transports " + f"{'on' if direct else 'off'} (wide_direct_transports)" + ) + return len(layers) + + +#: The modes of the TEP8 install (``install_tep8``): ``rows`` serves the router shard with +#: ``router_shard_rows``, ``wide`` with ``router_shard_wide``; None keeps the served +#: 16..32-row prologue. The levers are ``tep8_router_shard`` and ``tep8_router_shard_wide``. +TEP8_ROUTER_SHARD_MODES = ("rows", "wide") +#: The TP size and device capability ``install_tep8`` serves. +TEP8_WORLD = 8 +TEP8_CAPABILITY = (10, 3) +#: With the ``tep8_router_shard`` lever in ``rows`` mode, default off: ``1`` makes the router +#: fork's ``router_shard_rows`` launch programmatic (``router_shard_rows(enable_pdl=True)``: the +#: weight tile streams while the attention hop finishes, the rows and the stores follow the hop; +#: ``TEP8_ROUTER_SHARD_EARLY_TRIGGER`` also releases the routed pool kernel at the shard's entry) +#: and launches it with ``TEP8_ROUTER_SHARD_EARLY_TILING``. Same bits. Read once, in +#: ``install_tep8``; ``0``, empty or unset changes nothing; anything else raises on every rank. +TEP8_ROUTER_SHARD_EARLY_ENV = "TRTLLM_MUSESPARK_TEP8_ROUTER_SHARD_EARLY" +#: ``router_shard_rows``' ``rows:per:experts`` tiling under ``TEP8_ROUTER_SHARD_EARLY_ENV``: two rows +#: and two experts per program (128 / 192 / 256 programs at 16 / 24 / 32 rows, 64 registers, one +#: wave beside the pool GEMM). Chosen on GB300 in the served 62-layer TEP8 chain (8 ranks, two tray +#: pairs), paired against the plain launch: -10..-15 / -48..-62 / -104..-113 us per step at 16 / 24 / +#: 32 rows, the fastest at 32 rows and within the spread of the fastest at 16 and 24; one expert per +#: program saves less at 32 rows (-53 us), four save less (-59 / -47 us), and the same tilings as +#: plain launches lose or tie (+9 / -68 / +39 us for this one). +TEP8_ROUTER_SHARD_EARLY_TILING = "16:2:2,24:2:2,32:2:2" +#: Whether the early router shard releases its dependents at entry (``"entry"``) or at exit +#: (``"wait"``: only its weight tile is early). ``"wait"`` measured slower than the plain launch +#: (+43..+56 us per step): the routed pool kernel then launches only once the shard has exited. +TEP8_ROUTER_SHARD_EARLY_TRIGGER = "entry" + + +def tep8_router_shard_early(value: Optional[str] = None) -> bool: + """``TRTLLM_MUSESPARK_TEP8_ROUTER_SHARD_EARLY`` (or ``value``) as a bool. + + Raises: + ValueError: for a value other than unset, empty, ``0`` or ``1``. + """ + raw = (os.environ.get(TEP8_ROUTER_SHARD_EARLY_ENV, "") if value is None else value).strip() + if raw in ("", "0"): + return False + if raw == "1": + return True + raise ValueError(f"[musespark_open] {TEP8_ROUTER_SHARD_EARLY_ENV}={raw!r}: expected 0 or 1") + + +def tep8_router_shard_mode(wide: bool = False) -> str: + """The install mode the ``tep8_router_shard`` levers select: ``"wide"`` or ``"rows"``.""" + return "wide" if wide else "rows" + + +def tep8_refusal(model: MuseSparkOpenModel) -> Optional[str]: + """Why ``install_tep8`` cannot run on this rank, or None. Reads construction-time state only.""" + layers = list(model.layers) + if not layers: + return "the model has no layers" + world = layers[0].mapping.tp_size + if world != TEP8_WORLD: + return f"the TEP8 router shard is laid out for TP{TEP8_WORLD}, got TP{world}" + if not torch.cuda.is_available(): + return "no CUDA device" + capability = tuple(torch.cuda.get_device_capability()) + if capability != TEP8_CAPABILITY: + return f"the TEP8 router shard is validated on sm103 only, got {capability}" + for layer in layers: + moe = layer.feed_forward + if not isinstance(moe, MuseSparkOpenTepPrologueMoE) or moe._prologue_inbox is None: + return ( + f"the TEP8 router shard extends the one-chain MoE prologue; layer " + f"{layer.layer_idx} has {type(moe).__name__} without its inbox" + ) + if moe._router_stream is None: + return f"the TEP8 router shard runs on the router fork; layer {layer.layer_idx} has none" + if getattr(moe, "_wide", None) is None or moe._wide_pool_inbox is None: + return ( + "the TEP8 router shard takes the pool slices wide_projection_gathers defers; layer " + f"{layer.layer_idx} has no wide pool site" + ) + if moe._wide_inbox is not None: + return f"layer {layer.layer_idx} already has a wide route inbox" + if not _route.available(): + return "the TEP8 router shard needs CuTe DSL and flashinfer's RMSNormFP4QuantKernel" + return None + + +def install_tep8(model: MuseSparkOpenModel, mode: str) -> int: + """The TEP8 router shard (``tep8_router_shard``): give every prologue MoE + the TP8 wide route inbox, so ``_wide_prepare`` serves its 16..32-row decode steps. + + Called after ``moe_prologue.install``, ``install_router_fork`` and + ``wide_projection_gathers.install``. The refusal is agreed over the TP group before the + inbox (fabric memory, collective) is built; a refused install logs a warning on every + rank and changes nothing. Returns the number of layers given the inbox (0 when refused). + """ + if mode not in ("wide", "rows"): + raise ValueError(f"[musespark_open] install_tep8: mode {mode!r} is not rows or wide") + layers = list(model.layers) + mapping = layers[0].mapping if layers else None + reason = _tp_agreement.agreed_refusal(mapping, tep8_refusal(model)) + if reason is not None: + logger.warning( + "[musespark_open] tep8_router_shard is on, but the TEP8 router shard is not " + f"installed: {reason}; 16..32-row decode steps keep the served router" + ) + return 0 + early = mode == "rows" and tep8_router_shard_early() + if early: + _shard.configure(rows_tiling=TEP8_ROUTER_SHARD_EARLY_TILING) + route_inbox = ProjectionInbox( + mapping, row_words=_route.row_words(TEP8_WORLD), max_rows=WIDE_MAX_ROWS + ) + for layer in layers: + moe = layer.feed_forward + moe._wide_inbox = route_inbox + moe._wide_router_rows = mode == "rows" + moe._wide_router_early = TEP8_ROUTER_SHARD_EARLY_TRIGGER if early else None + router = ( + "router_shard_rows: each row's one-request bits" + if mode == "rows" + else "router_shard_wide: the built-in tensor-core router's bits" + ) + logger.info( + f"[musespark_open] tep8_router_shard ({mode}): TP{TEP8_WORLD} " + f"decode steps of {_shard.WIDE_MIN_ROWS}..{_shard.WIDE_MAX_ROWS} rows on {len(layers)} " + f"layers run this rank's {_shard.local_experts(TEP8_WORLD)} router logit columns " + f"({router}) on the router fork and the logit exchange + routing in pool_gather_route " + f"on a {WIDE_MAX_ROWS}-row TP{TEP8_WORLD} inbox, in place of router_gemm and the built-in " + "routing; the BF16-expert layers and every other row count keep the served path" + ) + if early: + logger.info( + f"[musespark_open] tep8 router shard: {TEP8_ROUTER_SHARD_EARLY_ENV}=1: router_shard_rows " + "on the router fork is a programmatic launch (weight tile before its dependency wait, " + f"dependents released at {'entry' if TEP8_ROUTER_SHARD_EARLY_TRIGGER == 'entry' else 'exit'}" + f"), tiling rows:per:experts {TEP8_ROUTER_SHARD_EARLY_TILING}; same bits" + ) + elif tep8_router_shard_early(): + logger.warning( + f"[musespark_open] {TEP8_ROUTER_SHARD_EARLY_ENV}=1 needs the tep8_router_shard lever " + f"in rows mode (router_shard_rows); with {mode} it changes nothing" + ) + return len(layers) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py index 6d37ed96839f..1d5ad30b9ab0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py @@ -946,6 +946,7 @@ class AdoptedLayers: published_projections: Optional[object] = None finalize_two_shot: Optional[object] = None batch_dense_gemms: Optional[object] = None + wide_direct: Optional[object] = None def adopt_model( @@ -1278,6 +1279,22 @@ def adopt_model( model.layers[0].mapping, lambda: _wide_projections.check(model, dense_gemm_plan) ) wide_projections = _wide_projections.install(model, dense_gemm_plan) + wide_direct = None + if wide_projections is not None and plan.enabled("tep8_wide_direct_transports"): + # Default off. After wide_projection_gathers (whose 9..32-row table and arming it rides) + # and unpool_into_boundary (whose joins, sites and producer policy its unpool site + # reuses). Collective: its refusal is agreed over the TP group first. + from .kernels.wide_direct_transports import transports as _wide_direct + + wide_direct = _wide_direct.install( + model, + dense_gemm_plan, + wide_projections, + unpool_direct, + capture_folded, + compute_capability, + sites=plan.param("tep8_wide_direct_transports", "SITES"), + ) published_projections = None if plan.enabled("projection_publish"): # After projection_gathers, whose MoE and join classes carry its hooks. Its refusal is @@ -1289,6 +1306,19 @@ def adopt_model( model.layers[0].mapping, lambda: _projection_publish.check(model, layer_kwargs) ) published_projections = _projection_publish.install(model) + if plan.enabled("tep8_router_shard"): + # Default off (TP8 sm103 gens): after moe_prologue, the router fork and + # wide_projection_gathers (whose deferred pool slices it routes), the 16..32-row prologue + # on a TP8 wide route inbox. Its refusal (another TP size or SM, a missing prerequisite) + # is agreed over the TP group and only warns. Off: nothing is imported or built. + # Placed after projection_publish -- which only adds hooks to the same MoE classes -- + # rather than beside wide_projection_gathers, to leave that edit site free. + from .kernels.wide_prologue import prologue as _wide_prologue + + _wide_prologue.install_tep8( + model, + _wide_prologue.tep8_router_shard_mode(plan.enabled("tep8_router_shard_wide")), + ) batch_dense_gemms = None if plan.enabled("batch_dense_gemms"): # After dense_gemms (whose re-classed projections and plan it arms) and every @@ -1315,4 +1345,5 @@ def adopt_model( published_projections=published_projections, finalize_two_shot=finalize_two_shot, batch_dense_gemms=batch_dense_gemms, + wide_direct=wide_direct, ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py new file mode 100644 index 000000000000..c3f810824ee5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The plan's lever decisions, pushed into the kernel modules that read them at launch. + +Several kernel modules of this target hold module-level settings their kernels read when they +build or launch (a row bound, a transport layout, whether a launch releases its dependents at +entry). Each module declares them with the default of the path it refines and one keyword-only +``configure``; ``apply`` calls every ``configure`` once per process, from the model constructor, +right after ``dispatch.compose`` and before any installer runs, with what the plan resolved +(``dispatch.Plan.enabled``, ``dispatch.Plan.param``). Under ``NO_PATHS`` it changes nothing, so a +deployment without a composition keeps every default. + +A module whose settings all keep their defaults is not imported for it: some kernel modules load +a compiler, or change what a worker does, by being imported. + +The levers that do not reach their site this way are read at their install instead, where the +plan is already in scope: ``tep8_router_shard`` and ``tep8_wide_direct_transports`` +(``layers.adopt_model``), ``tep8_persistent_batch_experts``, ``tep8_experts_early_start`` and +``tep8_wide_direct_pool`` (``weights.install_loaded_features``), and ``drafter_fused_prep`` +(``drafter``'s ``context_kv_fork`` install). +""" + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from . import dispatch as _dispatch + + +def _configure(module: str, settings: dict, defaults: dict) -> None: + """Import ``module`` (relative to this target) and call its ``configure`` with ``settings``, + unless every setting is its default: then the module is not imported for it, since some + kernel modules load their compilers or change a worker's behaviour by being imported.""" + if settings == defaults: + return + importlib.import_module(f"{__package__}.{module}").configure(**settings) + + +def _router_staged_build(plan: "_dispatch.Plan") -> Optional[str]: + """The staged router GEMM build ``router_staged_gemm__BUILD`` names, or None when the lever + is off. + + Raises: + ValueError: when ``BUILD`` names more than one build (the knob takes a subset of + ``levers.ROUTER_STAGED_BUILDS``; exactly one build runs). + """ + if not plan.enabled("router_staged_gemm"): + return None + build = tuple(plan.param("router_staged_gemm", "BUILD")) + if len(build) != 1: + from . import levers as _levers + + raise ValueError( + f"{_levers.knob_name('router_staged_gemm', 'BUILD')}={','.join(build)}: name exactly " + f"one of {list(_levers.ROUTER_STAGED_BUILDS)}" + ) + return build[0] + + +def apply(plan: "_dispatch.Plan") -> None: + """Configure the kernel modules from ``plan``, each with one ``configure`` call. + + A no-op for ``NO_PATHS``: every setting keeps its default. A module whose settings all keep + their defaults is not imported. + + Raises: + ValueError: for a lever parameter a module refuses, named with its knob. + """ + if not plan: + return + on = plan.enabled + # ``finalize_two_shot`` and ``wide_drafter_hops`` run at TP4 as well, but their parameters + # here are TP8 row bounds the modules ignore at any other size. Pushing them anyway would + # import those modules at construction for nothing, and both load the CuTe DSL by being + # imported, so at TP4 they keep the module default and the module is left alone. + tp = plan.levers.deployment.tp if plan.levers is not None else None + tep8 = tp is None or tp == 8 + fc_wide = on("tep8_fc_shard_wide") + two_shot = on("finalize_two_shot") and tep8 + wide_hops = on("wide_drafter_hops") and tep8 + transport_wide = on("tep8_finalize_transport_wide") + rb_wide = on("tep8_ring_buddy_wide") + rb_balance = on("tep8_wide_balance") + rb_stats = on("ring_buddy_stats") + settings = ( + # Dense GEMMs. + ( + "kernels.dense_gemms.dispatcher", + dict(batch_splitk=on("tep8_batch_splitk")), + dict(batch_splitk=False), + ), + # Attention. + ( + "ops.attention.multi_request_split_kv", + dict(combine_gate=on("multi_request_combine_gate")), + dict(combine_gate=False), + ), + # Experts. ``ring_buddy.configure`` also hands the wide band to ``ring_buddy_assign``, + # which owns it, so both modules stay in step from this one call. + ( + "kernels.replicated_experts.ring_buddy", + dict( + in_kernel_placement=on("tep8_experts_in_kernel_placement"), + tep8_max_rows=plan.param("tep8_ring_buddy_wide", "MAX_ROWS") if rb_wide else None, + wide_balance_rows=( + plan.param("tep8_wide_balance", "MAX_ROWS") if rb_balance else None + ), + stats_interval_s=( + plan.param("ring_buddy_stats", "INTERVAL_S") if rb_stats else None + ), + ), + dict( + in_kernel_placement=False, + tep8_max_rows=None, + wide_balance_rows=None, + stats_interval_s=None, + ), + ), + ( + "ops.moe.persistent_experts", + dict(interleaved_fetch=on("tep8_experts_interleaved_fetch")), + dict(interleaved_fetch=False), + ), + # Routing. + ( + "ops.moe.router_gemm", + dict(staged=_router_staged_build(plan)), + dict(staged=None), + ), + # Communication and step glue. + ( + "ops.moe.finalize_allreduce_transport", + dict( + tep8_max_tokens=( + plan.param("tep8_finalize_transport_wide", "MAX_TOKENS") + if transport_wide + else 16 + ), + trigger="entry" if on("finalize_entry_trigger") else "late", + ), + dict(tep8_max_tokens=16, trigger="late"), + ), + ( + "kernels.finalize_two_shot.install", + dict( + tep8_max_rows=( + plan.param("finalize_two_shot", "TP8_MAX_ROWS") if two_shot else None + ) + ), + dict(tep8_max_rows=None), + ), + # Drafter. + ( + "ops.allreduce.drafter_hop", + dict( + tep8_wide_rows=( + plan.param("wide_drafter_hops", "TP8_MAX_ROWS") if wide_hops else 64 + ) + ), + dict(tep8_wide_rows=64), + ), + ( + "spec_worker", + dict( + fc_shard_tp8_max_rows=plan.param("tep8_fc_shard_wide", "MAX_ROWS") if fc_wide else 8 + ), + dict(fc_shard_tp8_max_rows=8), + ), + ) + for module, values, defaults in settings: + _configure(module, values, defaults) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py index e4aee71bf58f..0443bc39e71f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py @@ -82,7 +82,62 @@ SOURCES = ("structural", "dependency", "knob", "baseline", "auto", "gate") #: Retired switches and what replaced them. Setting one, to any value, raises. -RETIRED: Mapping[str, str] = types.MappingProxyType({}) +#: +#: Every switch of the sm103 decode stack became a lever; a deployment that still exports one +#: of these names is told which knob to set instead, rather than running with the switch +#: silently ignored. A row-bound or mode switch names its lever's parameter knob, which has +#: no effect on its own: set the lever too. +RETIRED: Mapping[str, str] = types.MappingProxyType( + { + "TRTLLM_MUSESPARK_TEP8_FC_SHARD_MAX_ROWS": ( + PREFIX + "TEP8_FC_SHARD_WIDE=on, with its MAX_ROWS" + ), + "MUSESPARK_OPEN_TEP8_PERSISTENT_EXPERTS_MAX_ROWS": ( + PREFIX + "TEP8_PERSISTENT_BATCH_EXPERTS=on, with its MAX_ROWS" + ), + "TRTLLM_MUSESPARK_TEP8_EXPERTS_EARLY_START": PREFIX + "TEP8_EXPERTS_EARLY_START=on", + "MUSESPARK_OPEN_TEP8_ROUTER_SHARD": ( + PREFIX + "TEP8_ROUTER_SHARD=on (and " + PREFIX + "TEP8_ROUTER_SHARD_WIDE=on for wide)" + ), + "MUSESPARK_OPEN_TEP8_WIDE_DIRECT_TRANSPORTS": ( + PREFIX + "TEP8_WIDE_DIRECT_TRANSPORTS=on, with its SITES" + ), + "MUSESPARK_OPEN_TEP8_WIDE_DIRECT_POOL": PREFIX + "TEP8_WIDE_DIRECT_POOL=on", + "MUSESPARK_OPEN_DRAFTER_WIDE_HOPS": ( + PREFIX + "WIDE_DRAFTER_HOPS=on, with its TP8_MAX_ROWS" + ), + "TRTLLM_MUSESPARK_TEP8_FINALIZE_TWO_SHOT": PREFIX + "FINALIZE_TWO_SHOT=on", + "TRTLLM_MUSESPARK_TEP8_FINALIZE_TWO_SHOT_MAX_ROWS": ( + PREFIX + "FINALIZE_TWO_SHOT=on, with its TP8_MAX_ROWS" + ), + "TRTLLM_MUSESPARK_TEP8_FINALIZE_TRANSPORT_MAX_TOKENS": ( + PREFIX + "TEP8_FINALIZE_TRANSPORT_WIDE=on, with its MAX_TOKENS" + ), + "TRTLLM_MUSESPARK_FINALIZE_TRANSPORT_TRIGGER": PREFIX + "FINALIZE_ENTRY_TRIGGER=on", + "TRTLLM_MUSESPARK_TEP8_BATCH_SPLITK": PREFIX + "TEP8_BATCH_SPLITK=on", + "TRTLLM_MUSESPARK_MULTI_REQUEST_COMBINE_GATE": (PREFIX + "MULTI_REQUEST_COMBINE_GATE=on"), + "TRTLLM_MUSESPARK_ROUTER_STAGED_GEMM": (PREFIX + "ROUTER_STAGED_GEMM=on, with its BUILD"), + "MUSESPARK_OPEN_DRAFTER_FUSED_PREP": PREFIX + "DRAFTER_FUSED_PREP=on", + "MUSESPARK_OPEN_DRAFTER_FUSED_PREP_MIN_REQUESTS": ( + PREFIX + "DRAFTER_FUSED_PREP=on, with its MIN_REQUESTS" + ), + "MUSESPARK_OPEN_RING_BUDDY_MAX_ROWS": ( + PREFIX + "TEP8_RING_BUDDY_WIDE=on, with its MAX_ROWS" + ), + "TRTLLM_MUSESPARK_TEP8_WIDE_BALANCE_ROWS": ( + PREFIX + "TEP8_WIDE_BALANCE=on, with its MAX_ROWS" + ), + "MUSESPARK_OPEN_RING_BUDDY_STATS": ( + PREFIX + "RING_BUDDY_STATS=on, with its INTERVAL_S (whole seconds)" + ), + "TRTLLM_MUSESPARK_TEP8_EXPERTS_IN_KERNEL_PLACEMENT": ( + PREFIX + "TEP8_EXPERTS_IN_KERNEL_PLACEMENT=on" + ), + "TRTLLM_MUSESPARK_TEP8_EXPERTS_INTERLEAVED_FETCH": ( + PREFIX + "TEP8_EXPERTS_INTERLEAVED_FETCH=on" + ), + } +) #: Name prefixes only retired switches carry; any other name under them raises too. RETIRED_PREFIXES: Mapping[str, str] = types.MappingProxyType({}) @@ -384,6 +439,37 @@ def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, st return rule +def _knob_only(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + """Off under ``auto``: on only by its knob, until this target gates it.""" + return False, "flagged: on only by its knob, until it is gated on this target" + + +def _tep4_only(what: str) -> RuleFn: + """On at TP4, where the composition carries the path; off at any other size by its knob only. + + The two paths that read it (``wide_drafter_hops``, ``finalize_two_shot``) are the sm103 TP4 + composition's own, served there under ``auto`` as before; the wider composition serves them + too, but leaves them to their knob. + """ + + def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + if d.tp == 4: + return True, what + return False, f"tep{d.tp}: on only by its knob (the tep4 composition's path)" + + return rule + + +def _ported(kind: Kind, params: Tuple[Param, ...] = ()) -> Rule: + """A switch of the sm103 decode stack, ported from an environment variable: default off. + + Every one of them gates a decode step of at most a few hundred rows, so the serving region + is ``LL``; none is on under ``auto`` or ``baseline`` (``_knob_only``, and a non-upstream + origin), so a deployment that sets no knob runs exactly the paths it ran before. + """ + return Rule("LL", kind, _knob_only, params=params) + + def _hot_replicas(requires_decode: bool) -> RuleFn: def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: if requires_decode and not d.decodes(): @@ -398,16 +484,61 @@ def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, st LL_MAX_BATCH = 8 -#: The TP4 paths: on wherever served under ``auto``. +#: The TP4 paths: on wherever served under ``auto``. ``wide_drafter_hops`` was one until the +#: wider composition began to serve it too; it has its own rule now (``_tep4_only``). TEP4_PATHS = ( "multi_request_prep_append", "multi_request_split_kv", "multi_request_post_attention_hop", - "wide_drafter_hops", "router_gemm", "projection_publish", ) +# -------------------------------------------------------------------------- +# The parameter defaults of the ported switches +# -------------------------------------------------------------------------- +# Each is the value its environment variable named, or the bound the module it configures +# already serves. They are literals because this resolver imports no module of the target. +# +#: ``wide_drafter_hops__TP8_MAX_ROWS``: the rows the sm103 tep4 composition serves +#: (``ops/allreduce/drafter_hop.WIDE_SERVED_ROWS_BY_SM[(10, 3)]``), the drafter blocks of +#: 2..7 requests at K = 7. +WIDE_DRAFTER_HOPS_TP8_ROWS = 56 +#: ``finalize_two_shot__TP8_MAX_ROWS``: the last row of the tep8 band as +#: ``kernels/finalize_two_shot/install.ROWS_BY_SM_TP[(103, 8)]`` declares it (19 requests). +FINALIZE_TWO_SHOT_TP8_ROWS = 152 +#: ``tep8_fc_shard_wide__MAX_ROWS``: the five-request decode step (40 rows at K = 7) the +#: sharded ``fc``'s note names; its capacity is 128 (``spec_worker.FC_SHARD_TP8_ROW_CAPACITY``). +FC_SHARD_TP8_ROWS = 40 +#: ``tep8_persistent_batch_experts__MAX_ROWS``: the largest stacked call the persistent +#: kernel's batch builds serve (``ops/moe/persistent_experts.BATCH_MAX_ROWS``). +PERSISTENT_BATCH_ROWS = 32 +#: ``tep8_wide_direct_transports__SITES``: the sites the path can publish +#: (``kernels/wide_direct_transports/transports.SITES``); both, as its ``1`` meant. +WIDE_DIRECT_SITES = ("pool", "unpool") +#: ``tep8_finalize_transport_wide__MAX_TOKENS``: the widest TP8 layout +#: (``ops/moe/finalize_allreduce_transport.TP8_WIDE_MAX_TOKENS``), five requests at K = 7. +FINALIZE_TRANSPORT_TOKENS = 40 +#: ``router_staged_gemm__BUILD``: the staged GEMM's builds and the one its ``1`` selected +#: (``ops/moe/router_gemm.STAGED_BUILDS``, ``STAGED_DEFAULT``). Exactly one is named. +ROUTER_STAGED_BUILDS = ("bulk", "v4", "bulk-prewait", "v4-prewait") +ROUTER_STAGED_BUILD = "bulk-prewait" +#: ``drafter_fused_prep__MIN_REQUESTS``: ``ops/spec/drafter_step_prep.DEFAULT_MIN_REQUESTS``. +DRAFTER_PREP_MIN_REQUESTS = 2 +#: ``tep8_ring_buddy_wide__MAX_ROWS``: the stacked path's row cap on the tep8 ring, replacing +#: ``kernels/replicated_experts/ring_buddy.MAX_ROWS`` (16). 128 is ``8 * B`` for a 16-request +#: verify step at K = 7; the ceiling is one ``ring_buddy_assign`` program's rows +#: (``ops/moe/ring_buddy_assign.MAX_SLOTS // TOP_K``), which ``install`` refuses to exceed. +RING_BUDDY_MAX_ROWS = 128 +RING_BUDDY_ROW_CEILING = 256 +#: ``tep8_wide_balance__MAX_ROWS``: the band the balanced NVFP4 placement adds above the +#: stacked cap, from ``ring_buddy_assign.WIDE_BALANCE_MIN_ROWS`` (129) to the same ceiling. +WIDE_BALANCE_ROWS = 256 +WIDE_BALANCE_MIN_ROWS = 129 +#: ``ring_buddy_stats__INTERVAL_S``: seconds between the placement-census reports TP rank 0 +#: logs. Whole seconds only (the knob grammar has no fractional parameter). +RING_BUDDY_STATS_INTERVAL_S = 60 + def _tep4_path_rule() -> Rule: return Rule( @@ -459,6 +590,64 @@ def _vocab_parallel_head(d: Deployment, t: Resolving, p: Mapping[str, object]) - Param("MIN_ROWS", 16, 8, 16), ), ), + # The two tep4 composition paths the wider composition also serves, on its knob only. + # ``TP8_MAX_ROWS`` is read only where the tep8 workspace is built; at tep4 each path + # keeps the rows its own composition serves. + "wide_drafter_hops": Rule( + "upstream", + Kind.COMPOSITION, + _tep4_only("9..56-row block decodes"), + params=(Param("TP8_MAX_ROWS", WIDE_DRAFTER_HOPS_TP8_ROWS, 9, 64),), + ), + "finalize_two_shot": Rule( + "upstream", + Kind.COMPOSITION, + _tep4_only("72..512-row decode steps"), + params=(Param("TP8_MAX_ROWS", FINALIZE_TWO_SHOT_TP8_ROWS, 24, 152),), + ), + # The sm103 decode-stack switches, ported from environment variables: default off. + "tep8_fc_shard_wide": _ported( + Kind.FLAGGED, (Param("MAX_ROWS", FC_SHARD_TP8_ROWS, 9, 128),) + ), + "tep8_persistent_batch_experts": _ported( + Kind.FLAGGED, (Param("MAX_ROWS", PERSISTENT_BATCH_ROWS, 9, 32),) + ), + "tep8_experts_early_start": _ported(Kind.FLAGGED), + "tep8_router_shard": _ported(Kind.FLAGGED), + "tep8_router_shard_wide": _ported(Kind.FLAGGED), + "tep8_wide_direct_pool": _ported(Kind.FLAGGED), + "tep8_wide_direct_transports": _ported( + Kind.FLAGGED, (Param("SITES", WIDE_DIRECT_SITES, choices=WIDE_DIRECT_SITES),) + ), + "tep8_batch_splitk": _ported(Kind.BITWISE), + "finalize_entry_trigger": _ported(Kind.BITWISE), + "tep8_finalize_transport_wide": _ported( + Kind.FLAGGED, (Param("MAX_TOKENS", FINALIZE_TRANSPORT_TOKENS, 17, 40),) + ), + "multi_request_combine_gate": _ported(Kind.BITWISE), + "router_staged_gemm": _ported( + Kind.FLAGGED, + (Param("BUILD", (ROUTER_STAGED_BUILD,), choices=ROUTER_STAGED_BUILDS),), + ), + "drafter_fused_prep": _ported( + Kind.FLAGGED, (Param("MIN_REQUESTS", DRAFTER_PREP_MIN_REQUESTS, 1, 64),) + ), + "tep8_ring_buddy_wide": _ported( + Kind.FLAGGED, (Param("MAX_ROWS", RING_BUDDY_MAX_ROWS, 1, RING_BUDDY_ROW_CEILING),) + ), + "tep8_wide_balance": _ported( + Kind.FLAGGED, + (Param("MAX_ROWS", WIDE_BALANCE_ROWS, WIDE_BALANCE_MIN_ROWS, RING_BUDDY_ROW_CEILING),), + ), + # Diagnostics: one extra assign launch per placement, and the in-kernel placement + # stands aside for it (``ring_buddy``'s ``_stats is None``), so it is never on by + # itself under ``auto``. + "ring_buddy_stats": _ported( + Kind.FLAGGED, + (Param("INTERVAL_S", RING_BUDDY_STATS_INTERVAL_S, 1, 86400),), + ), + "tep8_experts_in_kernel_placement": _ported(Kind.FLAGGED), + "tep8_experts_interleaved_fetch": _ported(Kind.FLAGGED), **{name: _tep4_path_rule() for name in TEP4_PATHS}, } ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py index 3bd4e4864ea8..3728a13742ef 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py @@ -414,6 +414,7 @@ class acts only on sliding-window layers): their decode prep also writes the from . import dispatch as _dispatch from . import layers as _layers +from . import lever_settings as _lever_settings from . import weights as _weights from .attention import DecodeAttentionPlan from .layers import _NO_ALLREDUCE, PostAttnHopCall @@ -447,6 +448,8 @@ class MuseSparkOpenTepModelingV2(MuseSparkOpenModel): _finalize_two_shot: Optional[object] _wide_hop: Optional[object] _hot_replicas: Optional[object] = None + #: ``tep8_wide_direct_transports`` (``kernels/wide_direct_transports``), or None. + _wide_direct: Optional[object] = None @classmethod def adopt( @@ -489,6 +492,7 @@ def adopt( model._wide_hop = adopted.wide_hop model._wide_projections = adopted.wide_projections model._published_projections = adopted.published_projections + model._wide_direct = adopted.wide_direct model._batch_dense_gemms = adopted.batch_dense_gemms model._prefill = PrefillTarget(model, _dispatch.sm_of(compute_capability)) model._decode = DecodeTarget(model) @@ -734,6 +738,7 @@ def __init__(self, core: MuseSparkOpenTepModelingV2) -> None: self.wide_projections = getattr(core, "_wide_projections", None) self.published_projections = getattr(core, "_published_projections", None) self.batch_dense_gemms = getattr(core, "_batch_dense_gemms", None) + self.wide_direct = getattr(core, "_wide_direct", None) def forward( self, @@ -798,6 +803,10 @@ def forward( wide_ar_armed = wide_ar is not None and wide_ar.takes(rows) wide_proj = self.wide_projections wide_proj_armed = wide_proj is not None and wide_proj.takes(rows) + # tep8_wide_direct_transports: 16/24/32-row steps of the wide projection + # chain publish the pool / unpool slices from their GEMMs (kernels/wide_direct_transports). + wide_direct = self.wide_direct + wide_direct_armed = wide_proj_armed and wide_direct is not None and wide_direct.takes(rows) published = self.published_projections published_armed = published is not None and published.takes(rows) two_shot = self.finalize_two_shot @@ -824,6 +833,7 @@ def forward( kv_splits.armed() if kv_splits_armed else contextlib.nullcontext(), wide_ar.armed() if wide_ar_armed else contextlib.nullcontext(), wide_proj.armed() if wide_proj_armed else contextlib.nullcontext(), + wide_direct.armed(rows) if wide_direct_armed else contextlib.nullcontext(), published.armed(rows) if published_armed else contextlib.nullcontext(), two_shot.armed() if two_shot_armed else contextlib.nullcontext(), hot.armed() if hot_armed else contextlib.nullcontext(), @@ -875,30 +885,39 @@ def forward( #: verify head of 16..1024-row decode steps on each rank's own vocabulary shard over the #: tray's multicast, and the drafter hops of 9..56-row block decodes on a wide one-shot #: transport. +#: ``wide_drafter_hops`` and ``finalize_two_shot`` left this set when they became levers of +#: both compositions (``levers._tep4_only``: on at tep4 as before, on its knob at tep8). +#: The levers of the tep4 paths above are listed with them: the multi-request combine's +#: norm-gate, the staged router GEMM and the wide direct transports of the sharded pool and +#: unpool slices (whose ``wide_projection_gathers`` no tep8 composition serves). _SM103_TEP4_ONLY_PATHS = frozenset( { "hot_replicas", "multi_request_prep_append", "drafter_window_compact", "router_gemm", + "router_staged_gemm", "multi_request_split_kv", + "multi_request_combine_gate", "decode_kv_splits", "wide_moe_allreduce", "multi_request_post_attention_hop", "wide_projection_gathers", + "tep8_wide_direct_transports", "projection_publish", "context_sequence_parallel", - "finalize_two_shot", "batch_dense_gemms", "wide_drafter_window_attention", "wide_draft_topk", "wide_split_sampler", "context_prep", "vocab_parallel_head", - "wide_drafter_hops", } ) -#: TEP8: every base fast path (``dispatch.BASE_PATH_NAMES``) but the TP4-only ones. +#: TEP8: every base fast path (``dispatch.BASE_PATH_NAMES``) but the TP4-only ones. The +#: levers of the sm103 decode stack are base paths too, so this serves them; each is off +#: until its knob is set (``levers.RULES``), ``wide_drafter_hops`` and ``finalize_two_shot`` +#: among them. _TEP8_PATHS = _dispatch.BASE_PATH_NAMES - _TEP4_ONLY_PATHS - _SM103_TEP4_ONLY_PATHS #: The fast paths written for any rank count whose output is bitwise the built-in #: model's. @@ -973,6 +992,13 @@ def forward( "capture_fold", "moe_prologue", "finalize_transport", + # The two paths this composition has always carried; the wider one serves them + # too now, on its knob (``levers._tep4_only``). + "wide_drafter_hops", + "finalize_two_shot", + # The levers both compositions serve, each off until its knob is set. + "finalize_entry_trigger", + "drafter_fused_prep", } ) ) @@ -988,6 +1014,12 @@ def forward( #: DFlash drafting a linear chain of ``dispatch.DRAFT_LEN`` tokens on an NVFP4 #: checkpoint; every other deployment runs the built-in modules only. The lever table #: (``dispatch.compose``) decides which served levers a deployment is built with. +#: At TEP8 ``finalize_two_shot`` is the expert finalize and all-reduce of decode steps of +#: 24..``TP8_MAX_ROWS`` rows (3 to 19 requests) as one two-shot kernel over the two trays' +#: fabric (``kernels/finalize_two_shot``), in place of the built-in one-shot fused kernel at +#: 24..128 rows (bitwise its result) and of the unfused combine, BF16 cast and MNNVL two-shot +#: above (the fused kernel's rounding there, within one BF16 place of the unfused path's). The +#: composition serves it; its lever leaves it off there until its knob is set. TARGET_PATHS = { (103, 8): _dispatch.TargetPaths( served=_TEP8_PATHS | _PREFETCH_LEVERS, topology=_dispatch.TWO_TRAYS_OF_4 @@ -1091,6 +1123,8 @@ def __init__(self, model_config: ModelConfig, *args, **kwargs): ) self._plan = plan self._mapping = getattr(model_config, "mapping", None) + # Before any installer: the kernel modules read these settings when they build. + _lever_settings.apply(plan) if plan and self._runtime_hooks is None: # Every composition serves both executor fast paths # (dispatch.TargetPaths); NO_PATHS declares no hooks. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py index 42438ae7e6de..7b48e39f8d62 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py @@ -143,6 +143,42 @@ WIDE_SERVED_ROWS_BY_SM: dict[tuple[int, int], int] = {(10, 3): 56} +#: The widest TP8 bound the lever may name: the blocks of 8 requests at K = 7. +TP8_WIDE_MAX_ROWS = 64 +#: ``wide_drafter_hops__TP8_MAX_ROWS``: the rows of the wide workspace at TP8 (two trays), +#: which only the lever names (no ``WIDE_SERVED_ROWS_BY_SM`` entry covers that layout). +_TP8_WIDE_ROWS = TP8_WIDE_MAX_ROWS + + +def configure(*, tep8_wide_rows: int = TP8_WIDE_MAX_ROWS) -> None: + """Set the rows of the TP8 wide drafter-hop workspace. + + ``lever_settings.apply`` sets it from the model's plan before any install. It is read only + where that workspace is built (``kernels/drafter_hops.install_wide`` at TP8), so the + default changes nothing on a plan without ``wide_drafter_hops``. + + Args: + tep8_wide_rows: the workspace's rows (``_TP8_WIDE_ROWS``), + ``MAX_ROWS + 1``..``TP8_WIDE_MAX_ROWS``. + + Raises: + ValueError: for a row count outside that range. + """ + global _TP8_WIDE_ROWS + rows = int(tep8_wide_rows) + if not MAX_ROWS < rows <= TP8_WIDE_MAX_ROWS: + raise ValueError( + f"[musespark_open] wide_drafter_hops TP8 rows {rows}: expected " + f"{MAX_ROWS + 1}..{TP8_WIDE_MAX_ROWS}" + ) + _TP8_WIDE_ROWS = rows + + +def tp8_wide_rows() -> int: + """The TP8 wide workspace's rows, as ``configure`` set them.""" + return _TP8_WIDE_ROWS + + def wide_served_rows(capability: Optional[Tuple[int, int]] = None) -> int: """The served wide workspace's rows on ``capability`` (this process's device when None). diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py index f0a2b773a4e8..39f33a5765ff 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py @@ -310,6 +310,27 @@ def precompile( COMBINE_MAX_SLOTS = 256 COMBINE_SHAPES = ((4, 64, 8), (64, 16, 4), (128, 16, 4)) +#: Combine + output norm-gate in one launch (off by default, ``multi_request_combine_gate``): +#: on, the call site hands the call the layer's output gate (``gate``), and the combine stores +#: the built-in norm-gate of the attention output it merged -- the bits of the combine followed +#: by ``_attention_norm_gate_kernel`` -- so the norm-gate launch is gone. +COMBINE_GATE = False +#: Rows a fused gate is certified for: those the built-in output gate runs its Triton kernel on. +COMBINE_GATE_MAX_ROWS = 128 + + +def configure(*, combine_gate: bool = False) -> None: + """Set whether the combine also stores the attention output's norm-gate. + + ``lever_settings.apply`` sets it from the model's plan before any install; off, the + default, keeps the separate norm-gate launch. + + Args: + combine_gate: the ``multi_request_combine_gate`` lever (``COMBINE_GATE``). + """ + global COMBINE_GATE + COMBINE_GATE = bool(combine_gate) + @triton.jit(do_not_specialize=["num_requests", "num_work", "rows", "out_stride", "out_scale"]) def _combine_kernel( @@ -335,17 +356,41 @@ def _combine_kernel( MAX_SLOTS: tl.constexpr, DEVICE_SCALE: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, + gate_ptr=None, + gate_token_stride=0, + gate_head_stride=0, + gate_eps=0.0, + FUSED_GATE: tl.constexpr = False, + TRIGGER_AFTER_WAIT: tl.constexpr = False, ): """One program per (request, token, group of GH query heads): the request's slots, recomputed from ``seq_lens`` as the split kernel cuts the pages, merged in two passes (the maximum log2-sum-exp over all slots, then the weighted sum of O in chunks of CHUNK slots). The first chunk's loads are issued with the log2-sum-exp loads, before the maximum is - known.""" + known. + + ``FUSED_GATE``: the stored rows are the built-in output norm-gate of the BF16 attention + output (``_attention_norm_gate_kernel``'s statements on each of the GH heads, its summation + tree written out), not the attention output; the gate heads are loaded right after the + wait. ``TRIGGER_AFTER_WAIT``: the dependent may launch once this grid has passed its wait + (the PDL norm-gate's ``after_wait``).""" b = tl.program_id(0) tok = tl.program_id(1) hg = tl.program_id(2) if LAUNCH_WITH_PDL: tl.extra.cuda.gdc_wait() + if TRIGGER_AFTER_WAIT: + tl.extra.cuda.gdc_launch_dependents() + if FUSED_GATE: + # The prep's gate copy (complete before the split kernel passed its wait). + g_row = b * rows + tok + g_heads = hg * GH + tl.arange(0, GH) + gate = tl.load( + gate_ptr + + g_row * gate_token_stride + + g_heads[:, None] * gate_head_stride + + tl.arange(0, D)[None, :] + ) ib = tl.arange(0, BLOCK_B) okb = ib < num_requests L = tl.load(seq_lens_ptr + ib, mask=okb, other=0) @@ -415,7 +460,32 @@ def _combine_kernel( scale = out_scale res = acc * (scale / den)[:, None] row = b * rows + tok - tl.store(out_ptr + row * out_stride + hh[:, None] * D + dd[None, :], res.to(tl.bfloat16)) + if FUSED_GATE: + # The built-in _attention_norm_gate_kernel's statements on each head of the BF16 + # attention output this program would store. + x = res.to(tl.bfloat16).to(tl.float32) + variance = _heads_sum_of_squares(x, GH) / D + normed = (x * tl.rsqrt(variance + gate_eps)[:, None]).to(tl.bfloat16) + sigmoid = tl.sigmoid(gate.to(tl.float32)).to(tl.bfloat16) + result = normed.to(tl.float32) * sigmoid.to(tl.float32) + tl.store(out_ptr + row * out_stride + hh[:, None] * D + dd[None, :], result) + else: + tl.store(out_ptr + row * out_stride + hh[:, None] * D + dd[None, :], res.to(tl.bfloat16)) + + +@triton.jit +def _heads_sum_of_squares(x, GH: tl.constexpr): + """``sum(x * x)`` of each of ``GH`` 64-wide heads ``[GH, 64]``, in the built-in norm-gate + kernel's add order (``swa_block._norm_gate_sum_of_squares`` per head): element + ``32 w + 16 b4 + 8 b3 + 4 b2 + 2 b1 + b0`` of a head paired over ``b4``, ``b3``, ``b2``, + ``b1``, ``b0``, then ``w``; each size-2 sum is one IEEE add, which commutes.""" + sq = tl.reshape(x * x, [GH, 2, 2, 2, 2, 2, 2]) + sq = tl.sum(sq, axis=2) # b4 + sq = tl.sum(sq, axis=2) # b3 + sq = tl.sum(sq, axis=2) # b2 + sq = tl.sum(sq, axis=2) # b1 + sq = tl.sum(sq, axis=2) # b0 + return tl.sum(sq, axis=1) # w class MultiRequestSplitKv(OpWrapper): @@ -439,6 +509,15 @@ class MultiRequestSplitKv(OpWrapper): COMBINE_MAX_SLOTS)``; one sm_103 device. The first call per device and build compiles the kernel, which must not happen inside a CUDA-graph capture (``precompile``). + + ``gate`` (BF16 ``[rows, H 64]`` or ``[rows, H, 64]``, unit last stride, at most + ``COMBINE_GATE_MAX_ROWS`` rows): ``out`` holds the built-in output norm-gate of the attention + output (``bf16(rmsnorm(head, gate_eps)) * bf16(sigmoid(gate head))``, the bits of + ``_attention_norm_gate_kernel`` on it) instead of the attention output; ``gate_trigger`` + ``after_wait`` lets the dependent launch once the combine has passed its wait. + ``defer_combine``: only the split kernel is launched; the call returns + ``combine(gate=None, gate_eps=0.0, gate_trigger="exit")``, which launches the combine (and + returns ``out``) and must be the next launch on the stream. """ ARCHS = frozenset({Arch.SM_103}) @@ -471,6 +550,10 @@ def __call__( workspace: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, num_work: Optional[int] = None, nq: int = NQ, + gate: Optional[torch.Tensor] = None, + gate_eps: float = 0.0, + gate_trigger: str = "exit", + defer_combine: bool = False, ) -> torch.Tensor: q2 = q.reshape(q.shape[0], -1) if q.dim() == 3 else q rows = q2.shape[0] if q2.dim() == 2 else 0 @@ -594,33 +677,77 @@ def __call__( cuda.CUstream(stream.cuda_stream), ) chunk, warps = next(s[1:] for s in COMBINE_SHAPES if requests <= s[0]) - _combine_kernel[(requests, per, kv_heads)]( - part_o, - part_lse, - seq_lens, - out2, - out2.stride(0), - 1.0 if device_scale else float(bmm2_scale), - bmm2_scale if device_scale else part_lse, - requests, - num_work, - per, - HQ=hq, - GH=Q_PER_KV, - D=HEAD_DIM, - SLOT_ROWS=MAX_ROWS * hq, - BLOCK_B=request_build(requests), - P=PAGE, - WINDOW=window, - MIN_P=MIN_PAGES, - CHUNK=chunk, - MAX_SLOTS=COMBINE_MAX_SLOTS, - DEVICE_SCALE=device_scale, - LAUNCH_WITH_PDL=enable_pdl, - num_warps=warps, - launch_pdl=enable_pdl, - ) - return out + + def combine(gate=None, gate_eps=0.0, gate_trigger="exit"): + """The combine launch into ``out`` (with ``gate``: its norm-gate, ``COMBINE_GATE``).""" + fused_gate = {} + if gate is not None: + if ( + gate.dtype != torch.bfloat16 + or gate.device != q.device + or gate.ndim not in (2, 3) + or gate.stride(-1) != 1 + or gate.shape[0] != rows + or gate.numel() != rows * hq * HEAD_DIM + or (gate.ndim == 3 and gate.shape[-1] != HEAD_DIM) + or rows > COMBINE_GATE_MAX_ROWS + or gate_trigger not in ("exit", "after_wait") + ): + raise ValueError( + f"multi_request_split_kv's gate must be BF16 [rows, H * {HEAD_DIM}] or " + f"[rows, H, {HEAD_DIM}] with unit last stride on the query's device, at " + f"most {COMBINE_GATE_MAX_ROWS} rows, trigger exit or after_wait; got " + f"{gate.dtype} {tuple(gate.shape)} strides {tuple(gate.stride())}, " + f"trigger {gate_trigger!r}" + ) + fused_gate = dict( + gate_ptr=gate, + gate_token_stride=gate.stride(0), + gate_head_stride=gate.stride(1) if gate.ndim == 3 else HEAD_DIM, + gate_eps=gate_eps, + FUSED_GATE=True, + TRIGGER_AFTER_WAIT=enable_pdl and gate_trigger == "after_wait", + ) + from tensorrt_llm.logger import logger + + logger.info_once( + f"[musespark_open] multi_request_split_kv: combine + output norm-gate in one " + f"launch ({requests} requests x {per} rows, trigger {gate_trigger})", + key=f"musespark_open.multi_request_split_kv.combine_gate_{requests}_{per}", + ) + _combine_kernel[(requests, per, kv_heads)]( + part_o, + part_lse, + seq_lens, + out2, + out2.stride(0), + 1.0 if device_scale else float(bmm2_scale), + bmm2_scale if device_scale else part_lse, + requests, + num_work, + per, + HQ=hq, + GH=Q_PER_KV, + D=HEAD_DIM, + SLOT_ROWS=MAX_ROWS * hq, + BLOCK_B=request_build(requests), + P=PAGE, + WINDOW=window, + MIN_P=MIN_PAGES, + CHUNK=chunk, + MAX_SLOTS=COMBINE_MAX_SLOTS, + DEVICE_SCALE=device_scale, + LAUNCH_WITH_PDL=enable_pdl, + num_warps=warps, + launch_pdl=enable_pdl, + **fused_gate, + ) + return out + + if defer_combine: + # The caller launches it (``combine(...)``) as the next kernel on this stream. + return combine + return combine(gate, gate_eps, gate_trigger) def reference( self, @@ -636,6 +763,10 @@ def reference( workspace: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, num_work: Optional[int] = None, nq: int = NQ, + gate: Optional[torch.Tensor] = None, + gate_eps: float = 0.0, + gate_trigger: str = "exit", + defer_combine: bool = False, ) -> torch.Tensor: """fp32 attention over the cache with the op's Q (an E4M3 Q as is; a BF16 Q as ``Q_hi + Q_lo``, or one E4M3 term for ``nq=1``, formed exactly as the kernel forms them), @@ -686,6 +817,16 @@ def gather(kv: int) -> torch.Tensor: result[b * rows + t, h * Q_PER_KV : (h + 1) * Q_PER_KV] = ( torch.softmax(s, dim=-1) @ v[lo : pos + 1, h] ) * bmm2_scale + if gate is not None: + # The output norm-gate in float64 between its BF16 rounding points, on the BF16 + # attention output. + x = result.to(torch.bfloat16).double() + normed = (x * torch.rsqrt(x.square().mean(-1, keepdim=True) + gate_eps)).to( + torch.bfloat16 + ) + g = gate.reshape(result.shape).double() + sigmoid = torch.sigmoid(g).to(torch.bfloat16) + result = (normed.double() * sigmoid.double()).float() return result.reshape(q2.shape[0], -1) def is_valid( @@ -702,6 +843,10 @@ def is_valid( workspace: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, num_work: Optional[int] = None, nq: int = NQ, + gate: Optional[torch.Tensor] = None, + gate_eps: float = 0.0, + gate_trigger: str = "exit", + defer_combine: bool = False, ) -> None: """Inputs the kernels take and answer wrongly, with no error. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py index fe7ddda65b38..20a92629ea36 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py @@ -145,6 +145,55 @@ def _prep_rows( return y.to(tl.bfloat16) +@triton.jit +def _prep_rows_load(ptr, row_off, col_off, mask, N: tl.constexpr): + """``_prep_rows``' loads alone (BF16, as stored), so a caller can issue them early.""" + d = tl.arange(0, 64) + lane = tl.arange(0, 8) + head = ptr + row_off[:, None] + col_off[:, None] + m = mask[:, None] + x0 = tl.load(head + lane[None, :] * 8 + 0, mask=m, other=0.0) + x1 = tl.load(head + lane[None, :] * 8 + 1, mask=m, other=0.0) + x2 = tl.load(head + lane[None, :] * 8 + 2, mask=m, other=0.0) + x3 = tl.load(head + lane[None, :] * 8 + 3, mask=m, other=0.0) + x4 = tl.load(head + lane[None, :] * 8 + 4, mask=m, other=0.0) + x5 = tl.load(head + lane[None, :] * 8 + 5, mask=m, other=0.0) + x6 = tl.load(head + lane[None, :] * 8 + 6, mask=m, other=0.0) + x7 = tl.load(head + lane[None, :] * 8 + 7, mask=m, other=0.0) + x = tl.load(head + d[None, :], mask=m, other=0.0) + return x0, x1, x2, x3, x4, x5, x6, x7, x + + +@triton.jit +def _prep_rows_math(x0, x1, x2, x3, x4, x5, x6, x7, x, cos, sin, scale, eps, N: tl.constexpr, + APPLY_ROPE: tl.constexpr): # fmt: skip + """``_prep_rows``' arithmetic, statement for statement, on ``_prep_rows_load``'s values.""" + x0 = x0.to(tl.float32) + x1 = x1.to(tl.float32) + x2 = x2.to(tl.float32) + x3 = x3.to(tl.float32) + x4 = x4.to(tl.float32) + x5 = x5.to(tl.float32) + x6 = x6.to(tl.float32) + x7 = x7.to(tl.float32) + s = libdevice.add_rn(x0 * x0, x1 * x1) + s = libdevice.add_rn(s, x2 * x2) + s = libdevice.add_rn(s, x3 * x3) + s = libdevice.add_rn(s, x4 * x4) + s = libdevice.add_rn(s, x5 * x5) + s = libdevice.add_rn(s, x6 * x6) + s = libdevice.add_rn(s, x7 * x7) + s = tl.reshape(s, [N, 2, 2, 2]) + total = tl.sum(tl.sum(tl.sum(s, axis=3), axis=2), axis=1) + inverse_rms = _rsqrt_approx_ftz(total / 64.0 + eps) + y = _normalize(x.to(tl.float32), inverse_rms[:, None], scale) + if APPLY_ROPE: + even, odd = tl.split(tl.reshape(y, [N, 32, 2])) + even, odd = _rotate(even, odd, cos, sin) + y = tl.reshape(tl.join(even, odd), [N, 64]) + return y.to(tl.bfloat16) + + @triton.jit def _rope_rows(cos_sin_ptr, positions_ptr, index, mask, APPLY_ROPE: tl.constexpr): """``(cos, sin)`` ``[N, 32]`` of the positions ``positions[index]`` (zeros where masked).""" @@ -191,6 +240,7 @@ def _swa_split_kernel( APPLY_ROPE: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, REQS: tl.constexpr = 1, + EARLY_LOADS: tl.constexpr = False, ): G: tl.constexpr = HQ // HKV PAGE: tl.constexpr = HKV * P * D @@ -288,9 +338,33 @@ def _swa_split_kernel( tl.extra.cuda.gdc_wait() tl.extra.cuda.gdc_launch_dependents() + if EARLY_LOADS: + # Every operand load of this program at once, right after the wait (with the fused + # prep: the query rows, the key block's step K/V rows, the writer's step K/V rows; + # masked rows load nothing): one dependent round trip instead of one per prep. The + # arithmetic below is the same statements on the same values. + e_q = _prep_rows_load(qkv_ptr, t * qkv_stride, (h * G + r % G) * q_stride, row_ok, R) + e_k = _prep_rows_load(qkv_ptr, new_row * qkv_stride, tl.zeros([BLOCK], tl.int32) + k_col, + new, BLOCK) # fmt: skip + e_v = tl.load( + qkv_ptr + new_row[:, None] * qkv_stride + v_col + d[None, :], + mask=new[:, None], + other=0.0, + ) + e_wk = _prep_rows_load(qkv_ptr, tok * qkv_stride, tl.zeros([8], tl.int32) + k_col, write_ok, 8) + e_wv = tl.load( + qkv_ptr + tok[:, None] * qkv_stride + v_col + d[None, :], + mask=write_ok[:, None], + other=0.0, + ) + if writer: # The step's K/V into the paged cache: what applyBiasRopeUpdateKVCacheV2 stores. - if FUSED_PREP: + if EARLY_LOADS: + wk0, wk1, wk2, wk3, wk4, wk5, wk6, wk7, wkx = e_wk + k_new = _prep_rows_math(wk0, wk1, wk2, wk3, wk4, wk5, wk6, wk7, wkx, cos_w, sin_w, 1.0, + prep_eps, 8, APPLY_ROPE) # fmt: skip + elif FUSED_PREP: k_new = _prep_rows( qkv_ptr, tok * qkv_stride, @@ -309,11 +383,14 @@ def _swa_split_kernel( mask=tok_ok[:, None], other=0.0, ) - v_new = tl.load( - qkv_ptr + tok[:, None] * qkv_stride + v_col + d[None, :], - mask=tok_ok[:, None], - other=0.0, - ) + if EARLY_LOADS: + v_new = e_wv + else: + v_new = tl.load( + qkv_ptr + tok[:, None] * qkv_stride + v_col + d[None, :], + mask=tok_ok[:, None], + other=0.0, + ) tl.store( pool_ptr + (w_page_k * PAGE + slot)[:, None] + d[None, :], _to_cache_fp8(k_new, scale_bf16), @@ -326,7 +403,12 @@ def _swa_split_kernel( ) q_head = h * G + r % G - if FUSED_PREP: + if EARLY_LOADS: + qx0, qx1, qx2, qx3, qx4, qx5, qx6, qx7, qxx = e_q + q = _prep_rows_math(qx0, qx1, qx2, qx3, qx4, qx5, qx6, qx7, qxx, cos_q, sin_q, q_scale, prep_eps, R, + APPLY_ROPE) # fmt: skip + q = tl.where(row_ok[:, None], q, 0.0) + elif FUSED_PREP: q = _prep_rows( qkv_ptr, t * qkv_stride, @@ -357,7 +439,13 @@ def _swa_split_kernel( # Whether this program's key block holds any of the step's own keys. has_new = base + b0 * BLOCK + BLOCK > past for b in range(b0, b1): # at most once: one key block per program - if FUSED_PREP: + if EARLY_LOADS: + k_new = tl.zeros([BLOCK, D], tl.bfloat16) + if has_new: + kx0, kx1, kx2, kx3, kx4, kx5, kx6, kx7, kxx = e_k + k_new = _prep_rows_math(kx0, kx1, kx2, kx3, kx4, kx5, kx6, kx7, kxx, cos_k, sin_k, 1.0, + prep_eps, BLOCK, APPLY_ROPE) # fmt: skip + elif FUSED_PREP: k_new = tl.zeros([BLOCK, D], tl.bfloat16) if has_new: k_new = _prep_rows( @@ -378,11 +466,14 @@ def _swa_split_kernel( mask=new[:, None], other=0.0, ) - v_new = tl.load( - qkv_ptr + new_row[:, None] * qkv_stride + v_col + d[None, :], - mask=new[:, None], - other=0.0, - ) + if EARLY_LOADS: + v_new = e_v + else: + v_new = tl.load( + qkv_ptr + new_row[:, None] * qkv_stride + v_col + d[None, :], + mask=new[:, None], + other=0.0, + ) k = tl.where(new[:, None], _to_cache_fp8(k_new, scale_bf16), k_old) v = tl.where(new[:, None], _to_cache_fp8(v_new, scale_bf16), v_old) s = tl.dot(q, tl.trans(k.to(tl.bfloat16))) * qk_scale @@ -967,6 +1058,7 @@ def _swa_combine_kernel( FUSED_GATE: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, REQS: tl.constexpr = 1, + EARLY_LOADS: tl.constexpr = False, ): G: tl.constexpr = HQ // HKV row = tl.program_id(0) # token * G + query head within the KV group @@ -979,6 +1071,16 @@ def _swa_combine_kernel( out_ptr += req * M * out_stride if FUSED_GATE: gate_ptr += req * M * gate_token_stride + if EARLY_LOADS: + if FUSED_GATE: + # Before the wait: the gate is the QKV projection's (or the prep's), complete + # before this grid launched -- the split launches it only after its own wait, + # which returns once that kernel has completed. Read at L2. + g_early = tl.load( + gate_ptr + (row // G) * gate_token_stride + (h * G + row % G) * gate_head_stride + + tl.arange(0, D), + cache_modifier=".cg", + ) # fmt: skip if LAUNCH_WITH_PDL: tl.extra.cuda.gdc_wait() if FUSED_GATE: @@ -991,10 +1093,14 @@ def _swa_combine_kernel( part = (sp * HKV + h) * (M * G) + row m_split = tl.load(part_ml_ptr + part * 2, mask=s_ok, other=-1.0e30) l_split = tl.load(part_ml_ptr + part * 2 + 1, mask=s_ok, other=0.0) + if EARLY_LOADS: + # The partial O with the (max, sum) words: one dependent round trip. + o = tl.load(part_o_ptr + part[:, None] * D + d[None, :], mask=s_ok[:, None], other=0.0) m_max = tl.max(m_split, 0) w = tl.where(s_ok & (l_split > 0), tl.math.exp2(m_split - m_max), 0.0) den = tl.sum(w * l_split, 0) - o = tl.load(part_o_ptr + part[:, None] * D + d[None, :], mask=s_ok[:, None], other=0.0) + if not EARLY_LOADS: + o = tl.load(part_o_ptr + part[:, None] * D + d[None, :], mask=s_ok[:, None], other=0.0) num = tl.sum(o * w[:, None], 0) t = row // G head = h * G + row % G @@ -1004,7 +1110,10 @@ def _swa_combine_kernel( x = attention.to(tl.float32) variance = _norm_gate_sum_of_squares(x) / D normed = (x * tl.rsqrt(variance + gate_eps)).to(tl.bfloat16) - g = tl.load(gate_ptr + t * gate_token_stride + head * gate_head_stride + d).to(tl.float32) + if EARLY_LOADS: + g = g_early.to(tl.float32) + else: + g = tl.load(gate_ptr + t * gate_token_stride + head * gate_head_stride + d).to(tl.float32) sigmoid = tl.sigmoid(g).to(tl.bfloat16) result = normed.to(tl.float32) * sigmoid.to(tl.float32) tl.store(out_ptr + t * out_stride + head * D + d, result) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_cluster_block.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_cluster_block.py new file mode 100644 index 000000000000..5d9019ba20d4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_cluster_block.py @@ -0,0 +1,1071 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``swa_block_multi_attention`` of a 2..4-request decode step in one launch: a CTA cluster per (request, KV head). + +The multi-request sliding-window block (``swa_block``, the 64-key x 64-row tile) launches two dependent grids per +layer: a split grid (one program per 64-key block of the window, 34 per request and KV head) that writes FP32 +partials to global memory, and a combine grid that reads them back, merges them and applies the output norm-gate. +Here the 34 splits of one (request, KV head) are the 8 CTAs of one cluster: CTA ``r`` owns key blocks +``r, r + 8, r + 16, r + 24 (, r + 32)``, runs each block as the split program does, and the merge goes through +distributed shared memory instead of global memory and a second grid: + +1. Before ``griddepcontrol.wait``: the sequence length, the block table, the cached K/V of the CTA's blocks (FP8, + widened to BF16 in shared memory), the positions and RoPE rows, the KV scale. Tensor memory is allocated only + after the wait and before the dependents are triggered (no allocation can wait on a dependent). +2. After it: every CTA runs the fused prep of the 64 query rows of its KV head (8 tokens x 8 heads) and of the + step's 8 keys; cluster rank 0 writes the step's K/V into the paged cache (the split's writer program); the CTA + that owns a block holding a step key patches its tile. +3. Per block: S = Q K^T (tcgen05, M64 N64, four K16 steps, fresh accumulator), the split's softmax (scale, window + mask, row max, ``ex2.approx.ftz``, its row-sum tree), P to tensor memory in BF16, O = P V (four K16 steps onto + a zeroed accumulator). Each block's (max, sum) per row goes to every CTA of the cluster (``st.async``). +4. With every block's (max, sum) in: each CTA weighs its own blocks' O (``exp2(m - m_max)``, zero for an empty + split) and sums them in the combine's tree for its split class (splits ``r mod 8``: the combine's thread holds + exactly those), then sends the sums to the CTA owning the token (cluster rank ``t`` owns token ``t``'s 8 heads). +5. The owner finishes the combine's tree across the 8 classes, divides by the denominator (the combine's + butterfly over its 64 split lanes), and applies the output norm-gate. + +Numerics: each statement is the served Triton build's instruction (``_swa_split_kernel`` / ``_swa_combine_kernel`` +at ``MULTI_KEYS = MULTI_ROWS = 64`` with the fused prep, Q in FP8 and the fused gate, as Triton 3.8 compiles them +for sm_103a): the same tcgen05 instruction descriptors, operand layouts and K order, ``add.rn.ftz`` / +``rsqrt.approx.ftz`` / ``div.full`` / ``ex2.approx(.ftz)`` / ``cvt.rn.satfinite.e4m3x2`` / ``mul.bf16`` where it +has them, its packed-add trees in the same pairing, and its contracted FMA (the denominator's first butterfly +level). Every arithmetic statement is inline PTX with an explicit rounding mode, so nothing is re-contracted. The +KV cache bytes and the gated output are the two-grid build's. + +Only ``swa_block.SwaBlockMultiAttention`` calls this module, under ``TRTLLM_MUSESPARK_SWA_BLOCK_CLUSTER=1`` and the +conditions it checks (``fits``). +""" + +import math +from typing import Optional + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +#: CTAs per cluster, one cluster per (request, KV head). The 34 splits go round-robin, so CTA ``r`` holds every +#: split of the combine's tree class ``r mod 8``. +CLUSTER = 8 +#: Query rows of one KV head: 8 tokens x 8 query heads (row = token * 8 + head in group). +ROWS = 64 +TOKENS = 8 +GROUP = 8 +D = 64 +#: Keys per block (the multi-request tile), the window and its block count at that tile. +BLOCK = 64 +WINDOW = 2048 +NUM_SPLITS = 34 +#: Most blocks one CTA owns (ranks 0 and 1 own five, the others four). +NB_MAX = 5 +#: One warpgroup per owned block. +THREADS = 128 * NB_MAX + +# Shared memory (bytes; the tiles 1024-aligned, as the 128-byte swizzle needs). +TILE_BYTES = ROWS * D * 2 +Q_OFF = 0 +K_OFF = TILE_BYTES +V_OFF = K_OFF + NB_MAX * TILE_BYTES +ML_OFF = V_OFF + NB_MAX * TILE_BYTES # [NUM_SPLITS][ROWS] x (max, sum) fp32 +ML_BYTES = NUM_SPLITS * ROWS * 8 +T_OFF = ML_OFF + ML_BYTES # [CLUSTER classes][8 heads][64 dims] fp32: the owner's class sums +T_BYTES = CLUSTER * GROUP * D * 4 +BAR_OFF = T_OFF + T_BYTES # s_bar[5], o_bar[5], ml_bar, t_bar +SLOT_OFF = BAR_OFF + 16 * 8 # the TMEM base address +DEN_OFF = SLOT_OFF + 16 # per head, the two halves' denominators +SMEM_BYTES = DEN_OFF + 64 +# Tensor memory: block k at columns 96 k: O (lanes 0-15 of each subpartition, 64 columns), S (lanes 16-31, the +# same columns), P in BF16 (lanes 0-15, columns 64-95): the split's own placement, once per block. +TMEM_COLS = 512 +COLS_PER_BLOCK = 96 + +#: tcgen05 kind::f16 instruction descriptors of the split (BF16 x BF16 -> FP32, M64 N64; PV with B MN-major). +IDESC_QK = 0x04100490 +IDESC_PV = 0x04110490 +#: High word of the split's shared-memory descriptors: SBO 1024 B, version 1, 128-byte swizzle. +DESC_HI = 0x40004040 + +NEG_BIG = -1.0e30 + +# Debug dump (``debug=`` tensor, uint8): Q tiles, then each split's (max, sum), then each split's O * dequant. +DBG_REQS = 8 + + +STAMPS = 16 + + +def dbg_layout(kv_heads: int) -> dict: + q = DBG_REQS * kv_heads * ROWS * D * 2 + ml = DBG_REQS * NUM_SPLITS * kv_heads * ROWS * 8 + o = DBG_REQS * NUM_SPLITS * kv_heads * ROWS * D * 4 + ts = DBG_REQS * kv_heads * CLUSTER * STAMPS * 8 + return dict(q=0, ml=q, o=q + ml, ts=q + ml + o, bytes=q + ml + o + ts) + + +def owned_splits(rank: int) -> list: + """The splits (key blocks) cluster rank ``rank`` runs, in its block-slot order.""" + return [rank + CLUSTER * k for k in range(NB_MAX) if rank + CLUSTER * k < NUM_SPLITS] + + +def fits(rows: int, reqs: int, window: int, num_splits: int, fused_prep: bool, q_fp8: bool, has_gate: bool, + keys: int, tile_rows: int) -> Optional[str]: + """``None`` when a multi-request call is one this block computes bit for bit, else why not.""" + if keys != BLOCK or tile_rows != ROWS: + return f"tile {keys} keys x {tile_rows} rows (needs {BLOCK} x {ROWS})" + if rows != TOKENS: + return f"{rows} rows per request (needs {TOKENS})" + if not 2 <= reqs <= 4: + return f"{reqs} requests (needs 2..4)" + if window != WINDOW or num_splits != NUM_SPLITS: + return f"window {window} / {num_splits} splits (needs {WINDOW} / {NUM_SPLITS})" + if not (fused_prep and q_fp8 and has_gate): + return "needs the fused prep, an FP8 Q and the fused gate" + return None + + +AD = llvm.AsmDialect.AD_ATT + +def _v(x, loc, ip): + return x.ir_value(loc=loc, ip=ip) + +def _asm(ret, args, text, cons, side): + return llvm.inline_asm(ret, args, text, cons, has_side_effects=side, is_align_stack=False, asm_dialect=AD) + +# ------------------------------------------------------------------------------------------ PTX helpers +@dsl_user_op +def f1(text, a, *, loc=None, ip=None): + return Float32(_asm(T.f32(), [_v(Float32(a), loc, ip)], text, "=f,f", False)) + +@dsl_user_op +def f2(text, a, b, *, loc=None, ip=None): + return Float32(_asm(T.f32(), [_v(Float32(a), loc, ip), _v(Float32(b), loc, ip)], text, "=f,f,f", False)) + +@dsl_user_op +def f3(text, a, b, c, *, loc=None, ip=None): + return Float32(_asm(T.f32(), [_v(Float32(a), loc, ip), _v(Float32(b), loc, ip), _v(Float32(c), loc, ip)], + text, "=f,f,f,f", False)) + +def add(a, b): + return f2("add.rn.f32 $0, $1, $2;", a, b) + +def add_ftz(a, b): + return f2("add.rn.ftz.f32 $0, $1, $2;", a, b) + +def sub(a, b): + return f2("sub.rn.f32 $0, $1, $2;", a, b) + +def mul(a, b): + return f2("mul.rn.f32 $0, $1, $2;", a, b) + +def mul_ftz(a, b): + return f2("mul.rn.ftz.f32 $0, $1, $2;", a, b) + +def fma(a, b, c): + return f3("fma.rn.f32 $0, $1, $2, $3;", a, b, c) + +def fma_ftz(a, b, c): + return f3("fma.rn.ftz.f32 $0, $1, $2, $3;", a, b, c) + +def div_full(a, b): + return f2("div.full.f32 $0, $1, $2;", a, b) + +def fmax(a, b): + return f2("max.f32 $0, $1, $2;", a, b) + +def rsqrt_ftz(a): + return f1("rsqrt.approx.ftz.f32 $0, $1;", a) + +def ex2_ftz(a): + return f1("ex2.approx.ftz.f32 $0, $1;", a) + +def neg(a): + return f1("neg.f32 $0, $1;", a) + +def bf16r(a): + """fp32 -> BF16 (round to nearest even) -> fp32.""" + return f1("{\n\t.reg .b16 h;\n\tcvt.rn.bf16.f32 h, $1;\n\tcvt.f32.bf16 $0, h;\n\t}", a) + +def bf16_mul(a, b): + """``mul.bf16`` of two BF16-exact fp32 values, as fp32 (the served ``(x * s).to(bf16)``).""" + return f2("{\n\t.reg .b16 x, y, z;\n\tcvt.rn.bf16.f32 x, $1;\n\tcvt.rn.bf16.f32 y, $2;\n\t" + "mul.rn.bf16 z, x, y;\n\tcvt.f32.bf16 $0, z;\n\t}", a, b) + +def sigmoid_bf16(g): + """The combine's ``tl.sigmoid(g).to(bf16)``: ``1 / (1 + ex2.approx(g * -log2 e))``, ``div.full``.""" + return f1("{\n\t.reg .f32 x, e;\n\tmul.rn.f32 x, $1, 0fBFB8AA3B;\n\tex2.approx.f32 e, x;\n\t" + "add.rn.f32 e, e, 0f3F800000;\n\tdiv.full.f32 x, 0f3F800000, e;\n\t.reg .b16 h;\n\t" + "cvt.rn.bf16.f32 h, x;\n\tcvt.f32.bf16 $0, h;\n\t}", g) + +@dsl_user_op +def e4m3x2(hi, lo, *, loc=None, ip=None): + """``cvt.rn.satfinite.e4m3x2.f32`` (``hi`` in the upper byte) as the low 16 bits of an Int32.""" + return Int32(_asm(T.i32(), [_v(Float32(hi), loc, ip), _v(Float32(lo), loc, ip)], + "{\n\t.reg .b16 q;\n\tcvt.rn.satfinite.e4m3x2.f32 q, $1, $2;\n\tcvt.u32.u16 $0, q;\n\t}", + "=r,f,f", False)) + +@dsl_user_op +def e4m3x2_to_bf16x2(q, *, loc=None, ip=None): + """Two E4M3 (the low 16 bits of ``q``, element 0 in the low byte) -> BF16x2 (exact).""" + return Int32(_asm(T.i32(), [_v(Int32(q), loc, ip)], + "{\n\t.reg .b16 q, a, b;\n\t.reg .b32 h;\n\t.reg .f32 x, y;\n\tcvt.u16.u32 q, $1;\n\t" + "cvt.rn.f16x2.e4m3x2 h, q;\n\tmov.b32 {a, b}, h;\n\tcvt.f32.f16 x, a;\n\t" + "cvt.f32.f16 y, b;\n\tcvt.rn.bf16x2.f32 $0, y, x;\n\t}", "=r,r", False)) + +@dsl_user_op +def lo_bf16(w, *, loc=None, ip=None): + return Float32(_asm(T.f32(), [_v(Int32(w), loc, ip)], "shl.b32 $0, $1, 16;", "=f,r", False)) + +@dsl_user_op +def hi_bf16(w, *, loc=None, ip=None): + return Float32(_asm(T.f32(), [_v(Int32(w), loc, ip)], "and.b32 $0, $1, 0xFFFF0000;", "=f,r", False)) + +@dsl_user_op +def pack_bf16x2(hi, lo, *, loc=None, ip=None): + """``cvt.rn.bf16x2.f32`` (``lo`` in the low half).""" + return Int32(_asm(T.i32(), [_v(Float32(hi), loc, ip), _v(Float32(lo), loc, ip)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", "=r,f,f", False)) + +@dsl_user_op +def bits(x, *, loc=None, ip=None): + return Int32(_asm(T.i32(), [_v(Float32(x), loc, ip)], "mov.b32 $0, $1;", "=r,f", False)) + +@dsl_user_op +def as_f32(x, *, loc=None, ip=None): + return Float32(_asm(T.f32(), [_v(Int32(x), loc, ip)], "mov.b32 $0, $1;", "=f,r", False)) + +@dsl_user_op +def selp(a, b, p, *, loc=None, ip=None): + """``p != 0 ? a : b`` (fp32).""" + return Float32(_asm(T.f32(), [_v(Float32(a), loc, ip), _v(Float32(b), loc, ip), _v(Int32(p), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\tselp.f32 $0, $1, $2, q;\n\t}", + "=f,f,f,r", False)) + +@dsl_user_op +def isel(a, b, p, *, loc=None, ip=None): + return Int32(_asm(T.i32(), [_v(Int32(a), loc, ip), _v(Int32(b), loc, ip), _v(Int32(p), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\tselp.b32 $0, $1, $2, q;\n\t}", + "=r,r,r,r", False)) + +@dsl_user_op +def lt(a, b, *, loc=None, ip=None): + """``a < b`` (signed) as 0 / 1.""" + return Int32(_asm(T.i32(), [_v(Int32(a), loc, ip), _v(Int32(b), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.lt.s32 q, $1, $2;\n\tselp.u32 $0, 1, 0, q;\n\t}", + "=r,r,r", False)) + +@dsl_user_op +def gt0(x, *, loc=None, ip=None): + """``x > 0`` (fp32, ordered) as 0 / 1.""" + return Int32(_asm(T.i32(), [_v(Float32(x), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.gt.f32 q, $1, 0f00000000;\n\tselp.u32 $0, 1, 0, q;\n\t}", + "=r,f", False)) + +@dsl_user_op +def keep_sel(s, lo_e, hi_e, e, *, loc=None, ip=None): + """The split's window mask on one score: ``lo_e <= e <= hi_e ? s : -inf``.""" + return Float32(_asm(T.f32(), [_v(Float32(s), loc, ip), _v(Int32(lo_e), loc, ip), _v(Int32(hi_e), loc, ip)], + "{\n\t.reg .pred a, b;\n\tsetp.le.s32 a, $2, " + str(e) + ";\n\tsetp.ge.s32 b, $3, " + + str(e) + ";\n\tand.pred a, a, b;\n\tselp.f32 $0, $1, 0fFF800000, a;\n\t}", + "=f,f,r,r", False)) + +@dsl_user_op +def shfl_bfly(x, m, *, loc=None, ip=None): + return Float32(_asm(T.f32(), [_v(Float32(x), loc, ip)], f"shfl.sync.bfly.b32 $0, $1, {m}, 31, -1;", + "=f,f", True)) + +@dsl_user_op +def shfl_idx(x, src, *, loc=None, ip=None): + return Float32(_asm(T.f32(), [_v(Float32(x), loc, ip)], f"shfl.sync.idx.b32 $0, $1, {src}, 31, -1;", + "=f,f", True)) + +# ---- memory +@dsl_user_op +def ld_g_v4(addr, pred, *, loc=None, ip=None): + """16 bytes from global memory as 4 Int32 (zeros where ``pred`` is 0).""" + res = _asm(llvm.StructType.get_literal([T.i32()] * 4), [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $5, 0;\n\tmov.b32 $0, 0;\n\tmov.b32 $1, 0;\n\t" + "mov.b32 $2, 0;\n\tmov.b32 $3, 0;\n\t@q ld.global.v4.b32 {$0, $1, $2, $3}, [$4];\n\t}", + "=r,=r,=r,=r,l,r", True) + return [Int32(llvm.extractvalue(T.i32(), res, [i])) for i in range(4)] + +@dsl_user_op +def ld_g_v2(addr, pred, *, loc=None, ip=None): + res = _asm(llvm.StructType.get_literal([T.i32()] * 2), [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\tmov.b32 $0, 0;\n\tmov.b32 $1, 0;\n\t" + "@q ld.global.v2.b32 {$0, $1}, [$2];\n\t}", "=r,=r,l,r", True) + return [Int32(llvm.extractvalue(T.i32(), res, [i])) for i in range(2)] + +@dsl_user_op +def ld_g_b32(addr, pred, *, loc=None, ip=None): + return Int32(_asm(T.i32(), [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $2, 0;\n\tmov.b32 $0, 0;\n\t" + "@q ld.global.b32 $0, [$1];\n\t}", "=r,l,r", True)) + +@dsl_user_op +def st_g_v2(addr, a, b, pred, *, loc=None, ip=None): + _asm(None, [_v(Int64(addr), loc, ip), _v(Int32(a), loc, ip), _v(Int32(b), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\t@q st.global.v2.b32 [$0], {$1, $2};\n\t}", + "l,r,r,r", True) + +@dsl_user_op +def st_g_b32(addr, a, pred, *, loc=None, ip=None): + _asm(None, [_v(Int64(addr), loc, ip), _v(Int32(a), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $2, 0;\n\t@q st.global.b32 [$0], $1;\n\t}", "l,r,r", True) + +@dsl_user_op +def st_s_v4(addr, w, pred, *, loc=None, ip=None): + _asm(None, [_v(Int32(addr), loc, ip)] + [_v(Int32(x), loc, ip) for x in w] + [_v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $5, 0;\n\t@q st.shared.v4.b32 [$0], {$1, $2, $3, $4};\n\t}", + "r,r,r,r,r,r", True) + +def isel4(a, b, p): + """Element-wise ``p ? a : b`` of two lists of Int32.""" + return [isel(x, y, p) for x, y in zip(a, b)] + + +@dsl_user_op +def st_s_b32_pred(addr, a, pred, *, loc=None, ip=None): + _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(a), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $2, 0;\n\t@q st.shared.b32 [$0], $1;\n\t}", "r,r,r", True) + + +@dsl_user_op +def bar_named(bar_id, count, *, loc=None, ip=None): + """``bar.sync`` on named barrier ``bar_id`` (a register) for ``count`` threads.""" + _asm(None, [_v(Int32(bar_id), loc, ip)], f"bar.sync $0, {count};", "r", True) + + +@dsl_user_op +def stamp(addr, pred, *, loc=None, ip=None): + """``%clock64`` to global memory (the debug build's phase stamps).""" + _asm(None, [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\t.reg .b64 c;\n\tsetp.ne.s32 q, $1, 0;\n\tmov.u64 c, %clock64;\n\t" + "@q st.global.b64 [$0], c;\n\t}", "l,r", True) + + +@dsl_user_op +def st_s_b32(addr, a, *, loc=None, ip=None): + _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(a), loc, ip)], "st.shared.b32 [$0], $1;", "r,r", True) + +@dsl_user_op +def ld_s_b32(addr, *, loc=None, ip=None): + return Int32(_asm(T.i32(), [_v(Int32(addr), loc, ip)], "ld.shared.b32 $0, [$1];", "=r,r", True)) + +@dsl_user_op +def ld_s_v2f(addr, *, loc=None, ip=None): + res = _asm(llvm.StructType.get_literal([T.f32()] * 2), [_v(Int32(addr), loc, ip)], + "ld.shared.v2.f32 {$0, $1}, [$2];", "=f,=f,r", True) + return [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(2)] + +@dsl_user_op +def ld_s_v4f(addr, *, loc=None, ip=None): + res = _asm(llvm.StructType.get_literal([T.f32()] * 4), [_v(Int32(addr), loc, ip)], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", True) + return [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(4)] + +@dsl_user_op +def mapa(addr, rank, *, loc=None, ip=None): + return Int32(_asm(T.i32(), [_v(Int32(addr), loc, ip), _v(Int32(rank), loc, ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", False)) + +@dsl_user_op +def st_async(raddr, vals, rbar, pred, *, loc=None, ip=None): + """``st.async`` of 2 or 4 fp32 into a cluster peer's shared memory, crediting its mbarrier.""" + n = len(vals) + regs = ", ".join(f"${i + 1}" for i in range(n)) + _asm(None, [_v(Int32(raddr), loc, ip)] + [_v(bits(x), loc, ip) for x in vals] + + [_v(Int32(rbar), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $" + str(n + 2) + ", 0;\n\t" + f"@q st.async.shared::cluster.mbarrier::complete_tx::bytes.v{n}.b32 [$0], {{{regs}}}, [${n + 1}];\n\t}}", + ",".join(["r"] * (n + 3)), True) + +# ---- barriers +@dsl_user_op +def bar_init(addr, pred, *, loc=None, ip=None): + _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t@q mbarrier.init.shared::cta.b64 [$0], 1;\n\t}", + "r,r", True) + +@dsl_user_op +def bar_expect(addr, tx, pred, *, loc=None, ip=None): + _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t" + f"@q mbarrier.arrive.expect_tx.shared::cta.b64 _, [$0], {tx};\n\t}}", "r,r", True) + +@dsl_user_op +def bar_arrive(addr, pred, *, loc=None, ip=None): + """``mbarrier.arrive`` (release, CTA scope) when ``pred``.""" + _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t@q mbarrier.arrive.shared::cta.b64 _, [$0];\n\t}", + "r,r", True) + + +@dsl_user_op +def bar_wait(addr, pred, cluster_scope, *, loc=None, ip=None): + """Wait for phase 0 of the mbarrier at ``addr`` when ``pred``; acquire at cluster scope for peer data.""" + scope = ".acquire.cluster" if cluster_scope else "" + _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q, done;\n\t.reg .b32 ph;\n\tmov.b32 ph, 0;\n\tsetp.ne.s32 q, $1, 0;\n\t" + "@!q bra.uni WDONE;\n\tWLOOP:\n\t" + "mbarrier.try_wait.parity" + scope + ".shared::cta.b64 done, [$0], ph;\n\t" + "@!done bra.uni WLOOP;\n\tWDONE:\n\t}", "r,r", True) + +@dsl_user_op +def raw(text, *, loc=None, ip=None): + _asm(None, [], text, "", True) + +# ---- tensor memory and the MMA +@dsl_user_op +def tmem_alloc(slot, pred, *, loc=None, ip=None): + _asm(None, [_v(Int32(slot), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t" + f"@q tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [$0], {TMEM_COLS};\n\t" + "@q tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n\t}", "r,r", True) + +@dsl_user_op +def tmem_free(taddr, pred, *, loc=None, ip=None): + _asm(None, [_v(Int32(taddr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t" + f"@q tcgen05.dealloc.cta_group::1.sync.aligned.b32 $0, {TMEM_COLS};\n\t}}", "r,r", True) + +@dsl_user_op +def tmem_ld(taddr, n, *, loc=None, ip=None): + """``tcgen05.ld.16x32bx2.x{n}`` (half split 32): thread l gets lane l % 16, columns 32 (l / 16) + [0, n).""" + regs = ", ".join(f"${i}" for i in range(n)) + res = _asm(llvm.StructType.get_literal([T.f32()] * n), [_v(Int32(taddr), loc, ip)], + f"tcgen05.ld.sync.aligned.16x32bx2.x{n}.b32 {{{regs}}}, [${n}], 32;\n\t" + "tcgen05.wait::ld.sync.aligned;", ",".join(["=f"] * n) + ",r", True) + return [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(n)] + +@dsl_user_op +def tmem_st_zero32(taddr, *, loc=None, ip=None): + zeros = ", ".join(["$1"] * 32) + _asm(None, [_v(Int32(taddr), loc, ip), _v(Int32(0), loc, ip)], + "tcgen05.st.sync.aligned.16x32bx2.x32.b32 [$0], 32, {" + zeros + "};", "r,r", True) + +@dsl_user_op +def tmem_st16(taddr, words, *, loc=None, ip=None): + regs = ", ".join(f"${i + 1}" for i in range(16)) + _asm(None, [_v(Int32(taddr), loc, ip)] + [_v(Int32(w), loc, ip) for w in words], + "tcgen05.st.sync.aligned.16x32bx2.x16.b32 [$0], 16, {" + regs + "};", "r" + ",r" * 16, True) + +@dsl_user_op +def mma4(d_taddr, a_lo, b_lo, bar, pred, a_tmem, accumulate_first, idesc, a_step, b_step, *, loc=None, + ip=None): + """Four tcgen05 kind::f16 K16 steps as the split issues them, then a commit to ``bar``: D at ``d_taddr``; + A in shared memory (descriptor low word ``a_lo`` + ``a_step`` per step) or in tensor memory (column + ``a_lo`` + 8 per step); B descriptor low word ``b_lo`` + ``b_step`` per step; the first step accumulates iff + ``accumulate_first``, the others always.""" + lines = ["{", ".reg .pred q, acc0, acc1;", ".reg .b64 a, b;", ".reg .b32 al, bl, hi, id;", + "setp.ne.s32 q, $4, 0;", f"mov.pred acc0, {1 if accumulate_first else 0};", "mov.pred acc1, 1;", + f"mov.b32 hi, {DESC_HI};", f"mov.b32 id, {idesc};"] + for kk in range(4): + acc = "acc0" if kk == 0 else "acc1" + lines.append(f"add.s32 bl, $2, {b_step * kk};") + lines.append("mov.b64 b, {bl, hi};") + if a_tmem: + lines.append(f"add.s32 al, $1, {8 * kk};") + lines.append(f"@q tcgen05.mma.cta_group::1.kind::f16 [$0], [al], b, id, {acc};") + else: + lines.append(f"add.s32 al, $1, {a_step * kk};") + lines.append("mov.b64 a, {al, hi};") + lines.append(f"@q tcgen05.mma.cta_group::1.kind::f16 [$0], a, b, id, {acc};") + lines.append("@q tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [$3];") + lines.append("}") + _asm(None, [_v(Int32(d_taddr), loc, ip), _v(Int32(a_lo), loc, ip), _v(Int32(b_lo), loc, ip), + _v(Int32(bar), loc, ip), _v(Int32(pred), loc, ip)], "\n\t".join(lines), "r,r,r,r,r", True) + +@dsl_user_op +def pair2(text, a0, a1, b0, b1, *, loc=None, ip=None): + """A packed ``.f32x2`` op on (a0, a1), (b0, b1) (``text`` uses x, y -> z); two fp32.""" + res = _asm(llvm.StructType.get_literal([T.f32()] * 2), + [_v(Float32(a0), loc, ip), _v(Float32(a1), loc, ip), _v(Float32(b0), loc, ip), _v(Float32(b1), loc, ip)], + "{\n\t.reg .b64 x, y, z;\n\tmov.b64 x, {$2, $3};\n\tmov.b64 y, {$4, $5};\n\t" + text + + "\n\tmov.b64 {$0, $1}, z;\n\t}", "=f,=f,f,f,f,f", False) + return [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(2)] + + +def add2(a, b): + """``add.rn.f32x2`` of two (lo, hi) pairs: the split's and the combine's packed adds.""" + return pair2("add.rn.f32x2 z, x, y;", a[0], a[1], b[0], b[1]) + + +def max3(a, b, c): + return f3("max.f32 $0, $1, $2, $3;", a, b, c) + + +@dsl_user_op +def scale_mask2(a0, a1, qk_scale, nlo, span, e, *, loc=None, ip=None): + """Scores ``e``, ``e + 1`` of a row half: ``mul.rn.f32x2`` by the QK scale, then the split's window mask as one + unsigned range test (``e - lo <= span``) per score.""" + res = _asm(llvm.StructType.get_literal([T.f32()] * 2), + [_v(Float32(a0), loc, ip), _v(Float32(a1), loc, ip), _v(Float32(qk_scale), loc, ip), + _v(Int32(nlo), loc, ip), _v(Int32(span), loc, ip)], + "{\n\t.reg .b64 x, y, z;\n\t.reg .f32 s0, s1;\n\t.reg .b32 t;\n\t.reg .pred p;\n\t" + "mov.b64 x, {$2, $3};\n\tmov.b64 y, {$4, $4};\n\tmul.rn.f32x2 z, x, y;\n\tmov.b64 {s0, s1}, z;\n\t" + f"add.s32 t, $5, {e};\n\tsetp.le.u32 p, t, $6;\n\tselp.f32 $0, s0, 0fFF800000, p;\n\t" + f"add.s32 t, $5, {e + 1};\n\tsetp.le.u32 p, t, $6;\n\tselp.f32 $1, s1, 0fFF800000, p;\n\t}}", + "=f,=f,f,f,f,r,r", False) + return [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(2)] + + +@dsl_user_op +def tmem_ld8x(taddrs, *, loc=None, ip=None): + """``tcgen05.ld.16x32bx2.x8`` (half split 32) at each address, one ``wait::ld`` for all.""" + n = len(taddrs) + lines = [] + for b in range(n): + regs = ", ".join(f"${8 * b + i}" for i in range(8)) + lines.append(f"tcgen05.ld.sync.aligned.16x32bx2.x8.b32 {{{regs}}}, [${8 * n + b}], 32;") + lines.append("tcgen05.wait::ld.sync.aligned;") + res = _asm(llvm.StructType.get_literal([T.f32()] * (8 * n)), [_v(Int32(a), loc, ip) for a in taddrs], + "\n\t".join(lines), ",".join(["=f"] * (8 * n)) + "," + ",".join(["r"] * n), True) + vals = [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(8 * n)] + return [vals[8 * b:8 * b + 8] for b in range(n)] + + +def desc_lo(saddr): + """Low word of the split's shared-memory descriptor: the start address >> 4 (14 bits).""" + return (saddr >> 4) & 0x3FFF + +# ------------------------------------------------------------------------------------------ math blocks +def unpack8(words): + out = [] + for w in words: + out.append(lo_bf16(w)) + out.append(hi_bf16(w)) + return out + +def prep8(words, cosv, sinv, eps, scale, rope): + """``_prep_rows_math`` of this lane's 8 elements (``words``: BF16x2 of elements 8 l .. 8 l + 7 of a head; + 8 lanes per head): the squares (``mul.f32``) summed in order with ``add.rn.ftz``, the lanes paired over + bits 0, 1, 2 (``add.f32``), ``rsqrt.approx.ftz(total / 64 + eps)`` with ``div.full``, the two BF16 + roundings, then (RoPE) the pair rotation with ``mul.rn.ftz`` / ``fma.rn.ftz``. 8 BF16-exact fp32.""" + x = unpack8(words) + s = add_ftz(mul(x[0], x[0]), mul(x[1], x[1])) + for i in range(2, 8): + s = add_ftz(s, mul(x[i], x[i])) + s = add(s, shfl_bfly(s, 1)) + s = add(s, shfl_bfly(s, 2)) + s = add(s, shfl_bfly(s, 4)) + inv = rsqrt_ftz(add(eps, div_full(s, Float32(64.0)))) + y = [bf16r(mul(bf16r(mul(v, inv)), scale)) for v in x] + if rope: + out = [] + for i in range(4): + ev, od = y[2 * i], y[2 * i + 1] + out.append(bf16r(fma_ftz(ev, cosv[i], mul_ftz(neg(od), sinv[i])))) + out.append(bf16r(fma_ftz(od, cosv[i], mul_ftz(ev, sinv[i])))) + y = out + return y + +def to_cache_fp8(vals, scale_b): + """``_to_cache_fp8`` of 8 BF16-exact values: ``mul.bf16`` by the BF16 scale, then E4M3 (RN, satfinite). + 4 Int32, each two E4M3 in its low 16 bits (element 2 i in the low byte).""" + y = [bf16_mul(v, scale_b) for v in vals] + return [e4m3x2(y[2 * i + 1], y[2 * i]) for i in range(4)] + +def fp8_to_bf16(pairs): + return [e4m3x2_to_bf16x2(p) for p in pairs] + +def fp8_words_to_bf16(words): + """4 Int32 of 4 E4M3 each (16 elements) -> 8 BF16x2 words.""" + out = [] + for w in words: + out.append(e4m3x2_to_bf16x2(w & 0xFFFF)) + out.append(e4m3x2_to_bf16x2((w >> 16) & 0xFFFF)) + return out + +def store_tile_half(tile_row, key, kh, words8, valid): + """32 BF16 (16 words) of a 64-wide tile row into chunks 4 kh .. 4 kh + 3 of the 128-byte swizzle.""" + for c in range(4): + chunk = kh * 4 + c + st_s_v4(tile_row + ((chunk ^ (key % 8)) * 16), words8[4 * c:4 * c + 4], valid) + +def rope_row(g_cos, pos, l8): + """``cos[pos, 4 l8 .. 4 l8 + 3]`` and ``sin[...]`` of the fp32 ``[positions, 2, 32]`` table.""" + row = g_cos + Int64(pos) * 256 + Int64(l8 * 16) + return [as_f32(x) for x in ld_g_v4(row, Int32(1))], [as_f32(x) for x in ld_g_v4(row + 128, Int32(1))] + +def softmax(acc, qk_scale, nlo, span): + """The split's online-softmax step on one block (one row half per thread, 32 scores): the scale, the window + mask (``nlo`` = -lo_e, ``span`` = hi_e - lo_e, or an empty range), the row max (exact in any order), ``ex2``, + the row sum in the split's packed tree (pairs (p 2k, p 2k+1) added over k = 0..15 pairwise, then even + odd, + then the other half). Returns (m_new, l_new, P as 16 BF16x2 words).""" + sc = [] + for e in range(0, 32, 2): + sc += scale_mask2(acc[e], acc[e + 1], qk_scale, nlo, span, e) + mx = max_tree(sc) + m_new = max3(Float32(NEG_BIG), mx, shfl_bfly(mx, 16)) + alpha = ex2_ftz(sub(Float32(NEG_BIG), m_new)) + p = [ex2_ftz(sub(x, m_new)) for x in sc] + v = [(p[2 * k], p[2 * k + 1]) for k in range(16)] + while len(v) > 1: + v = [add2(v[2 * i], v[2 * i + 1]) for i in range(len(v) // 2)] + ssum = add(v[0][0], v[0][1]) + ssum = add(ssum, shfl_bfly(ssum, 16)) + l_new = fma(Float32(0.0), alpha, ssum) + return m_new, l_new, [pack_bf16x2(p[2 * i + 1], p[2 * i]) for i in range(16)] + + +def max_tree(v): + """Max of a list (exact in any order), three at a time.""" + v = list(v) + while len(v) > 1: + nxt = [max3(v[i], v[i + 1], v[i + 2]) for i in range(0, len(v) - 2, 3)] + if len(v) % 3 == 1: + nxt.append(v[-1]) + elif len(v) % 3 == 2: + nxt.append(fmax(v[-2], v[-1])) + v = nxt + return v[0] + + +def row_max(sml, row, half): + """Max over the splits' maxima of ``row`` and -1e30 (the padded lanes'), the two lanes of a row (``half`` 0 + and 1, lanes l and l ^ 16) each loading every other split first, then a tree and one shuffle.""" + v = [ld_s_v2f(sml + ((2 * i + half) * ROWS + row) * 8)[0] for i in range(NUM_SPLITS // 2)] + m = max_tree([Float32(NEG_BIG)] + v) + return fmax(m, shfl_bfly(m, 16)) + +def class_sums(xs): + """The combine's in-thread tree over the splits ``r + 8 k``, k = 0..7 (absent ones +0): + ((x0 + x1) + (x2 + x3)) + ((x4 + x5) + (x6 + x7)), on column pairs with ``add.rn.f32x2``.""" + out = [] + zero = (Float32(0.0), Float32(0.0)) + for i in range(0, len(xs[0]), 2): + x = [(xs[k][i], xs[k][i + 1]) for k in range(len(xs))] + [zero] * (8 - len(xs)) + out += add2(add2(add2(x[0], x[1]), add2(x[2], x[3])), add2(add2(x[4], x[5]), add2(x[6], x[7]))) + return out + +def denominators(sml, rows, lane, den_lane): + """The combine's denominator of each row in ``rows`` over its 64 split lanes (two 32-lane halves): ``w l`` + (``mul.f32``) shuffled 16, contracted into ``fma(w, l, partner)``, then lanes 8, 4, 2, 1 (``add.f32``); per + half the value of lane ``den_lane`` (the one the combine's shared-memory store keeps); the two halves added. + The chains are emitted level by level, interleaved.""" + mm = [] + for row in rows: + lo_ = ld_s_v2f(sml + (lane * ROWS + row) * 8)[0] + hi_ = selp(ld_s_v2f(sml + ((isel(lane + 32, Int32(0), lt(lane, Int32(NUM_SPLITS - 32)))) * ROWS + row) * 8)[0], + Float32(NEG_BIG), lt(lane, Int32(NUM_SPLITS - 32))) + mm.append(max3(Float32(NEG_BIG), lo_, hi_)) + for m in (16, 8, 4, 2, 1): + mm = [fmax(x, shfl_bfly(x, m)) for x in mm] + chains = [] + for i, row in enumerate(rows): + for hh in range(2): + s = 32 * hh + lane + sok = lt(s, Int32(NUM_SPLITS)) + ml = ld_s_v2f(sml + (isel(s, Int32(0), sok) * ROWS + row) * 8) + chains.append((i, sok, ml)) + + ws = [] + for i, sok, ml in chains: + m_s = selp(ml[0], Float32(NEG_BIG), sok) + l_s = selp(ml[1], Float32(0.0), sok) + ws.append((selp(selp(ex2_ftz(sub(m_s, mm[i])), Float32(0.0), gt0(l_s)), Float32(0.0), sok), l_s)) + parts = [shfl_bfly(mul(w, lv), 16) for w, lv in ws] + a = [fma(w, lv, q) for (w, lv), q in zip(ws, parts)] + for m in (8, 4, 2, 1): + sh = [shfl_bfly(x, m) for x in a] + a = [add(x, y) for x, y in zip(a, sh)] + win = [shfl_idx(x, den_lane) for x in a] + return [add(win[2 * i], win[2 * i + 1]) for i in range(len(rows))] + + +def gate_rows(num, den, sig, gate_eps): + """The combine's ``(num / den).to(bf16)`` and its output norm-gate on this thread's 4 dims (16 lanes per + head, dims 4 (lane % 16) + [0, 4)): the square tree over dims bits 4, 3, 2 (lanes 4, 2, 1), then 1, 0 + (in thread), then 5 (lane 8); the BF16 sigmoid; ``mul.bf16``. 4 BF16-exact fp32.""" + att = [bf16r(div_full(n, den)) for n in num] + lv = [add(mul(a, a), shfl_bfly(mul(a, a), 4)) for a in att] + lv = [add(x, shfl_bfly(x, 2)) for x in lv] + lv = [add(x, shfl_bfly(x, 1)) for x in lv] + tot = add(add(lv[0], lv[2]), add(lv[1], lv[3])) + tot = add(tot, shfl_bfly(tot, 8)) + rs = rsqrt_ftz(add(gate_eps, div_full(tot, Float32(64.0)))) + return [bf16_mul(bf16r(mul(rs, att[i])), sig[i]) for i in range(4)] + +class SwaClusterBlockKernel: + """The kernel of one build: PDL, the rank's KV heads, the cache page, the QKV layout, RoPE, the debug output + (0 none, 1 the partials dump, 2 per-phase clock stamps) and the denominator lane are compile-time constants. + + Thread roles (``THREADS`` = one warpgroup per owned block): warpgroup ``g`` widens block ``g``'s cached K/V, + runs its softmax and issues its PV; warpgroups 0-3 run the query prep (one 64-wide head per 8 lanes) and the + class sums (8 of each half's 32 O columns each); warpgroup 4 runs the step keys' prep and the cache write, then + the owner's denominators; warpgroup 0 finishes the owner's rows.""" + + def __init__(self, enable_pdl: bool, kv_heads: int, page: int, permuted: bool, rope: bool, debug: int, + den_lane: int): + self.enable_pdl = enable_pdl + self.hkv = kv_heads + self.page = page + self.rope = rope + self.debug = debug + self.den_lane = den_lane + self.q_stride = D if permuted else 2 * D + self.k_col0 = (GROUP * kv_heads if permuted else 2 * GROUP * kv_heads) * D + self.v_col0 = self.k_col0 + kv_heads * D + self.dbg = dbg_layout(kv_heads) + + @cute.jit + def __call__(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Pointer, + seq_ptr: cute.Pointer, scale_ptr: cute.Pointer, cos_ptr: cute.Pointer, pos_ptr: cute.Pointer, + gate_ptr: cute.Pointer, out_ptr: cute.Pointer, dbg_ptr: cute.Pointer, reqs: Int32, + qkv_stride: Int32, bt_stride_req: Int32, bt_stride_kv: Int32, gate_stride: Int32, + gate_hstride: Int32, out_stride: Int32, prep_eps: Float32, q_scale: Float32, + gate_eps: Float32, sm_scale_log2: Float32, stream): + self.kernel(qkv_ptr, pool_ptr, bt_ptr, seq_ptr, scale_ptr, cos_ptr, pos_ptr, gate_ptr, out_ptr, dbg_ptr, + qkv_stride, bt_stride_req, bt_stride_kv, gate_stride, gate_hstride, out_stride, prep_eps, + q_scale, gate_eps, sm_scale_log2).launch( + grid=(CLUSTER, self.hkv, reqs), block=[THREADS, 1, 1], cluster=(CLUSTER, 1, 1), smem=SMEM_BYTES + 1024, + stream=stream, use_pdl=self.enable_pdl) + + @cute.kernel + def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Pointer, seq_ptr: cute.Pointer, + scale_ptr: cute.Pointer, cos_ptr: cute.Pointer, pos_ptr: cute.Pointer, gate_ptr: cute.Pointer, + out_ptr: cute.Pointer, dbg_ptr: cute.Pointer, qkv_stride: Int32, bt_stride_req: Int32, + bt_stride_kv: Int32, gate_stride: Int32, gate_hstride: Int32, out_stride: Int32, + prep_eps: Float32, q_scale: Float32, gate_eps: Float32, sm_scale_log2: Float32): + tid, _, _ = cute.arch.thread_idx() + rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + _, h, b = cute.arch.block_idx() + wg = tid // 128 # warpgroup = this thread's block slot + wt = tid % 128 + warp = wt // 32 # TMEM subpartition + lane = tid % 32 + l16 = lane % 16 + half = lane // 16 + l8 = tid % 8 + is_t0 = lt(tid, Int32(1)) + # warp 16 (warpgroup 4's first) owns the tensor memory: it is idle while warpgroup 0 finishes the rows + is_w16 = lt(tid, Int32(544)) * (1 - lt(tid, Int32(512))) + lead = lt(wt, Int32(1)) + prep_q = lt(tid, Int32(512)) # warpgroups 0-3 + one = Int32(1) + + smem = utils.SmemAllocator() + sbase = smem.allocate(SMEM_BYTES, byte_alignment=1024).toint() + g_qkv = qkv_ptr.toint() + g_pool = pool_ptr.toint() + g_bt = bt_ptr.toint() + Int64(b) * Int64(bt_stride_req) * 4 + g_cos = cos_ptr.toint() + g_pos = pos_ptr.toint() + g_out = out_ptr.toint() + g_dbg = dbg_ptr.toint() + g_ts = g_dbg + Int64(self.dbg["ts"]) + Int64(((b * self.hkv + h) * CLUSTER + rank) * STAMPS * 8) + sml = sbase + ML_OFF + s_bar = sbase + BAR_OFF + o_bar = s_bar + 40 + ml_bar = s_bar + 80 + t_bar = s_bar + 88 + if cutlass.const_expr(self.debug == 2): + stamp(g_ts, is_t0) + for k in cutlass.range_constexpr(NB_MAX): + bar_init(s_bar + 8 * k, is_t0) + bar_init(o_bar + 8 * k, is_t0) + bar_init(ml_bar, is_t0) + bar_init(t_bar, is_t0) + raw("fence.mbarrier_init.release.cluster;") + bar_expect(ml_bar, ML_BYTES, is_t0) + bar_expect(t_bar, T_BYTES, is_t0) + + # ------------------------------------------------ before the wait: graph inputs and the cached K/V + # round trip 1: the length, the scale, this thread's RoPE position (query token tid / 64, or the step key + # (tid % 128) / 8 of warpgroup 4) + L = ld_g_b32(seq_ptr.toint() + Int64(b) * 4, one) + scale = as_f32(ld_g_b32(scale_ptr.toint(), one)) + tq = isel(tid // 64, (wt // 8) % TOKENS, prep_q) + tw = (wt // 8) % TOKENS + if cutlass.const_expr(self.rope): + pos = ld_g_b32(g_pos + Int64(b * TOKENS + tq) * 4, one) + past = L - TOKENS + lo = cutlass.max(past - (WINDOW - 1), Int32(0)) + base = (lo // BLOCK) * BLOCK + nblk = (L - base + (BLOCK - 1)) // BLOCK + scale_b = bf16r(scale) + dequant = div_full(Float32(1.0), scale) + qk_scale = mul(mul(sm_scale_log2, dequant), dequant) + s = rank + CLUSTER * wg # this warpgroup's split + live = lt(s, nblk) + valid = lt(s, Int32(NUM_SPLITS)) + # round trip 2: the block's page ids, the writer's page ids, the RoPE row + key = wt // 2 + kh = wt % 2 + j = base + s * BLOCK + key + cached = live * lt(j, past) * (1 - lt(j, lo)) + pg = j // self.page + pk = ld_g_b32(g_bt + Int64(pg) * 4, cached) + pv = ld_g_b32(g_bt + (Int64(bt_stride_kv) + Int64(pg)) * 4, cached) + jw = past + tw + writer = lt(rank, one) * (1 - prep_q) * lt(wt, Int32(64)) + wpk = ld_g_b32(g_bt + Int64(jw // self.page) * 4, writer) + wpv = ld_g_b32(g_bt + (Int64(bt_stride_kv) + Int64(jw // self.page)) * 4, writer) + if cutlass.const_expr(self.rope): + cosv, sinv = rope_row(g_cos, pos, l8) + else: + cosv = None + sinv = None + # round trip 3: the block's cached K/V (FP8, 32 bytes of a key row per thread) + inpage = Int64((j % self.page) * D + kh * 32) + ka = g_pool + (Int64(pk) * self.hkv + Int64(h)) * (self.page * D) + inpage + va = g_pool + (Int64(pv) * self.hkv + Int64(h)) * (self.page * D) + inpage + kw = ld_g_v4(ka, cached) + ld_g_v4(ka + 16, cached) + vw = ld_g_v4(va, cached) + ld_g_v4(va + 16, cached) + tile_off = wg * TILE_BYTES + key * 128 + store_tile_half(sbase + K_OFF + tile_off, key, kh, fp8_words_to_bf16(kw), valid) + store_tile_half(sbase + V_OFF + tile_off, key, kh, fp8_words_to_bf16(vw), valid) + + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + if cutlass.const_expr(self.debug == 2): + stamp(g_ts + 8, is_t0) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(self.debug == 2): + stamp(g_ts + 16, is_t0) + + # ------------------------------------------------ after the wait: the projection, one round trip + qr = tid // 8 # query row of warpgroups 0-3 (token qr / 8, head qr % 8) + qa = g_qkv + (Int64(b * TOKENS + qr // 8) * Int64(qkv_stride) + + Int64((h * GROUP + qr % 8) * self.q_stride + l8 * 8)) * 2 + row_k = Int64(b * TOKENS + tw) * Int64(qkv_stride) + step = 1 - prep_q + qw = ld_g_v4(qa, prep_q) + kraw = ld_g_v4(g_qkv + (row_k + Int64(self.k_col0 + h * D + l8 * 8)) * 2, step) + vraw = ld_g_v4(g_qkv + (row_k + Int64(self.v_col0 + h * D + l8 * 8)) * 2, step) + og = wt // 16 + od = wt % 16 + ga = gate_ptr.toint() + (Int64(b * TOKENS + rank) * Int64(gate_stride) + + Int64((h * GROUP + og) * gate_hstride + od * 4)) * 2 + gw = ld_g_v2(ga, lt(tid, Int32(128))) + + # Tensor memory only now, with the projection loads in flight: after the wait every predecessor has + # completed (and freed its columns), and before the trigger no dependent can hold any, so the allocation + # never waits on a kernel that waits on this one. The dependents launch once every CTA holds its columns. + tmem_alloc(sbase + SLOT_OFF, is_w16) + raw("tcgen05.fence::before_thread_sync;") + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + raw("tcgen05.fence::after_thread_sync;") + tmem = ld_s_b32(sbase + SLOT_OFF) + + # warpgroups 0-3: the query rows; warpgroup 4: the step keys (one prep each: 8 lanes per head) + y = prep8(isel4(qw, kraw, prep_q), cosv, sinv, prep_eps, selp(q_scale, Float32(1.0), prep_q), self.rope) + qb = fp8_to_bf16(to_cache_fp8(y, scale_b)) + st_s_v4(sbase + Q_OFF + qr * 128 + ((l8 ^ (qr % 8)) * 16), qb, prep_q) + if cutlass.const_expr(self.debug == 1): + dq = g_dbg + Int64(self.dbg["q"]) + Int64(((b * self.hkv + h) * ROWS + qr) * D + l8 * 8) * 2 + st_g_v2(dq, qb[0], qb[1], prep_q) + st_g_v2(dq + 8, qb[2], qb[3], prep_q) + k8 = to_cache_fp8(y, scale_b) + v8 = to_cache_fp8(unpack8(vraw), scale_b) + inw = Int64((jw % self.page) * D + l8 * 8) + st_g_v2(g_pool + (Int64(wpk) * self.hkv + Int64(h)) * (self.page * D) + inw, k8[0] | (k8[1] << 16), + k8[2] | (k8[3] << 16), writer) + st_g_v2(g_pool + (Int64(wpv) * self.hkv + Int64(h)) * (self.page * D) + inw, v8[0] | (v8[1] << 16), + v8[2] | (v8[3] << 16), writer) + sj = (jw - base) // BLOCK + kj = (jw - base) % BLOCK + mine = step * lt(wt, Int32(64)) * (1 - lt(sj % CLUSTER, rank)) * (1 - lt(rank, sj % CLUSTER)) + prow = (sj // CLUSTER) * TILE_BYTES + kj * 128 + ((l8 ^ (kj % 8)) * 16) + st_s_v4(sbase + K_OFF + prow, qb, mine) + st_s_v4(sbase + V_OFF + prow, fp8_to_bf16(v8), mine) + raw("fence.proxy.async.shared::cta;") + cute.arch.barrier() + if cutlass.const_expr(self.debug == 2): + stamp(g_ts + 24, is_t0) + + # ------------------------------------------------ QK: each warpgroup's lead thread issues its block's + mma4(tmem + (16 << 16) + COLS_PER_BLOCK * wg, desc_lo(sbase + Q_OFF), desc_lo(sbase + K_OFF + wg * TILE_BYTES), + s_bar + 8 * wg, lead * live, False, False, IDESC_QK, 2, 2) + + # ------------------------------------------------ each warpgroup: its block's softmax, PV, (max, sum) out + row = 16 * warp + l16 + q_pos = past + row // 8 + lane_tm = (32 * warp) << 16 + if valid != 0: + bar_wait(s_bar + 8 * wg, live, False) + raw("tcgen05.fence::after_thread_sync;") + col = tmem + lane_tm + COLS_PER_BLOCK * wg + jb = base + s * BLOCK + 32 * half + lo_e = cutlass.max(lo, q_pos - (WINDOW - 1)) - jb + hi_e = q_pos - jb + empty = lt(hi_e, lo_e) + m_new, l_new, pw = softmax(tmem_ld(col + (16 << 16), 32), qk_scale, isel(Int32(-64), 0 - lo_e, empty), + isel(Int32(0), hi_e - lo_e, empty)) + m_k = selp(m_new, Float32(NEG_BIG), live) + l_k = selp(l_new, Float32(0.0), live) + tmem_st16(col + 64, pw) + tmem_st_zero32(col) + raw("tcgen05.wait::st.sync.aligned;") + raw("tcgen05.fence::before_thread_sync;") + bar_named(1 + wg, 128) + raw("tcgen05.fence::after_thread_sync;") + mma4(tmem + COLS_PER_BLOCK * wg, tmem + COLS_PER_BLOCK * wg + 64, + desc_lo(sbase + V_OFF + wg * TILE_BYTES), o_bar + 8 * wg, lead * live, True, True, IDESC_PV, 0, 128) + # a split with no keys: O stays the zeros just stored (the split's acc * dequant of an empty loop) + bar_arrive(o_bar + 8 * wg, lead * (1 - live)) + slot = sml + (s * ROWS + row) * 8 + for dd in cutlass.range_constexpr(4): + st_async(mapa(slot, half * 4 + dd), [m_k, l_k], mapa(ml_bar, half * 4 + dd), valid) + if cutlass.const_expr(self.debug == 1): + dml = g_dbg + Int64(self.dbg["ml"]) + Int64((((b * NUM_SPLITS + s) * self.hkv + h) * ROWS + row) * 8) + st_g_v2(dml, bits(m_k), bits(l_k), valid * (1 - half)) + if cutlass.const_expr(self.debug == 2): + stamp(g_ts + 32, is_t0) + + if valid == 0: + # an absent block slot (rank >= 2, slot 4): O = 0, so the class sums read it as an empty split + tmem_st_zero32(tmem + lane_tm + COLS_PER_BLOCK * wg) + raw("tcgen05.wait::st.sync.aligned;") + raw("tcgen05.fence::before_thread_sync;") + bar_named(1 + wg, 128) + bar_arrive(o_bar + 8 * wg, lead) + + # ------------------------------------------------ every split's (max, sum) in + sig = [sigmoid_bf16(x) for x in unpack8(gw)] + bar_wait(ml_bar, one, True) + if cutlass.const_expr(self.debug == 2): + stamp(g_ts + 40, is_t0) + # warpgroups 0-3: the class sums of 8 of each half's 32 O columns, to the token's owner + if wg < 4: + chunk = wg % 4 + cls_ok = prep_q + xs = [] + mls = [] + for k in cutlass.range_constexpr(NB_MAX): + sk = rank + CLUSTER * k + mls.append(ld_s_v2f(sml + (isel(sk, Int32(0), lt(sk, Int32(NUM_SPLITS))) * ROWS + row) * 8)) + m_max = row_max(sml, row, half) + wks = [] + for k in cutlass.range_constexpr(NB_MAX): + wks.append(selp(ex2_ftz(sub(mls[k][0], m_max)), Float32(0.0), gt0(mls[k][1]))) + for k in cutlass.range_constexpr(NB_MAX): + bar_wait(o_bar + 8 * k, cls_ok, False) + raw("tcgen05.fence::after_thread_sync;") + if cutlass.const_expr(self.debug == 2): + stamp(g_ts + 72, is_t0) + os_ = tmem_ld8x([tmem + lane_tm + COLS_PER_BLOCK * k + 8 * chunk for k in range(NB_MAX)]) + for k in cutlass.range_constexpr(NB_MAX): + sk = rank + CLUSTER * k + lk = lt(sk, nblk) + w_k = wks[k] + o = os_[k] + xk = [] + for i in cutlass.range_constexpr(0, 8, 2): + okp = pair2("mul.rn.f32x2 z, x, y;", o[i], o[i + 1], dequant, dequant) + xk += pair2("mul.rn.f32x2 z, x, y;", w_k, w_k, okp[0], okp[1]) + for i in cutlass.range_constexpr(8): + ok = mul(o[i], dequant) + if cutlass.const_expr(self.debug == 1): + dob = g_dbg + Int64(self.dbg["o"]) + Int64( + ((((b * NUM_SPLITS + sk) * self.hkv + h) * ROWS + row) * D + 32 * half + 8 * chunk + i) * 4) + st_g_b32(dob, bits(selp(ok, Float32(0.0), lk)), lt(sk, Int32(NUM_SPLITS)) * cls_ok) + xs.append(xk) + tsum = class_sums(xs) + dst = sbase + T_OFF + ((rank * GROUP + row % 8) * D + 32 * half + 8 * chunk) * 4 + owner = row // 8 + st_async(mapa(dst, owner), tsum[0:4], mapa(t_bar, owner), cls_ok) + st_async(mapa(dst + 16, owner), tsum[4:8], mapa(t_bar, owner), cls_ok) + if cutlass.const_expr(self.debug == 2): + stamp(g_ts + 88, is_t0) + # warpgroup 4: the owner's denominators (warp i: heads 2 i, 2 i + 1), the combine's butterfly + if wg == 4: + dens = denominators(sml, [rank * GROUP + 2 * warp, rank * GROUP + 2 * warp + 1], lane, self.den_lane) + st_s_b32_pred(sbase + DEN_OFF + (warp * 2) * 4, bits(dens[0]), step) + st_s_b32_pred(sbase + DEN_OFF + (warp * 2 + 1) * 4, bits(dens[1]), step) + if cutlass.const_expr(self.debug == 2): + stamp(g_ts + 80, lt(wt, one) * step) + raw("tcgen05.fence::before_thread_sync;") + cute.arch.barrier() + raw("tcgen05.fence::after_thread_sync;") + tmem_free(tmem, is_w16) # every tensor-memory read is done (the barrier above) + if cutlass.const_expr(self.debug == 2): + stamp(g_ts + 48, is_t0) + + # ------------------------------------------------ warpgroup 0, the owner (token = rank): finish the rows + fin = lt(tid, Int32(128)) + bar_wait(t_bar, fin, True) + if cutlass.const_expr(self.debug == 2): + stamp(g_ts + 56, is_t0) + den = as_f32(ld_s_b32(sbase + DEN_OFF + og * 4)) + tv = [] + for cls in cutlass.range_constexpr(CLUSTER): + tv.append(ld_s_v4f(sbase + T_OFF + ((cls * GROUP + og) * D + od * 4) * 4)) + num = [] + for i in cutlass.range_constexpr(4): + pw4 = [] + for w2 in cutlass.range_constexpr(4): + pw4.append(add(tv[2 * w2][i], tv[2 * w2 + 1][i])) + num.append(add(add(pw4[0], pw4[2]), add(pw4[1], pw4[3]))) + res = gate_rows(num, den, sig, gate_eps) + oa = g_out + (Int64(b * TOKENS + rank) * Int64(out_stride) + Int64((h * GROUP + og) * D + od * 4)) * 2 + st_g_v2(oa, pack_bf16x2(res[1], res[0]), pack_bf16x2(res[3], res[2]), fin) + if cutlass.const_expr(self.debug == 2): + stamp(g_ts + 64, is_t0) + + +_compiled: dict = {} + + +def compiled_key(device: int, enable_pdl: bool, kv_heads: int, page: int, permuted: bool, rope: bool, + debug: int = 0, den_lane: int = 0) -> tuple: + return (device, bool(enable_pdl), int(kv_heads), int(page), bool(permuted), bool(rope), int(debug), + int(den_lane)) + + +def _compile(key: tuple): + fn = _compiled.get(key) + if fn is not None: + return fn + import cuda.bindings.driver as cuda + from cutlass.cute.runtime import make_ptr + + device, pdl, hkv, page, permuted, rope, debug, den_lane = key + kernel = SwaClusterBlockKernel(pdl, hkv, page, permuted, rope, debug, den_lane) + ptrs = [make_ptr(dt, 16, cute.AddressSpace.gmem, assumed_align=16) for dt in ( + cutlass.BFloat16, cutlass.Uint8, cutlass.Int32, cutlass.Int32, cutlass.Float32, cutlass.Float32, + cutlass.Int32, cutlass.BFloat16, cutlass.BFloat16, cutlass.Uint8)] + with torch.cuda.device(device): + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + fn = cute.compile(kernel, *ptrs, *[cutlass.Int32(1)] * 7, *[cutlass.Float32(0.0)] * 4, stream, + options="--enable-tvm-ffi") + _compiled[key] = fn + return fn + + +def swa_cluster_block(qkv: torch.Tensor, kv_pool: torch.Tensor, block_tables: torch.Tensor, seq_lens: torch.Tensor, + kv_scale: torch.Tensor, num_kv_heads: int, cos_sin: Optional[torch.Tensor], + positions: Optional[torch.Tensor], prep_eps: float, q_scale: float, qkv_permuted: bool, + gate: torch.Tensor, gate_head_stride: int, gate_eps: float, enable_pdl: bool, + debug: Optional[torch.Tensor] = None, den_lane: int = 0, stamps: bool = False) -> torch.Tensor: + """The block of ``seq_lens.numel()`` sequences of 8 rows each; the caller checked ``fits`` and the multi-request + op's preconditions. ``debug`` (uint8, ``dbg_layout(G)["bytes"]``) receives the Q tiles and the split partials + in the split kernel's layout; ``den_lane`` exists for the build's tests.""" + reqs = seq_lens.numel() + out = torch.empty((qkv.shape[0], GROUP * num_kv_heads * D), dtype=torch.bfloat16, device=qkv.device) + rope = cos_sin is not None + key = compiled_key(qkv.device.index, enable_pdl, num_kv_heads, kv_pool.shape[2], qkv_permuted, rope, + 0 if debug is None else (2 if stamps else 1), den_lane) + if torch.cuda.is_current_stream_capturing() and key not in _compiled: + raise RuntimeError("swa_cluster_block compiles on its first call per device and build; make that call " + "before capturing a CUDA graph") + fn = _compile(key) + import cuda.bindings.driver as cuda + + fn(qkv.data_ptr(), kv_pool.data_ptr(), block_tables.data_ptr(), seq_lens.data_ptr(), kv_scale.data_ptr(), + (cos_sin if rope else kv_scale).data_ptr(), (positions if rope else seq_lens).data_ptr(), gate.data_ptr(), + out.data_ptr(), (debug if debug is not None else out).data_ptr(), reqs, qkv.stride(0), block_tables.stride(0), + block_tables.stride(1), gate.stride(0), gate_head_stride, out.stride(0), float(prep_eps), float(q_scale), + float(gate_eps), math.log2(math.e) / math.sqrt(D), + cuda.CUstream(torch.cuda.current_stream(qkv.device).cuda_stream)) + return out + + +__all__ = ["CLUSTER", "NUM_SPLITS", "dbg_layout", "fits", "owned_splits", "swa_cluster_block", "compiled_key"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_direct_route.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_direct_route.py index 937b6b7405e7..3792cacf136c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_direct_route.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_direct_route.py @@ -61,8 +61,10 @@ records, after its own dependency wait. A poll that has seen no arrival for ``guard_ns`` sets a guard word in host memory (and keeps polling). -Wide transport (TP4, ``max_rows=WIDE_MAX_ROWS``): the decode steps of two to four -requests (16..32 rows) run the wide tagged kernel's statements, in +Wide transport (``max_rows=WIDE_MAX_ROWS``; upstream b5ca973 / #21377 at TP4, here also TP8 +routing-only, ``tep8_wide_direct_pool`` beside the TEP8 router shard's wide +route): the decode steps of two to four +requests (16..32 rows, ``wide_prologue``) run the wide tagged kernel's statements, in which each CTA routes only its own two rows (no placement, no early ids). A record is then the source rank's ``[32, C]`` pool slot followed by one ``[32, L]`` FP32 logit block, the router's output as it is (``router_shard_rows`` with the block as its @@ -150,11 +152,11 @@ def pool_row_words(world: int) -> int: def check_layout(world: int, max_rows: int) -> None: - """Raise unless ``(world, max_rows)`` is a served layout: 8 rows at TP8 or TP4, 32 rows at TP4.""" + """Raise unless ``(world, max_rows)`` is a served layout: 8 rows at TP8 or TP4, 32 rows at TP4 + or TP8 (at TP8 routing only, ``PLACE_NONE``: the tagged kernel's wide statements are the + TP-generic ones, each CTA's ``[2, L]`` logit words; the half window's census is TP4's).""" if world not in WORLDS or max_rows not in (MAX_ROWS, WIDE_MAX_ROWS): raise ValueError(f"no direct pool layout for TP{world} and {max_rows} rows") - if max_rows == WIDE_MAX_ROWS and world != 4: - raise ValueError("the wide direct pool layout is TP4 only") def pool_slot_words(world: int, max_rows: int = MAX_ROWS) -> int: @@ -364,7 +366,9 @@ def __init__( check_layout(world, max_rows) wide = max_rows > MAX_ROWS if wide and ( - place not in (PLACE_NONE, PLACE_HALF) or (place == PLACE_HALF) != bool(census) + place not in (PLACE_NONE, PLACE_HALF) + or (place == PLACE_HALF) != bool(census) + or (place == PLACE_HALF and world != 4) ): raise ValueError( "the wide layout routes only, or places the half window with its census words" @@ -1280,7 +1284,7 @@ class PoolDirectWorkspace: slot of ``buffer`` (512 columns at TP8, 1024 at TP4) and ``boxes(buffer, rows)`` its ``[publish_copies(rows), box_words(world)]`` FP32 logit boxes (the 8-row layout) or ``block(buffer, rows)`` its ``[rows, L]`` FP32 logit block (the wide - layout, ``max_rows=WIDE_MAX_ROWS``, TP4), all on the multicast address: a producer + layout, ``max_rows=WIDE_MAX_ROWS``, TP4 or TP8), all on the multicast address: a producer given them as its output publishes into every rank's copy. ``kernel(variant)`` is the compiled routed kernel for a call's settings. Construction is collective over the group (fabric allocation and sentinel fill), and so is ``build_variants``; both succeed on @@ -1848,7 +1852,11 @@ def is_valid( layout routes only (``PLACE_NONE``, no early ids).""" if not 1 <= rows <= workspace.max_rows: raise ValueError(f"rows must be in 1..{workspace.max_rows}, got {rows}") - if workspace.wide and (place not in (PLACE_NONE, PLACE_HALF) or early is not None): + if workspace.wide and ( + place not in (PLACE_NONE, PLACE_HALF) + or early is not None + or (place == PLACE_HALF and workspace.world != 4) + ): raise ValueError( "the wide layout routes only or places the half window (PLACE_NONE or PLACE_HALF), " "without early ids" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_gather_route.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_gather_route.py index ecdbc1704e5a..c1c1cfa2b7eb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_gather_route.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_gather_route.py @@ -58,7 +58,8 @@ it, which the quantizing warps advance only after the routing warps signalled that they have read it. -Wide calls (TP4): on an inbox of ``WIDE_MAX_ROWS`` rows +Wide calls (``wide_prologue``: TP4, and TP8 with ``tep8_router_shard``): on an +inbox of ``WIDE_MAX_ROWS`` rows (``ProjectionInbox(max_rows=WIDE_MAX_ROWS)``) a call may carry 1..32 rows, routing only (``PLACE_NONE``; at more than 8 rows the experts are the built-in ones, which read only the routing memo). Each CTA then exchanges only its own rows' logit blocks (the @@ -1764,7 +1765,7 @@ class PoolGatherRoute(OpWrapper): logits, a delayed rank, row-count changes with tag wrap-around and CUDA-graph replays. Rows 1..8 only. At TP4 the same claims hold against the built-in chain at TP4 (the ring-buddy placement against ``reference``'s ring flow over - four 64-expert blocks). On a TP4 wide inbox (``WIDE_MAX_ROWS``): rows 1..32, + four 64-expert blocks). On a TP4 (TP8) wide inbox (``WIDE_MAX_ROWS``): rows 1..32, routing only, the same claims for the logits, the routing and the quantization. """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/projection_publish.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/projection_publish.py index 41177bf3f1f2..0a4b2f190234 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/projection_publish.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/projection_publish.py @@ -100,6 +100,11 @@ WORLDS = (4, 8) #: The decode rows the transports serve: five to 64 requests at K = 7 (and the counts between). MIN_ROWS, MAX_ROWS = 33, 512 +#: The row capacity of a wide transport: the decode steps of two to four requests at K = 7 +#: (16..32 rows), whose slice GEMMs are flashinfer's split-K cells of ``wide_projection_gathers`` +#: (``wide_direct_transports``, sm103 TEP8, behind ``tep8_wide_direct_transports``). +#: Its own allocation: the 33..512-row transports are unchanged. +WIDE_TRANSPORT_ROWS = 32 #: Width of the pooled (expert) activation and of the residual stream. POOL_HIDDEN, UNPOOL_HIDDEN = 4096, 8192 #: Buffers of a transport (``buffer_sequence``). @@ -240,8 +245,10 @@ def __init__(self, mapping: Mapping, columns: int, max_rows: int = MAX_ROWS, com world = mapping.tp_size if columns not in (local_columns(POOL_HIDDEN, world), local_columns(UNPOOL_HIDDEN, world)): raise ValueError(f"no projection site has {columns} columns per rank at TP{world}") - if not MIN_ROWS <= max_rows <= MAX_ROWS: - raise ValueError(f"max_rows must be in {MIN_ROWS}..{MAX_ROWS}, got {max_rows}") + if not (max_rows == WIDE_TRANSPORT_ROWS or MIN_ROWS <= max_rows <= MAX_ROWS): + raise ValueError( + f"max_rows must be {WIDE_TRANSPORT_ROWS} or in {MIN_ROWS}..{MAX_ROWS}, got {max_rows}" + ) self.rank = mapping.tp_rank self.world = mapping.tp_size self.columns = int(columns) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/unpool_direct_boundary.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/unpool_direct_boundary.py index 3a78a6ebc487..7fdc403a283f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/unpool_direct_boundary.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/unpool_direct_boundary.py @@ -56,11 +56,14 @@ keeps polling): a call sequence that differs between ranks, or a peer that stopped, shows up there, and ``start_guard_watch`` logs it from the host. -Wide transport (TP4, ``max_rows=WIDE_MAX_ROWS``): the decode steps of two to four -requests (16..32 rows) take a separate workspace whose slots hold 32 -rows (``3 buffers x 4 source ranks x 32 rows x 2048``), with its own buffer sequence; -the kernel is the gather boundary's at those row counts (a cluster of 4 CTAs, one per -peer slice) with the slot stride of 32 rows. The 8-row layout is unchanged. +Wide transport (``max_rows=WIDE_MAX_ROWS``): the decode steps of two to four +requests (16..32 rows) take a separate workspace whose slots hold 32 rows +(``3 buffers x TP source ranks x 32 rows x 8192 / TP``), with its own buffer sequence; +the kernel is the gather boundary's at those row counts (a cluster of 4 CTAs of 256 +threads: one peer slice per CTA at TP4, two at TP8 -- each thread picks its one peer from +its own columns, ``unpool_gather_boundary._UnpoolGatherBoundaryKernel``) with the slot stride of +32 rows. TP8 takes it under ``wide_direct_transports`` (sm103 TEP8, +behind ``tep8_wide_direct_transports``). The 8-row layout is unchanged. PDL memory order: before ``griddepcontrol.wait`` the kernel touches shared memory and the cluster barrier only (and, with ``trigger="entry"``, releases @@ -115,11 +118,10 @@ def workspace_bytes(max_rows: int = MAX_ROWS) -> int: def check_layout(world: int, max_rows: int) -> None: - """Raise unless ``(world, max_rows)`` is a served layout: 8 rows at TP4 or TP8, 32 rows at TP4.""" + """Raise unless ``(world, max_rows)`` is a served layout: 8 or ``WIDE_MAX_ROWS`` rows per slot + at TP4 or TP8 (the gather boundary feeds both TP sizes at 9..32 rows).""" if world not in WORLDS or max_rows not in (MAX_ROWS, WIDE_MAX_ROWS): raise ValueError(f"no direct unpool layout for TP{world} and {max_rows} rows") - if max_rows == WIDE_MAX_ROWS and world != 4: - raise ValueError("the wide direct unpool layout is TP4 only") #: Per-rank control words, in pinned host memory the kernel stores to at system scope: @@ -590,7 +592,7 @@ class UnpoolDirectWorkspace: ``slot(buffer, rows)`` is this rank's ``[rows, 8192 / TP]`` BF16 slot of buffer ``buffer`` on the multicast address: a GEMM given it as its output publishes into every rank's copy. ``max_rows`` is the slots' row capacity: 8, or - ``WIDE_MAX_ROWS`` for the wide layout (TP4). ``kernel(variant)`` is the compiled boundary for a + ``WIDE_MAX_ROWS`` for the wide layout. ``kernel(variant)`` is the compiled boundary for a call's settings. Construction is collective over the group (fabric allocation and sentinel fill), and so is ``build_variants``; both succeed on every rank or raise on every rank. Nothing compiles inside a CUDA graph diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/batch_splitk_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/batch_splitk_gemm.py new file mode 100644 index 000000000000..ff9d58a6f64b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/batch_splitk_gemm.py @@ -0,0 +1,805 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Blackwell BF16 cluster split-K GEMM for decode steps of 33..256 rows (``batch_splitk``). + +flashinfer's low-M split-K kernel (``flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py``, +flashinfer-python 0.6.18; ``dense_gemm``'s ``splitk`` family) serves at most 32 public rows: its +kernel-N tile, which carries the public rows, is 8, 16 or 32 wide, and its owner CTA receives every +peer's whole FP32 tile through DSMEM. Above 32 rows the built-in ``Linear`` sends a TP8 rank's +projections to cuBLAS, which picks a split-K kernel plus a separate ``cublasLt::splitKreduce`` +launch for most of them. + +This is the same kernel with the public rows in one kernel-N tile of up to 256 (any multiple of 8 +at ``mma_m`` 64, of 16 at ``mma_m`` 128: the tcgen05 MMA shapes), so each weight slab is streamed +from HBM once whatever the row count, and with the reduce-scatter epilogue of this target's +upstream copy of that kernel (TensorRT-LLM ``feat/musespark-open-agentx``, +``muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/splitk_gemm.py``, ``rs_split_k``): each of +the ``split_k`` CTAs of a cluster owns ``128 / split_k`` of the epilogue's threads, receives those +threads' FP32 partials from the other ranks through DSMEM, sums the ``split_k`` partials in +cluster-rank order, rounds once to BF16 and stores them. A CTA's mailbox is ``(split_k - 1) / +split_k`` of a tile, so the wide tiles keep a deep pipeline, and no separate reduction kernel runs. +It also carries flashinfer's ordering fixes of the TMEM and DSMEM hand-offs (flashinfer PR #7443, +854a896b: ``tcgen05.fence`` around the mbarrier hand-offs, the TMEM release arrival in phase 1, an +``.acquire.cluster`` wait on the reduction barrier). + +Kernel coordinates: the public ``x[M, K] @ weight[N, K]^T`` is swapped, so kernel-M carries the +public N (weight rows, ``mma_m``) and kernel-N the public M (rows, ``mma_n``). Each CTA of a +cluster of ``split_k`` accumulates one equal K slice in FP32 in TMEM. + +Numerics: every output element is the FP32 sum of the ``split_k`` K-slice partials in cluster-rank +order, each slice accumulated by the tensor core tile by tile along K, rounded once to BF16. It is +deterministic and does not depend on the row count or ``mma_n`` (a row's sum never involves +another row), only on ``split_k``; it is not cuBLAS's summation order, so it can differ from the +built-in's output in the last BF16 place. + +Stores: each output element is stored exactly once, by the CTA that owns it, after the +reduction, which follows the kernel's dependency wait (the activation is loaded only after +``griddepcontrol.wait``); there is no read-back of the output. So the output may be a slot of a +multicast transport, as flashinfer's split-K kernel's output is (``unpool_direct``). + +Programmatic dependent launch (``pdl``): the weight's TMA loads are issued before +``griddepcontrol.wait`` (immutable bytes), the activation's after it; the weight warp triggers +the dependents once it has issued its last load. +""" + +from __future__ import annotations + +import dataclasses +import functools +from typing import Optional, Tuple + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass import Int32 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA shared memory CuTe DSL reports on SM100 / SM103. +SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") +#: K extent of one CTA k-tile. +CTA_K = 128 +#: Kernel-M tiles (public N per CTA). +SUPPORTED_MMA_M = (64, 128) +#: Kernel-N (public rows per CTA): multiples of these, up to ``MAX_MMA_N`` (tcgen05 MMA shapes). +MMA_N_STEP = {64: 8, 128: 16} +MAX_MMA_N = 256 +#: Cluster sizes along K (portable cluster sizes; 128 epilogue threads split evenly). +SUPPORTED_SPLIT_K = (1, 2, 4, 8) +#: Largest public row count a call may have (several kernel-N tiles above ``mma_n``). +MAX_ROWS = 256 +#: FP32 accumulator values each epilogue thread holds: at most this many (registers). +MAX_VALUES_PER_THREAD = 128 +MAX_AB_STAGES = 16 + +_EPILOG_THREADS = 128 +_FP32_BYTES = 4 +_AB_ELEMENT_BYTES = 2 +_AB_BUFFER_ALIGN_BYTES = 1024 +_MAILBOX_ALIGN_BYTES = 128 +_MBARRIER_BYTES = 8 +_TMEM_POINTER_BYTES = 4 +_TMEM_COLUMNS = 256 + + +@dataclasses.dataclass(frozen=True) +class BatchSplitKTactic: + """One specialization: ``mma_m`` weight rows x ``mma_n`` public rows per CTA, ``split_k`` CTAs + per cluster along K, ``ab_stages`` k-tiles in flight.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def parse(kernel: str) -> BatchSplitKTactic: + """``"bsplitk:,,,"`` -> the tactic. + + Raises: + ValueError: for any other string. + """ + family, _, rest = str(kernel).partition(":") + fields = rest.split(",") + if family == FAMILY and len(fields) == 4 and all(f.strip().isdigit() for f in fields): + return BatchSplitKTactic(*(int(f) for f in fields)) + raise ValueError( + f"unknown batch split-K kernel {kernel!r}; expected '{FAMILY}:,,,'" + ) + + +#: The kernel string family (``parse``). +FAMILY = "bsplitk" + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def values_per_thread(tactic: BatchSplitKTactic) -> int: + """FP32 accumulator values one epilogue thread holds (``mma_m * mma_n / 128``).""" + return tactic.mma_m * tactic.mma_n // _EPILOG_THREADS + + +def mailbox_bytes(tactic: BatchSplitKTactic) -> int: + """Bytes of FP32 partials a CTA receives: its thread group's values from each other rank.""" + if tactic.split_k == 1: + return 0 + group = _EPILOG_THREADS // tactic.split_k + return (tactic.split_k - 1) * group * values_per_thread(tactic) * _FP32_BYTES + + +def smem_bytes(tactic: BatchSplitKTactic, ab_stages: Optional[int] = None) -> int: + """Shared memory the kernel's allocator lays out (A ring, B ring, barriers, mailbox).""" + stages = tactic.ab_stages if ab_stages is None else ab_stages + cursor = _align_up(tactic.mma_m * CTA_K * _AB_ELEMENT_BYTES * stages, _AB_BUFFER_ALIGN_BYTES) + cursor += tactic.mma_n * CTA_K * _AB_ELEMENT_BYTES * stages + cursor = _align_up(cursor, _MBARRIER_BYTES) + (2 * stages + 3) * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + _TMEM_POINTER_BYTES + if tactic.split_k == 1: + return cursor + cursor = _align_up(cursor, _MAILBOX_ALIGN_BYTES) + mailbox_bytes(tactic) + return _align_up(cursor, _MBARRIER_BYTES) + _MBARRIER_BYTES + + +def max_stages(tactic: BatchSplitKTactic, smem_capacity: int = SMEM_CAPACITY_BYTES) -> int: + """The deepest ring that fits ``smem_capacity`` (0 if none).""" + return next( + (s for s in range(MAX_AB_STAGES, 0, -1) if smem_bytes(tactic, s) <= smem_capacity), 0 + ) + + +def validate( + tactic: BatchSplitKTactic, m: int, n: int, k: int, smem_capacity: int = SMEM_CAPACITY_BYTES +) -> None: + """Reject a tactic that cannot serve public ``x[m, k] @ weight[n, k]^T``. + + Raises: + ValueError: naming the first violated bound. + """ + if tactic.mma_m not in SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + step = MMA_N_STEP[tactic.mma_m] + if not step <= tactic.mma_n <= MAX_MMA_N or tactic.mma_n % step: + raise ValueError( + f"mma_n={tactic.mma_n} must be a multiple of {step} in [{step}, {MAX_MMA_N}] at " + f"mma_m={tactic.mma_m}" + ) + if values_per_thread(tactic) > MAX_VALUES_PER_THREAD: + raise ValueError( + f"tile {tactic.mma_m}x{tactic.mma_n} gives {values_per_thread(tactic)} FP32 values per " + f"epilogue thread (at most {MAX_VALUES_PER_THREAD})" + ) + if tactic.split_k not in SUPPORTED_SPLIT_K: + raise ValueError(f"unsupported split_k={tactic.split_k} (one of {SUPPORTED_SPLIT_K})") + if not 1 <= tactic.ab_stages <= MAX_AB_STAGES: + raise ValueError(f"ab_stages must be in [1, {MAX_AB_STAGES}], got {tactic.ab_stages}") + if not 1 <= m <= MAX_ROWS: + raise ValueError(f"this kernel serves 1 <= M <= {MAX_ROWS}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % (CTA_K * tactic.split_k): + raise ValueError( + f"K={k} does not split into {tactic.split_k} slices of {CTA_K}-wide k-tiles" + ) + if tactic.ab_stages > k // CTA_K // tactic.split_k: + raise ValueError( + f"ab_stages={tactic.ab_stages} exceeds the {k // CTA_K // tactic.split_k} k-tiles of a " + "CTA's K slice" + ) + need = smem_bytes(tactic) + if need > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {need} B of shared memory, {smem_capacity} B available; at most " + f"{max_stages(tactic, smem_capacity)} stages fit" + ) + + +@dsl_user_op +def _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, *, loc=None, ip=None): + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [smem_ptr.toint(loc=loc, ip=ip).ir_value(), peer_cta_rank_in_cluster.ir_value()], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr, + mbar_ptr, + peer_cta_rank_in_cluster, + *, + loc=None, + ip=None, +): + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes to its mbarrier.""" + llvm.inline_asm( + None, + [ + _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _mbarrier_wait_acquire_cluster(mbar_ptr, phase, *, loc=None, ip=None) -> None: + """``mbarrier.try_wait.parity`` with ``.acquire.cluster``: the peers' ``st.async`` complete-tx + operations release at cluster scope (flashinfer #7443).""" + llvm.inline_asm( + None, + [mbar_ptr.toint(loc=loc, ip=ip).ir_value(), phase.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred P1;\n\tLAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [$0], $1, 10000000;\n\t" + "@P1 bra DONE;\n\tbra LAB_WAIT;\n\tDONE:\n\t}", + "r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _tcgen05_fence(*, before: bool, loc=None, ip=None) -> None: + """``tcgen05.fence::{before,after}_thread_sync`` around an mbarrier hand-off of tcgen05 work.""" + llvm.inline_asm( + None, + [], + f"tcgen05.fence::{'before' if before else 'after'}_thread_sync;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +class BatchSplitKGemmKernel: + """BF16 GEMM with a cluster-local reduce-scatter split-K reduction (see the module docstring).""" + + def __init__(self, *, tactic: BatchSplitKTactic, use_pdl: bool) -> None: + self.acc_dtype = cutlass.Float32 + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.threads_per_cta = 256 + self.epilog_threads = _EPILOG_THREADS + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + vpt = values_per_thread(tactic) + if vpt % 4: + raise ValueError( + f"tile {tactic.mma_m}x{tactic.mma_n}: {vpt} values per thread, not a multiple of 4" + ) + # Reduce-scatter: epilogue thread group g (rs_group threads) of every CTA belongs to cluster + # rank g, which receives that group's values from the other split_k - 1 ranks. + self.rs_group = self.epilog_threads // tactic.split_k + self.rs_mailbox_elements = max(1, (tactic.split_k - 1) * self.rs_group * vpt) + self.rs_expected_bytes = (tactic.split_k - 1) * self.rs_group * vpt * _FP32_BYTES + + @cute.experimental.jit + def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: _cuda.CUstream): + # Grid-y packs the kernel-N (public row) tile and the cluster rank. + self.kernel(a, b, c).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(SMEM_CAPACITY_BYTES), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major: the weight + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major: the activation + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major: the output + ): + stages = self.num_ab_stage + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + bar_spare = cute_ext.allocate( + cutlass.Int64, cute.AddressSpace.smem, cute.make_layout(1), alignment=_MBARRIER_BYTES + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.rs_mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_spare + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + if cutlass.const_expr(self.split_k > 1): + # The receiving group's arrival plus the peers' transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + # Publish every CTA's barriers before any cross-CTA store. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 0 loads the weight, warp 1 the activation, warp 2 issues the MMAs, warp 3 idles, + # warps 4-7 run the epilogue. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + ) + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + # The activation is the predecessor's output: load it only after the dependency wait. + cute.arch.griddepcontrol_wait() + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + bar_full + stage, + cute.size_in_bytes( + s_tile.element_type, cute.slice_(s_tile.layout, (None, None, None, 0)) + ), + ) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + if cutlass.const_expr(is_a and self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def _drain_producer(self, bar_empty, empty_phase: cutlass.Int32, k_tile_count: cutlass.Int32): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, + sB: cute.Tensor, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): + cute.arch.alloc_tmem(_TMEM_COLUMNS, tmem_base_ptr, is_two_cta=False) + _tcgen05_fence(before=True) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + _tcgen05_fence(before=False) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + # The release arrival must land in phase 1 (flashinfer #7443). + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + _tcgen05_fence(before=False) + cute.arch.dealloc_tmem(tmem_ptr, _TMEM_COLUMNS, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + ): + # Wait until the MMA warp publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + _tcgen05_fence(before=False) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), acc_layout + )[((None, None), 0, 0, 0)] + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate(self.acc_dtype, cute.AddressSpace.rmem, rmem_layout, alignment=32) + rD = cute_ext.allocate(c_dtype, cute.AddressSpace.rmem, rmem_layout, alignment=32) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + _tcgen05_fence(before=False) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before the TMEM release and the RMEM use. + cute.arch.fence_view_async_tmem_load() + _tcgen05_fence(before=True) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + if cutlass.const_expr(self.split_k > 1): + # Reduce-scatter: thread group g = epi_tid // rs_group of every CTA sends its FP32 + # partial to cluster rank g (unless it is rank g); rank g's group g sums the split_k + # partials of its values in cluster-rank order, rounds once and stores them. + vpt = cutlass.const_expr(cute.size(rmem_layout)) + group_values = cutlass.const_expr(self.rs_group * vpt) + dest = epi_tid // self.rs_group + lt = epi_tid % self.rs_group + if dest != split_rank: + # This sender's slot in rank dest's mailbox: ranks below dest keep their index, + # ranks above shift down by one. + j = split_rank - cutlass.select_(split_rank > dest, Int32(1), Int32(0)) + for value_idx in cutlass.range_constexpr(0, vpt, 4): + slot = (value_idx // 4) * (self.rs_group * 4) + lt * 4 + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + j * group_values + slot, + bar_reduce, + dest, + ) + else: + if lt == 0: + cute.arch.mbarrier_arrive_and_expect_tx(bar_reduce, self.rs_expected_bytes) + _mbarrier_wait_acquire_cluster(bar_reduce, Int32(0)) + for value_idx in cutlass.range_constexpr(vpt): + slot = (value_idx // 4) * (self.rs_group * 4) + lt * 4 + value_idx % 4 + own = rAcc[value_idx] + acc = own + for r in cutlass.range_constexpr(self.split_k): + # Rank r's partial: this CTA's own, or its slot in the mailbox (the index + # clamped where it is unused). + jr = cutlass.min( + Int32(r) - cutlass.select_(Int32(r) > split_rank, Int32(1), Int32(0)), + Int32(self.split_k - 2), + ) + term = cutlass.select_( + Int32(r) == split_rank, own, mailbox[jr * group_values + slot] + ) + if cutlass.const_expr(r == 0): + acc = term + else: + acc = acc + term + rAcc[value_idx] = acc + rD.store(rAcc.load().to(c_dtype)) + # Preserve the TMEM coordinates; the copy predicates the output's tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + else: + rD.store(rAcc.load().to(c_dtype)) + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + # The reduction mbarrier covers the remote stores; no cluster barrier is needed. + + +@cute.experimental.jit +def _bmm( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + stream, + ) + + +def _from_dlpack_dynamic(tensor: torch.Tensor, leading_dim: int, assumed_align: int = 32): + return from_dlpack(tensor.detach(), assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: torch.Tensor) -> int: + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor(shape, dtype, leading_dim: int) -> torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] + return torch.empty(tuple(shape[d] for d in permutation), dtype=dtype, device="cuda").permute( + [permutation.index(d) for d in range(len(shape))] + ) + + +def _swapped(x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor): + """The kernel's operands: A = the weight [1, N, K], B = x as [1, K, M], C = out as [1, N, M].""" + tensors = ( + weight.unsqueeze(0), + x.unsqueeze(0).transpose(-2, -1), + out.unsqueeze(0).transpose(-2, -1), + ) + leading = tuple(_detect_leading_dim(t) for t in tensors) + return tuple(_from_dlpack_dynamic(t, d) for t, d in zip(tensors, leading)), leading + + +@functools.lru_cache(maxsize=None) +def _compiled( + dtype, tactic: BatchSplitKTactic, use_pdl: bool, leading: Tuple[int, int, int], device: int +): + # Compile-time stand-ins (flashinfer's): every extent and stride but the unit one is dynamic. + m, n, k = 64, 8, CTA_K + with torch.cuda.device(device): + reprs = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, d), d) + for shape, d in ( + ((1, n, k), leading[0]), + ((1, k, m), leading[1]), + ((1, n, m), leading[2]), + ) + ) + return cute_ext.compile( + _bmm, + BatchSplitKGemmKernel(tactic=tactic, use_pdl=use_pdl), + *reprs, + _cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +def run( + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + tactic: BatchSplitKTactic, + pdl: bool = True, +) -> torch.Tensor: + """``out[M, N] = x[M, K] @ weight[N, K]^T`` in BF16: row-major, contiguous, 32-byte aligned. + + Raises: + ValueError: for operands or a tactic the kernel does not take. + """ + if x.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16 or out.dtype != torch.bfloat16: + raise ValueError("batch split-K takes BF16 operands only") + if x.dim() != 2 or weight.dim() != 2 or out.dim() != 2: + raise ValueError("batch split-K takes 2-D x [M, K], weight [N, K], out [M, N]") + m, k = (int(d) for d in x.shape) + n = int(weight.shape[0]) + if int(weight.shape[1]) != k or tuple(out.shape) != (m, n): + raise ValueError( + f"shapes x {tuple(x.shape)}, weight {tuple(weight.shape)}, out {tuple(out.shape)}" + ) + if not (x.is_contiguous() and weight.is_contiguous() and out.is_contiguous()): + raise ValueError("batch split-K takes contiguous row-major operands") + if any(t.data_ptr() % 32 for t in (x, weight, out)) or not ( + x.is_cuda and weight.device == x.device == out.device + ): + raise ValueError("batch split-K takes 32-byte aligned CUDA operands on one device") + validate(tactic, m, n, k) + with torch.cuda.device(x.device): + tensors, leading = _swapped(x, weight, out) + compiled = _compiled(x.dtype, tactic, bool(pdl), leading, x.get_device()) + compiled(*tensors, _cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream)) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/_persistent_experts/kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/_persistent_experts/kernel.py index 6ae9980e4690..a118c8730569 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/_persistent_experts/kernel.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/_persistent_experts/kernel.py @@ -7,6 +7,22 @@ scale with it (``geometry``); everything below is written for the ``ntok = 8`` build, whose constants are the module's ``NTOK`` ones. +Batch builds (``overlap``; ``ntok`` 16 or 32, windows of at most 64 experts): a CTA +runs all its FC1 tiles before its first FC2 tile, so once the MMAs of its last FC1 tile +have read activation k-block ``kb`` that B stage holds FC2 operand k-block ``kb`` +instead. The FC2 image is one buffer over the activation's 16 B stages (and their SFB +stages), each k-block with its own empty barrier, which the MMA warp commits after the +CTA's last FC1 tile and after each FC2 tile; the fetch warps write k-block ``kb`` of FC2 +tile ``j2`` once commit ``j2 + had_fc1`` there has completed (the stage's previous +reader; the parity wait is exact because the commit before it was waited for by the +previous tile's copy). The B operand then costs 16 stages, not 48, which leaves the +weight ring its six stages at 32 token columns (the 16-column wide build keeps four). +The activation is loaded by the epilogue warps as well (idle until the first +accumulator) and lands in groups of ``X_GROUP_KB`` k-blocks, each with its own barrier, +so the first FC1 MMAs wait for the first group only. Everything else -- the tile list, +the arithmetic, the tagged hand-off, the row map -- is the ``ntok = 8`` kernel's, so a +routed row is bitwise the same in every build. + One launch per MoE layer replaces TRTLLM-gen's routing-indices kernel, FC1 and FC2 for the rows of one expert-parallel window of ``L`` local experts (32 = this rank's own block, 64 = the ring-buddy stack, 96 = the @@ -59,6 +75,30 @@ the activation operand, then fetch the FC2 operand images (k-blocks 0-7 and 8-15). +Interleaved fetch (``fetch_kb != 0``, a batch build with the tagged hand-off only; compiled +out otherwise). The two fetch warps take the FC2 k-blocks alternately (warp ``h`` serves +``h, h + 2, ..., h + 14``, in the MMA's order) instead of halves, and a 32-column build polls +``fetch_kb`` k-blocks per pass (one, as without it, at ``fetch_kb = 1``). The k-blocks an FC2 +tile waits for after the last FC1 round (its last ones) are then polled by both warps instead +of one warp in series, and both warps deliver in the order the MMA consumes. The same words +land in the same shared-memory bytes behind the same barriers: only which warp moves a +k-block, and when, changes. + +In-kernel placement (``place_rank >= 0``, a batch build on the 64-expert ring-buddy window; +compiled out otherwise). ``ids`` are the step's selected global expert ids (0..255) instead of +``ring_buddy_assign``'s window ids, and every CTA derives the window ids itself: the balanced +ring placement of ``ops/moe/ring_buddy_assign`` (8 ranks of 32-expert blocks, ``balance`` on) +for rank ``place_rank``, statement for statement in integers (``_place_census``, +``_place_flow``, ``_place_local``), so the expert call needs no placement launch before it. +Ids outside 0..255 are treated as not local. The scratch (census, per-block maxima, moves) is +the FC1 staging buffer, which nothing reads or writes before the first accumulator. + +Early start (``early_start != 0``, a batch build only; compiled out otherwise). The +kernel's same-stream predecessor (``ops/moe/ring_buddy_assign`` with ``early_trigger``) +triggers its dependents right after its own dependency wait, so this grid launches, +allocates and initializes while the predecessor computes the ids, and reads its epoch +(its own workspace word) before its wait; everything else is read after it. + Early ids (``early_expert_ids``, ``early_ids != 0``; compiled out otherwise). The kernel's same-stream predecessor (``ops/comm/pool_gather_route``) publishes the window ids as 64 single-copy-atomic 64-bit words ``(id, tag)`` at the @@ -136,9 +176,22 @@ #: Workspace bytes per slot, enough for either hand-off. WS_SLOT_BYTES = max(IMG_BYTES, TAG_SLOT_BYTES) +# mbarriers (besides the ring and accumulator pipelines) +X_BAR = 0 +#: A batch build's activation lands in groups of this many k-blocks, each with its own +#: barrier (``X_BAR + g``), so the first FC1 MMAs need only the first group. +X_GROUP_KB = 4 +IMG_FULL_BAR = 1 # IMG_BUFS x FC2_KB, one per (buffer, k-block) +IMG_EMPTY_BAR = IMG_FULL_BAR + IMG_BUFS * FC2_KB # one per buffer +NUM_BAR = IMG_EMPTY_BAR + IMG_BUFS + #: Token columns a build is compiled for: the 8-row decode step, and the wide build of a #: TP4 decode step of two requests. NTOKS = (8, 16) +#: Token columns of the batch builds (``overlap``): TP8 decode steps of two (16) and of +#: three or four requests (32) at K = 7. 32 token rows fill each 128-row scale atom's +#: 32-row group, so the compact SFB chunk is the whole 512-byte atom there. +BATCH_NTOKS = (16, 32) class Geometry: @@ -148,16 +201,29 @@ class Geometry: ``0 .. ntok - 1`` of each 128-row scale chunk (``16 ntok`` bytes apart); the FC1 tile's staged output is ``ntok`` rows x 32 code bytes then ``ntok`` x 4 scale bytes; a tagged tile is those ``9 ntok`` words. + + ``overlap`` (a batch build): the FC2 operand image reuses the activation's B stages + (16 B stages, not 48) through one image buffer whose k-blocks each have an empty + barrier (``img_empty_bar + kb``); see the module docstring. """ - def __init__(self, ntok: int): - assert ntok in NTOKS, ntok + def __init__(self, ntok: int, overlap: bool = False): + assert ntok in (BATCH_NTOKS if overlap else NTOKS), (ntok, overlap) self.ntok = ntok + self.overlap = bool(overlap) + self.b_stages = FC1_KB if overlap else B_STAGES + self.img_bufs = 1 if overlap else IMG_BUFS + self.img_stage0 = X_STAGE0 if overlap else IMG_STAGE0 + # the activation's barriers: one, or one per X_GROUP_KB k-blocks in a batch build + self.x_bars = FC1_KB // X_GROUP_KB if overlap else 1 + self.img_full_bar = X_BAR + self.x_bars + self.img_empty_bar = self.img_full_bar + self.img_bufs * FC2_KB + self.num_bar = self.img_empty_bar + (FC2_KB if overlap else self.img_bufs) self.route = ntok * TOP_K # (token, k) routing lanes self.b_stage_bytes = ntok * KTILE // 2 self.sfb_chunk = 16 * ntok self.sfb_stage = 4 * self.sfb_chunk - self.sfb_bytes = (B_STAGES - 1) * self.sfb_stage + 3 * self.sfb_chunk + 512 + self.sfb_bytes = (self.b_stages - 1) * self.sfb_stage + 3 * self.sfb_chunk + 512 self.img_b_bytes = FC2_KB * self.b_stage_bytes self.img_sf_bytes = FC2_KB * self.sfb_stage self.img_bytes = self.img_b_bytes + self.img_sf_bytes @@ -186,6 +252,12 @@ def __init__(self, ntok: int): _G8.hstage_sf, _G8.hstage_bytes, _G8.ws_slot_bytes, + _G8.b_stages, + _G8.img_bufs, + _G8.img_stage0, + _G8.img_empty_bar, + _G8.num_bar, + _G8.img_full_bar, ) == ( B_STAGE_BYTES, SFB_CHUNK_STRIDE, @@ -198,6 +270,12 @@ def __init__(self, ntok: int): HSTAGE_SF, HSTAGE_BYTES, WS_SLOT_BYTES, + B_STAGES, + IMG_BUFS, + IMG_STAGE0, + IMG_EMPTY_BAR, + NUM_BAR, + IMG_FULL_BAR, ) # Warp roles @@ -211,12 +289,9 @@ def __init__(self, ntok: int): BAR_EPI = 1 # named barrier of the epilogue warps BAR_X = 2 # named barrier of the activation loaders BAR_ROUTE = 3 # named barrier of every warp but the weight TMA warp (routing) - -# mbarriers (besides the ring and accumulator pipelines) -X_BAR = 0 -IMG_FULL_BAR = 1 # IMG_BUFS x FC2_KB, one per (buffer, k-block) -IMG_EMPTY_BAR = IMG_FULL_BAR + IMG_BUFS * FC2_KB # one per buffer -NUM_BAR = IMG_EMPTY_BAR + IMG_BUFS +#: Threads that load the activation in a batch build (``_load_x_batch``): the epilogue +#: warps and warps 6-7. +X_LOADERS = 32 * (EPI_WARPS + NUM_WARPS - X_WARP0) TMEM_COLS = 128 TMEM_SFA_COL = 64 @@ -234,6 +309,15 @@ def __init__(self, ntok: int): # Workspace words per CTA (launch epochs) MAX_CTAS = 1024 +# In-kernel placement (``place_rank``): ``ring_buddy_assign``'s tep8 ring of 8 ranks of 32-expert +# blocks, and its scratch words in the FC1 staging buffer (``hstage``): the per-block census, +# each block's smallest feasible maximum, and each block's moves to its successor. +RING = 8 +RING_BLOCK = 32 +PLACE_CENSUS = 0 +PLACE_NEED = RING +PLACE_MOVED = 2 * RING + # Stress-test delay sites (``delay`` bits; compiled out when 0). A # pseudo-random third of the CTAs, re-drawn every launch, spins 1..6 units of # ``delay_ns`` on the global timer at each enabled site. @@ -284,16 +368,48 @@ def __init__( early_deadline_ns: int = 100_000, early_poll_ns: int = 32, ntok: int = NTOK, + overlap: int = 0, + early_start: int = 0, + fetch_kb: int = 0, + place_rank: int = -1, ): assert num_ctas <= MAX_CTAS assert num_local in (32, 64, 96) # fastroute's warp-local bit set is two words wide assert num_local <= 64 or fastroute == 0, "fastroute supports windows up to 64" - self.g = Geometry(ntok) + # Early start is built for the batch builds; the early-ids variant has its own entry. + assert early_start in (0, 1), early_start + assert early_start == 0 or (overlap != 0 and early_ids == 0), ( + "the persistent kernel's early start is built for the batch builds without early ids" + ) + # Interleaved fetch: k-blocks per polling pass of a 32-column build (a divisor of the + # eight k-blocks a fetch warp serves), on the batch builds' tagged hand-off. + assert fetch_kb in (0, 1, 2, 4), fetch_kb + assert fetch_kb == 0 or (overlap != 0 and handoff != 0), ( + "the persistent kernel's interleaved fetch is built for the batch builds' tagged " + "hand-off" + ) + # In-kernel placement: the batch builds on the two-window (64-expert) ring-buddy stack, + # with the routing through shared memory, the default entry and plain ids. + assert -1 <= place_rank < RING, place_rank + assert place_rank < 0 or ( + overlap != 0 + and num_local == 2 * RING_BLOCK + and early_ids == 0 + and early_start == 0 + and fastroute == 0 + ), "the persistent kernel's in-kernel placement is built for the batch builds' 64-window" + self.g = Geometry(ntok, overlap != 0) + assert place_rank < 0 or (PLACE_MOVED + RING) * 4 <= self.g.hstage_bytes self.ntok = ntok assert acc_stages * ntok <= TMEM_SFA_COL # A wide build's slot table keeps 64 slots: at most 64 distinct local experts. assert ntok == NTOK or (num_local <= 64 and early_ids == 0) + # A batch build's fetch warps wait on per-k-block empty barriers (``_fetch_images``, + # ``_fetch_tagged``); the fast route (the TMA warp's own routing) is not built for it. + assert overlap == 0 or fastroute == 0, ( + "the persistent kernel's batch builds take the routing through shared memory" + ) if early_ids != 0 and (fastroute != 0 or early_fallbacks == 0): # The fast route reads the ids from global memory in the TMA warp, # which never waits in early mode. @@ -321,6 +437,9 @@ def __init__( self.early_fallbacks = early_fallbacks self.early_deadline_ns = early_deadline_ns self.early_poll_ns = early_poll_ns + self.early_start = early_start + self.fetch_kb = fetch_kb + self.place_rank = place_rank self.sf_vec_size = 16 self.cta_group = tcgen05.CtaGroup.ONE self.mma_inst_mn = (128, ntok) @@ -355,12 +474,12 @@ def _setup(self): self.tiled_mma, self.mma_tiler, self.sf_vec_size, self.num_stages ) self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( - self.tiled_mma, self.mma_tiler, cutlass.Float4E2M1FN, B_STAGES + self.tiled_mma, self.mma_tiler, cutlass.Float4E2M1FN, self.g.b_stages ) # Standard SFB layout (defines the TMEM layout) and the compact SMEM # layout actually used for storage (see SFB_CHUNK_STRIDE and ``Geometry``). self.sfb_smem_layout_std = blockscaled_utils.make_smem_layout_sfb( - self.tiled_mma, self.mma_tiler, self.sf_vec_size, B_STAGES + self.tiled_mma, self.mma_tiler, self.sf_vec_size, self.g.b_stages ) self.sfb_smem_layout_staged = cute.make_layout( self.sfb_smem_layout_std.shape, @@ -435,7 +554,7 @@ def run( class SharedStorage: ring_mbar: cute.struct.MemRange[cutlass.Int64, 2 * self.num_stages] acc_mbar: cute.struct.MemRange[cutlass.Int64, 2 * self.acc_stages] - b_mbar: cute.struct.MemRange[cutlass.Int64, NUM_BAR] + b_mbar: cute.struct.MemRange[cutlass.Int64, self.g.num_bar] tmem_holding: cutlass.Int32 hstage: cute.struct.Align[ cute.struct.MemRange[cutlass.Int64, self.g.hstage_bytes // 8], 16 @@ -578,7 +697,7 @@ def kernel( b_mbar = storage.b_mbar.data_ptr() if warp_idx == MMA_WARP: with cute.arch.elect_one(): - for i in cutlass.range_constexpr(NUM_BAR): + for i in cutlass.range_constexpr(self.g.num_bar): cute.arch.mbarrier_init(b_mbar + i, 1) cute.arch.mbarrier_init_fence() if warp_idx == 0: @@ -601,51 +720,84 @@ def kernel( s_slot_k[tidx + i * NUM_THREADS] = Int32(TOP_K) # above every slot index (atomic min) if tidx < self.bitmap_words: s_bitmap[tidx] = Uint32(0) + if cutlass.const_expr(self.place_rank >= 0): + if tidx < RING: + ptx.sts_u32( + Int32(storage.hstage.data_ptr().toint()) + (PLACE_CENSUS + tidx) * 4, Uint32(0) + ) cute.arch.sync_threads() if tidx == 0: self._stamp(trace_addr, cta, TR_ENTRY) - # Everything below reads tensors written by earlier kernels or touches - # the workspace shared with the previous launch. With early ids only the - # epoch (this grid's own workspace word; the previous launch completed - # before the predecessor's wait returned, and the predecessor triggers - # after that wait) and the tagged ids are read before the wait below. - if cutlass.const_expr(self.early_ids == 0): - cute.arch.griddepcontrol_wait() - # With ``fastroute`` the weight TMA warp routes on its own - # (``_tma_producer``) and skips the CTA's routing barriers, so its first - # loads go out one ids load after the wait. - epoch = Uint32(0) - if cutlass.const_expr(self.fastroute == 0): + if cutlass.const_expr(self.early_start != 0): + # Early start: the predecessor (``ring_buddy_assign``) released this grid right + # after its own dependency wait, so every grid before it has completed, the + # previous launch on this workspace included: the epoch word (this grid's own) + # is final and is read before the wait. The ids the predecessor writes, the + # activation and everything else are read after it. epoch = self._next_epoch(epoch_addr, cta) - else: - if warp_idx != TMA_WARP: - epoch = self._next_epoch(epoch_addr, cta) - if cutlass.const_expr(self.early_ids == 0): + cute.arch.griddepcontrol_wait() if tidx == 0: self._stamp(trace_addr, cta, TR_WAITED) + else: + # Everything below reads tensors written by earlier kernels or touches + # the workspace shared with the previous launch. With early ids only the + # epoch (this grid's own workspace word; the previous launch completed + # before the predecessor's wait returned, and the predecessor triggers + # after that wait) and the tagged ids are read before the wait below. + if cutlass.const_expr(self.early_ids == 0): + cute.arch.griddepcontrol_wait() + # With ``fastroute`` the weight TMA warp routes on its own + # (``_tma_producer``) and skips the CTA's routing barriers, so its first + # loads go out one ids load after the wait. + epoch = Uint32(0) + if cutlass.const_expr(self.fastroute == 0): + epoch = self._next_epoch(epoch_addr, cta) + else: + if warp_idx != TMA_WARP: + epoch = self._next_epoch(epoch_addr, cta) + if cutlass.const_expr(self.early_ids == 0): + if tidx == 0: + self._stamp(trace_addr, cta, TR_WAITED) if cutlass.const_expr(self.delay & DELAY_ENTRY): self._delay_spin(DELAY_ENTRY * 53, cta) # -------- window ids -> sorted slots of the active local experts -------- # Warps 0-1, one lane per (token, slot). - if tidx < self.g.route: - t = tidx // TOP_K - gid = Int32(-1) - if t < num_tokens: - if cutlass.const_expr(self.early_ids != 0): - gid = self._early_id(tidx, ids_addr, epoch, trace_addr, cta) - else: - gid = Int32(ptx.ld_u32(ids_addr + tidx * 4)) - local = (gid >= 0) & (gid < self.num_local) - lid = Int32(-1) - if local: - lid = gid - ptx.atom_or_shared_u32( - Int32(storage.bitmap.data_ptr().toint()) + (gid // 32) * 4, - Uint32(1) << Uint32(gid % 32), - ) - s_ids[tidx] = lid + if cutlass.const_expr(self.place_rank >= 0): + # In-kernel placement: the step's census, then the ring flow, then this lane's id. + hst = Int32(storage.hstage.data_ptr().toint()) + gid = self._place_census(tidx, ids_addr, num_tokens, hst) + cute.arch.sync_threads() + if warp_idx == 0: + self._place_flow(lane, hst) + cute.arch.sync_threads() + if tidx < self.g.route: + lid = self._place_local(gid, hst) + if lid >= 0: + ptx.atom_or_shared_u32( + Int32(storage.bitmap.data_ptr().toint()) + (lid // 32) * 4, + Uint32(1) << Uint32(lid % 32), + ) + s_ids[tidx] = lid + else: + if tidx < self.g.route: + t = tidx // TOP_K + gid = Int32(-1) + if t < num_tokens: + if cutlass.const_expr(self.early_ids != 0): + gid = self._early_id(tidx, ids_addr, epoch, trace_addr, cta) + else: + gid = Int32(ptx.ld_u32(ids_addr + tidx * 4)) + local = (gid >= 0) & (gid < self.num_local) + lid = Int32(-1) + if local: + lid = gid + ptx.atom_or_shared_u32( + Int32(storage.bitmap.data_ptr().toint()) + (gid // 32) * 4, + Uint32(1) << Uint32(gid % 32), + ) + s_ids[tidx] = lid self._route_barrier(warp_idx) # Every thread has read its epoch word: advance it for the next launch. if tidx == 0: @@ -750,7 +902,18 @@ def kernel( trace_addr, ) if warp_idx >= X_WARP0: - self._load_x(tidx - 32 * X_WARP0, x_addr, xsf_addr, num_tokens, storage, b_mbar) + if cutlass.const_expr(self.g.overlap): + # batch build: the epilogue warps load the activation too (threads 0..127) + self._load_x_batch( + tidx - 32 * X_WARP0 + 32 * EPI_WARPS, + x_addr, + xsf_addr, + num_tokens, + storage, + b_mbar, + ) + else: + self._load_x(tidx - 32 * X_WARP0, x_addr, xsf_addr, num_tokens, storage, b_mbar) if tidx == 32 * X_WARP0: self._stamp(trace_addr, cta, TR_X) if warp_idx >= FETCH_WARP0: @@ -786,6 +949,8 @@ def kernel( trace_addr, ) if warp_idx < EPI_WARPS: + if cutlass.const_expr(self.g.overlap): + self._load_x_batch(tidx, x_addr, xsf_addr, num_tokens, storage, b_mbar) # off the critical path and the tail: the epilogue warps are idle # until the first accumulator if cta == 0: @@ -878,17 +1043,24 @@ def _write_row_map(self, tidx, index_addr, num_tokens, s_ids, s_expert_slot, s_s its token naming that expert, i.e. the ``k`` the slot table kept, so every valid entry names a row the FC2 epilogue writes. Runs after the routing barriers; the finalize reads the map only after this grid has - completed. + completed. The 128 epilogue threads take the ``8 ntok`` lanes in passes + (one up to 16 token columns, two at 32). """ - if tidx < self.g.route: - t = tidx // TOP_K - if t < num_tokens: - out = Int32(-1) - lid = s_ids[tidx] - if lid >= 0: - if s_slot_k[s_expert_slot[lid] * self.ntok + t] == tidx % TOP_K: - out = tidx - ptx.st_u32(index_addr + tidx * 4, Uint32(out)) + for part in cutlass.range_constexpr( + (self.g.route + 32 * EPI_WARPS - 1) // (32 * EPI_WARPS) + ): + lane = tidx + if cutlass.const_expr(part > 0): + lane = tidx + part * 32 * EPI_WARPS + if lane < self.g.route: + t = lane // TOP_K + if t < num_tokens: + out = Int32(-1) + lid = s_ids[lane] + if lid >= 0: + if s_slot_k[s_expert_slot[lid] * self.ntok + t] == lane % TOP_K: + out = lane + ptx.st_u32(index_addr + lane * 4, Uint32(out)) @cute.jit def _retag_rotation_slot(self, tidx, flag_addr, img_addr, epoch, cta, grid0, num_slots): @@ -992,6 +1164,75 @@ def _route_barrier(self, warp_idx): else: cute.arch.sync_threads() + @cute.jit + def _place_census(self, tidx, ids_addr, num_tokens, hst): + """In-kernel placement, step 1: this lane's selected global id (-1 past the last row or + outside 0..255), OR-ed into its block's presence word (``ring_buddy_assign``'s census: + a repeated expert sets the same bit).""" + gid = Int32(-1) + if tidx < self.g.route: + if tidx // TOP_K < num_tokens: + g = Int32(ptx.ld_u32(ids_addr + tidx * 4)) + if (g >= 0) & (g < RING * RING_BLOCK): + gid = g + ptx.atom_or_shared_u32( + hst + (PLACE_CENSUS + g // RING_BLOCK) * 4, + Uint32(1) << Uint32(g % RING_BLOCK), + ) + return gid + + @cute.jit + def _place_flow(self, lane, hst): + """In-kernel placement, step 2 (warp 0; lane ``r < 8`` is block ``r``): the smallest + feasible maximum ``T`` and the smallest flow ``x[r]`` (experts block ``r`` moves to rank + ``r + 1``), ``ring_buddy_assign``'s balanced placement in integers: ``arc`` is the load of + the ``n`` blocks ending at ``r``; ``T`` is the largest ``ceil(arc / (n + 1))`` for + ``n < 8`` and ``ceil(D / 8)`` for the whole ring; ``x[r] = max(0, max_n(arc - n T))``.""" + if lane < RING: + arc = Int32(0) + need = Int32(0) + for n in cutlass.range_constexpr(1, RING + 1): + b = (lane + RING + 1 - n) % RING + arc = arc + ptx.popc_u32(ptx.lds_u32(hst + (PLACE_CENSUS + b) * 4)) + if cutlass.const_expr(n == RING): + need = cutlass.max(need, (arc + RING - 1) // RING) + else: + need = cutlass.max(need, (arc + n) // (n + 1)) + ptx.sts_u32(hst + (PLACE_NEED + lane) * 4, Uint32(need)) + cute.arch.sync_warp() + if lane < RING: + t = Int32(0) + for r in cutlass.range_constexpr(RING): + t = cutlass.max(t, Int32(ptx.lds_u32(hst + (PLACE_NEED + r) * 4))) + arc = Int32(0) + moved = Int32(0) + for n in cutlass.range_constexpr(1, RING + 1): + b = (lane + RING + 1 - n) % RING + arc = arc + ptx.popc_u32(ptx.lds_u32(hst + (PLACE_CENSUS + b) * 4)) + moved = cutlass.max(moved, arc - n * t) + ptx.sts_u32(hst + (PLACE_MOVED + lane) * 4, Uint32(moved)) + + @cute.jit + def _place_local(self, gid, hst): + """In-kernel placement, step 3: the window id of a lane's selected expert ``gid`` on + rank ``place_rank`` (own block ``i``, the predecessor's ``32 + i``), or -1 when another + rank runs it: the expert moves when fewer than ``x[block]`` present experts of its block + sit above it, as ``ring_buddy_assign`` decides.""" + lid = Int32(-1) + if gid >= 0: + blk = gid // RING_BLOCK + i = gid % RING_BLOCK + word = ptx.lds_u32(hst + (PLACE_CENSUS + blk) * 4) + above = ptx.popc_u32((word >> Uint32(i)) >> Uint32(1)) + runs_on = blk + if above < Int32(ptx.lds_u32(hst + (PLACE_MOVED + blk) * 4)): + runs_on = (blk + 1) % RING + if runs_on == self.place_rank: + lid = i + if blk != self.place_rank: + lid = RING_BLOCK + i + return lid + @cute.jit def _nth_expert(self, lo, hi, slot): """The local expert of ``slot``: the ``slot``-th set bit of the 64-bit set.""" @@ -1030,6 +1271,16 @@ def _early_slots(self, num_slots, grid): early = cap return early + @cute.jit + def _fc1_end(self, num_slots, total): + """One past the last FC1 position a CTA with tile-list length ``total`` runs: + ``min(num_slots * FC1_MT, total)`` (positions ``p`` with ``p + grid`` at or past it are + the CTA's last FC1 tile and its FC2 tiles).""" + fc1_end = num_slots * FC1_MT + if fc1_end > total: + fc1_end = total + return fc1_end + @cute.jit def _decode(self, p, num_slots, n_early): """Tile position -> (is FC2, slot, 128-row tile index within the slot).""" @@ -1267,25 +1518,36 @@ def _mma_issuer( ring.consumer_wait(ring_state) if lane == 0: self._stamp(trace_addr, cta, TR_W0) - # the activation operand (FC1 B) is in shared memory - cute.arch.mbarrier_wait(b_mbar + X_BAR, 0) + # the activation operand (FC1 B) is in shared memory (a batch build waits per + # group of k-blocks, below) + if cutlass.const_expr(not self.g.overlap): + cute.arch.mbarrier_wait(b_mbar + X_BAR, 0) + if cutlass.const_expr(self.g.overlap): + fc1_end = self._fc1_end(num_slots, total) j2 = Int32(0) # FC2 tiles of this CTA so far p = cta while p < total: is_fc2, slot, mt = self._decode(p, num_slots, n_early) - buf = j2 % IMG_BUFS - use_phase = (j2 // IMG_BUFS) % 2 + buf = j2 % self.g.img_bufs + use_phase = (j2 // self.g.img_bufs) % 2 b0 = Int32(X_STAGE0) if is_fc2 != 0: - b0 = Int32(IMG_STAGE0) + buf * FC2_KB + b0 = Int32(self.g.img_stage0) + buf * FC2_KB acc_pipe.producer_acquire(acc_state) tCtAcc = tCtAcc_base[(None, None, None, acc_state.index)] tiled_mma.set(tcgen05.Field.ACCUMULATE, False) kb = Int32(0) while kb < FC1_KB: + if cutlass.const_expr(self.g.overlap): + if is_fc2 == 0: + if kb % X_GROUP_KB == 0: + # this group of activation k-blocks has landed + cute.arch.mbarrier_wait(b_mbar + X_BAR + kb // X_GROUP_KB, 0) if is_fc2 != 0: # this k-block's intermediate columns have landed - cute.arch.mbarrier_wait(b_mbar + IMG_FULL_BAR + buf * FC2_KB + kb, use_phase) + cute.arch.mbarrier_wait( + b_mbar + self.g.img_full_bar + buf * FC2_KB + kb, use_phase + ) ring.consumer_wait(ring_state) if cutlass.const_expr(self.trace != 0): if (kb == 0) & (lane == 0): @@ -1308,15 +1570,24 @@ def _mma_issuer( ) tiled_mma.set(tcgen05.Field.ACCUMULATE, True) ring.consumer_release(ring_state) + if cutlass.const_expr(self.g.overlap): + # batch build: B stage kb (activation, then image) is free once + # this tile's MMAs on it have completed; committed after the last + # FC1 tile (which completes after every earlier one) and after + # every FC2 tile, so the fetch waits on one known commit each + if p + grid >= fc1_end: + with cute.arch.elect_one(): + tcgen05.commit(b_mbar + self.g.img_empty_bar + kb) ring_state.advance() kb = kb + 1 if cutlass.const_expr(self.trace != 0): if lane == 0: self._tile_stamp(trace_addr, cta, (p - cta) // grid, 2) if is_fc2 != 0: - # the image buffer is free once these MMAs have completed - with cute.arch.elect_one(): - tcgen05.commit(b_mbar + IMG_EMPTY_BAR + buf) + if cutlass.const_expr(not self.g.overlap): + # the image buffer is free once these MMAs have completed + with cute.arch.elect_one(): + tcgen05.commit(b_mbar + self.g.img_empty_bar + buf) j2 = j2 + 1 acc_pipe.producer_commit(acc_state) acc_state.advance() @@ -1394,6 +1665,97 @@ def _load_x(self, q, x_addr, xsf_addr, num_tokens, storage, b_mbar): if q == 0: cute.arch.mbarrier_arrive(b_mbar + X_BAR) + @cute.jit + def _load_x_batch(self, q, x_addr, xsf_addr, num_tokens, storage, b_mbar): + """Batch builds: ``_load_x`` spread over the epilogue warps and warps 6-7, in groups. + + Thread ``q`` (``0 .. X_LOADERS - 1``: the epilogue warps, then warps 6-7) loads, + group by group (k-blocks ``4 g .. 4 g + 3``), the 16-byte code chunks + ``q, q + X_LOADERS, ...`` of the group (k-block-major: k-block, row, 16-byte column) + and the group's scale words (row, then its 16 words). Every load goes out before + the first store, in group order, so the first group's bytes come back first (every + CTA reads the same activation lines). Then, group by group, the thread stores its + chunks and words; after a named barrier over the loaders one thread arrives on + ``X_BAR + g``, and the MMA starts on the first group while the others land. Same + placement as ``_load_x``, with zero codes and scales for rows at or beyond + ``num_tokens``; a thread's slot past the group's last chunk loads a valid word and + stores nothing. + """ + sb = Int32(storage.sB.data_ptr().toint()) + ssfb = Int32(storage.sSFB.data_ptr().toint()) + last = num_tokens - 1 + kb_chunks = self.ntok * (KTILE // 2 // 16) # 16-byte chunks per k-block + g_chunks = X_GROUP_KB * kb_chunks # code chunks per group + g_words = X_GROUP_KB * (KTILE // 64) # scale words per row per group + n_words = self.ntok * g_words # scale words per group + c_iters = (g_chunks + X_LOADERS - 1) // X_LOADERS + w_iters = (n_words + X_LOADERS - 1) // X_LOADERS + codes = [] + scales = [] + for g in cutlass.range_constexpr(self.g.x_bars): + for i in cutlass.range_constexpr(c_iters): + c = q + X_LOADERS * i + kb = g * X_GROUP_KB + c // kb_chunks + n = (c % kb_chunks) // 8 + c16 = c % 8 + nn = n + kk = kb + if (n > last) | (c >= g_chunks): + nn = last + if c >= g_chunks: + kk = g * X_GROUP_KB + codes.append(ptx.ld_v4_u32(x_addr + nn * X_ROW_BYTES + (kk * 8 + c16) * 16)) + for i in cutlass.range_constexpr(w_iters): + w = q + X_LOADERS * i + n = w // g_words + j = g * g_words + w % g_words + nn = n + if n > last: + nn = last + scales.append(ptx.ld_u32(xsf_addr + nn * X_SF_ROW_BYTES + j * 4)) + for g in cutlass.range_constexpr(self.g.x_bars): + for i in cutlass.range_constexpr(c_iters): + c = q + X_LOADERS * i + if c < g_chunks: + kb = g * X_GROUP_KB + c // kb_chunks + n = (c % kb_chunks) // 8 + c16 = c % 8 + d = codes[g * c_iters + i] + d0 = d[0] + d1 = d[1] + d2 = d[2] + d3 = d[3] + if n > last: + d0 = Uint32(0) + d1 = Uint32(0) + d2 = Uint32(0) + d3 = Uint32(0) + ptx.sts_v4_u32( + sb + + (X_STAGE0 + kb) * self.g.b_stage_bytes + + n * 128 + + ((c16 ^ self._sw(n)) * 16), + d0, + d1, + d2, + d3, + ) + for i in cutlass.range_constexpr(w_iters): + w = q + X_LOADERS * i + if w < n_words: + n = w // g_words + j = g * g_words + w % g_words + v = scales[g * w_iters + i] + if n > last: + v = Uint32(0) + ptx.sts_u32( + ssfb + X_STAGE0 * self.g.sfb_stage + j * self.g.sfb_chunk + n * 16, v + ) + ptx.fence_proxy_async_shared() + ptx.bar_sync(BAR_X, X_LOADERS) + if q == 0: + cute.arch.mbarrier_arrive(b_mbar + X_BAR + g) + # ------------------------------------------------------------------ @cute.jit def _fetch_images( @@ -1425,20 +1787,34 @@ def _fetch_images( k-block is requested as soon as its own producers have published, independent of the others. A buffer is refilled only after the MMAs of its previous tile have completed. + + Batch builds (``overlap``): one buffer over the activation's B stages. The MMA warp + commits a stage's empty barrier after the MMAs on it of this CTA's last FC1 tile + (if it has one: ``had_fc1``) and of each FC2 tile, so the lane that copies + k-block ``kb`` of FC2 tile ``j2`` waits for commit number ``j2 + had_fc1`` there, + the stage's previous reader. That parity wait is exact: the commit before it was + waited for by the copy of tile ``j2 - 1`` (or is the first), and the one after + needs this copy. """ sb = Int32(storage.sB.data_ptr().toint()) ssfb = Int32(storage.sSFB.data_ptr().toint()) - full0 = Int32((b_mbar + IMG_FULL_BAR).toint()) + full0 = Int32((b_mbar + self.g.img_full_bar).toint()) kb_base = half * (FC2_KB // 2) + had_fc1 = Int32(0) + if cutlass.const_expr(self.g.overlap): + # this CTA's first position is an FC1 tile + if cta < self._fc1_end(num_slots, total): + had_fc1 = Int32(1) j2 = Int32(0) p = cta while p < total: is_fc2, slot, mt = self._decode(p, num_slots, n_early) if is_fc2 != 0: - buf = j2 % IMG_BUFS - use = j2 // IMG_BUFS - if use > 0: - cute.arch.mbarrier_wait(b_mbar + IMG_EMPTY_BAR + buf, (use - 1) % 2) + buf = j2 % self.g.img_bufs + use = j2 // self.g.img_bufs + if cutlass.const_expr(not self.g.overlap): + if use > 0: + cute.arch.mbarrier_wait(b_mbar + self.g.img_empty_bar + buf, (use - 1) % 2) if cutlass.const_expr(self.delay & DELAY_CONSUMER): self._delay_spin(DELAY_CONSUMER * 53, cta) fa = flag_addr + (slot * FC1_MT + 4 * kb_base + lane) * 4 @@ -1462,12 +1838,20 @@ def _fetch_images( if lane < FC2_KB // 2: if ((new >> Uint32(lane)) & Uint32(1)) != Uint32(0): kb = kb_base + lane + if cutlass.const_expr(self.g.overlap): + # B stage kb is free: the MMAs of this CTA's last + # tile on it (FC1 or FC2) have completed + commits = j2 + had_fc1 + if commits > 0: + cute.arch.mbarrier_wait( + b_mbar + self.g.img_empty_bar + kb, (commits - 1) % 2 + ) # the image was written by generic stores of other # CTAs; the bulk copy reads it through the async proxy ptx.fence_proxy_async_global() bar = full0 + (buf * FC2_KB + kb) * 8 ptx.mbar_arrive_expect_tx(bar, Int32(self.g.img_kb_tx)) - stage = IMG_STAGE0 + buf * FC2_KB + kb + stage = self.g.img_stage0 + buf * FC2_KB + kb src = img_addr + Int64(slot) * self.g.img_bytes ptx.bulk_g2s( sb + stage * self.g.b_stage_bytes, @@ -1521,9 +1905,26 @@ def _fetch_tagged( shared stores, made visible to the tensor core (``fence.proxy.async.shared::cta``, warp barrier) and released to the MMA warp with one arrive on its (buffer, k-block) mbarrier. + + Batch builds (``overlap``): one buffer over the activation's B stages; before + unpacking k-block ``kb`` of FC2 tile ``j2`` the warp waits for commit + ``j2 + had_fc1`` on that stage's empty barrier, as ``_fetch_images`` does, and a + 32-column build polls one k-block per pass (36 words per lane). + + Interleaved fetch (``fetch_kb``): warp ``half`` serves k-blocks ``half, half + 2, ...`` + in passes of ``pass_kb`` (``fetch_kb`` at 32 columns), see the module docstring. """ sb = Int32(storage.sB.data_ptr().toint()) ssfb = Int32(storage.sSFB.data_ptr().toint()) + pass_kb = TAG_PASS_KB + if cutlass.const_expr(self.g.overlap and self.ntok > 16): + pass_kb = 1 + if cutlass.const_expr(self.fetch_kb != 0): + pass_kb = self.fetch_kb + had_fc1 = Int32(0) + if cutlass.const_expr(self.g.overlap): + if cta < self._fc1_end(num_slots, total): + had_fc1 = Int32(1) # this lane's 9 words of a k-block: source offset and destination src_off = [] dst_off = [] @@ -1550,21 +1951,29 @@ def _fetch_tagged( while p < total: is_fc2, slot, mt = self._decode(p, num_slots, n_early) if is_fc2 != 0: - buf = j2 % IMG_BUFS - use = j2 // IMG_BUFS - if use > 0: - cute.arch.mbarrier_wait(b_mbar + IMG_EMPTY_BAR + buf, (use - 1) % 2) + buf = j2 % self.g.img_bufs + use = j2 // self.g.img_bufs + if cutlass.const_expr(not self.g.overlap): + if use > 0: + cute.arch.mbarrier_wait(b_mbar + self.g.img_empty_bar + buf, (use - 1) % 2) if cutlass.const_expr(self.delay & DELAY_CONSUMER): self._delay_spin(DELAY_CONSUMER * 53, cta) slot_base = img_addr + Int64(slot) * self.g.tag_slot_bytes - for ps in cutlass.range_constexpr(FC2_KB // 2 // TAG_PASS_KB): - kb0 = half * (FC2_KB // 2) + ps * TAG_PASS_KB + for ps in cutlass.range_constexpr(FC2_KB // 2 // pass_kb): + kb0 = half * (FC2_KB // 2) + ps * pass_kb + # the pass's k-blocks: a run of the warp's half, or every other one + kbs = [] + for q in cutlass.range_constexpr(pass_kb): + if cutlass.const_expr(self.fetch_kb != 0): + kbs.append(2 * (ps * pass_kb + q) + half) + else: + kbs.append(kb0 + q) delivered = Int32(0) - while delivered != Int32((1 << TAG_PASS_KB) - 1): + while delivered != Int32((1 << pass_kb) - 1): payload = [] valid = [] - for q in cutlass.range_constexpr(TAG_PASS_KB): - src = slot_base + Int64((kb0 + q) * self.g.tag_kb_bytes) + for q in cutlass.range_constexpr(pass_kb): + src = slot_base + Int64(kbs[q] * self.g.tag_kb_bytes) ok = Int32(1) pl = [] for i in cutlass.range_constexpr(self.g.tag_lane_words): @@ -1574,10 +1983,19 @@ def _fetch_tagged( ok = Int32(0) payload.append(pl) valid.append(ok) - for q in cutlass.range_constexpr(TAG_PASS_KB): + for q in cutlass.range_constexpr(pass_kb): if ((delivered >> q) & 1) == 0: if ptx.ballot_u32(valid[q]) == Uint32(0xFFFFFFFF): - stage = IMG_STAGE0 + buf * FC2_KB + kb0 + q + if cutlass.const_expr(self.g.overlap): + # B stage kbs[q] is free: the MMAs of its previous + # reader (the last FC1 tile, or FC2 tile j2 - 1) are done + commits = j2 + had_fc1 + if commits > 0: + cute.arch.mbarrier_wait( + b_mbar + self.g.img_empty_bar + kbs[q], + (commits - 1) % 2, + ) + stage = self.g.img_stage0 + buf * FC2_KB + kbs[q] for i in cutlass.range_constexpr(self.g.tag_lane_words): if is_sf[i] != 0: ptx.sts_u32( @@ -1593,7 +2011,7 @@ def _fetch_tagged( cute.arch.sync_warp() if lane == 0: cute.arch.mbarrier_arrive( - b_mbar + IMG_FULL_BAR + buf * FC2_KB + kb0 + q + b_mbar + self.g.img_full_bar + buf * FC2_KB + kbs[q] ) delivered = delivered | (1 << q) if cutlass.const_expr(self.trace != 0): @@ -1818,23 +2236,78 @@ def _swiglu_quant_to_stage(self, v, warp_idx, lane, mask, s_gate, s_c, num_token rcp6 = Float32(RCP6_F32) for p in cutlass.range_constexpr(self.ntok // 2): n0 = 2 * p - g0 = cute.arch.shuffle_sync_down(v[n0], 8) - g1 = cute.arch.shuffle_sync_down(v[n0 + 1], 8) - x = ptx.mul_rn_x2(g0, g1, t_neg, t_neg) - d = ptx.add_rn_x2(ptx.ex2_approx(x[0]), ptx.ex2_approx(x[1]), one, one) - gs = ptx.mul_rn_x2(g0, g1, ptx.rcp_approx_ftz(d[0]), ptx.rcp_approx_ftz(d[1])) - u = ptx.fma_rn_zero_x2(v[n0], v[n0 + 1], s_gate, s_gate) - hh = ptx.mul_rn_x2(u[0], u[1], gs[0], gs[1]) - h0 = Float32(0.0) - h1 = Float32(0.0) - if is_up: - h0 = hh[0] - h1 = hh[1] - # (amax|h|) * sC == max(|h| * sC): rounding is monotone and sC > 0 - a = ptx.mul_rn_x2(self._block_amax(h0), self._block_amax(h1), s_c, s_c) - q = ptx.mul_rn_x2(a[0], a[1], rcp6, rcp6) - self._quant_store(h0, q[0], s_c, n0, mask, num_tokens, warp_idx, lane, hst) - self._quant_store(h1, q[1], s_c, n0 + 1, mask, num_tokens, warp_idx, lane, hst) + if cutlass.const_expr(self.g.overlap): + # Batch builds: a pair with neither token routed to this expert (``mask`` + # is the same in every thread) gets the zero bytes and scales + # ``_quant_store`` would give it, without the arithmetic. + if ((mask >> Uint32(n0)) & Uint32(3)) != Uint32(0): + self._swiglu_pair( + v, + n0, + is_up, + t_neg, + one, + rcp6, + s_gate, + s_c, + mask, + num_tokens, + warp_idx, + lane, + hst, + ) + else: + self._zero_stage(n0, warp_idx, lane, hst) + self._zero_stage(n0 + 1, warp_idx, lane, hst) + else: + self._swiglu_pair( + v, + n0, + is_up, + t_neg, + one, + rcp6, + s_gate, + s_c, + mask, + num_tokens, + warp_idx, + lane, + hst, + ) + + @cute.jit + def _swiglu_pair( + self, v, n0, is_up, t_neg, one, rcp6, s_gate, s_c, mask, num_tokens, warp_idx, lane, hst + ): + """Tokens ``n0`` and ``n0 + 1`` of ``_swiglu_quant_to_stage``, through the packed + instructions (``*_x2``).""" + g0 = cute.arch.shuffle_sync_down(v[n0], 8) + g1 = cute.arch.shuffle_sync_down(v[n0 + 1], 8) + x = ptx.mul_rn_x2(g0, g1, t_neg, t_neg) + d = ptx.add_rn_x2(ptx.ex2_approx(x[0]), ptx.ex2_approx(x[1]), one, one) + gs = ptx.mul_rn_x2(g0, g1, ptx.rcp_approx_ftz(d[0]), ptx.rcp_approx_ftz(d[1])) + u = ptx.fma_rn_zero_x2(v[n0], v[n0 + 1], s_gate, s_gate) + hh = ptx.mul_rn_x2(u[0], u[1], gs[0], gs[1]) + h0 = Float32(0.0) + h1 = Float32(0.0) + if is_up: + h0 = hh[0] + h1 = hh[1] + # (amax|h|) * sC == max(|h| * sC): rounding is monotone and sC > 0 + a = ptx.mul_rn_x2(self._block_amax(h0), self._block_amax(h1), s_c, s_c) + q = ptx.mul_rn_x2(a[0], a[1], rcp6, rcp6) + self._quant_store(h0, q[0], s_c, n0, mask, num_tokens, warp_idx, lane, hst) + self._quant_store(h1, q[1], s_c, n0 + 1, mask, num_tokens, warp_idx, lane, hst) + + @cute.jit + def _zero_stage(self, n, warp_idx, lane, hst): + """Token ``n``'s staged bytes as ``_quant_store`` leaves an unrouted token: zero codes + and a zero scale.""" + if lane < 8: + ptx.sts_u8(hst + n * 32 + warp_idx * 8 + lane, Uint32(0)) + if lane == 0: + ptx.sts_u8(hst + self.g.hstage_sf + n * 4 + warp_idx, Uint32(0)) @cute.jit def _block_amax(self, h): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_allreduce_transport.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_allreduce_transport.py index e2b00832865d..6463d3dac83b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_allreduce_transport.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_allreduce_transport.py @@ -157,12 +157,17 @@ LAYOUT_RANK_COUNTS = {MAX_TOKENS: (4, 8), WIDE_MAX_TOKENS: (4,)} #: The built-in TP8 kernel's cluster-8 range; above it the built-in runs one CTA per token. CLUSTER8_TOKENS = 16 +#: The most tokens the TP8 layout may be widened to: five requests at K = 7 (40 rows), whose +#: 320 CTAs fit the ``grid_shift`` of 9 (512) the arrival counter then advances by. +TP8_WIDE_MAX_TOKENS = 40 + #: The layout of a TP8 workspace, whose calls of 17..TP8_MAX_TOKENS tokens add their #: slot sums flat, in the built-in one-CTA-per-token order (``flat_sum``, as at TP4; the built-in #: runs one CTA per token up to 128). TP4 workspaces ignore it (their 17..32-token calls are #: ``serve_wide_calls``'). Kept out of ``LAYOUT_RANK_COUNTS`` (``has_layout``), so a 32-token TP8 -#: layout never makes the TP4 wide layout's entry admit TP8. -TP8_MAX_TOKENS = 16 +#: layout never makes the TP4 wide layout's entry admit TP8. ``tep8_finalize_transport_wide`` +#: widens it to 17..``TP8_WIDE_MAX_TOKENS``; off, it stays ``MAX_TOKENS`` and nothing changes. +TP8_MAX_TOKENS = MAX_TOKENS #: CTAs per token (one cluster), threads per CTA, BF16 values per thread. CLUSTER = 8 VECTOR = 8 @@ -263,6 +268,34 @@ def label(self) -> str: ENTRY_TRIGGER = Transport(trigger="entry") +def configure(*, tep8_max_tokens: int = MAX_TOKENS, trigger: str = "late") -> None: + """Set the TP8 workspace layout and the trigger of the transport ``finalize_transport`` builds. + + ``lever_settings.apply`` sets them from the model's plan before any install; a workspace + built afterwards takes them. The defaults are the built-in layout and trigger place. + + Args: + tep8_max_tokens: the token layout of a TP8 workspace (``TP8_MAX_TOKENS``), + ``MAX_TOKENS``..``TP8_WIDE_MAX_TOKENS``. Above ``MAX_TOKENS`` its calls of + 17..``tep8_max_tokens`` tokens run the flat build (``tep8_finalize_transport_wide``). + trigger: the dependent-launch trigger of ``DEFAULT``: ``"late"`` or ``"entry"`` + (``finalize_entry_trigger``). + + Raises: + ValueError: for a token count outside the range, or a trigger ``Transport`` refuses. + """ + global TP8_MAX_TOKENS, DEFAULT + tokens = int(tep8_max_tokens) + if not MAX_TOKENS <= tokens <= TP8_WIDE_MAX_TOKENS: + raise ValueError( + f"[musespark_open] finalize transport token cap {tokens}: expected " + f"{MAX_TOKENS}..{TP8_WIDE_MAX_TOKENS}" + ) + transport = Transport(trigger=str(trigger).strip() or "late") + TP8_MAX_TOKENS = tokens + DEFAULT = transport + + try: from tensorrt_llm._torch.models import musespark_open_finalize_ar as _legacy except ImportError: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_two_shot.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_two_shot.py index 69fe2c2cfffc..225254bfcb7f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_two_shot.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_two_shot.py @@ -25,6 +25,12 @@ (buffer, ``t``, column vector) of every other rank, where that rank's CTA of the token reads it. Ingress per rank: ``(ranks - 1) / ranks * tokens * 8 KiB`` twice, behind two dependent exchanges. +* **Rank counts** (``RANK_COUNTS``) -- four ranks on one tray (CUDA IPC), or eight over two + trays of four in one NVLink domain (the receive memory is a fabric allocation every rank + maps, as the TP8 built-in's). The built-in TP8 kernel runs the thread map above from 17 + tokens on (below that it splits a token over a cluster of 8 CTAs), so at eight ranks a + call takes 17 or more tokens (``MIN_ROWS_BY_RANKS``) and stays bitwise the built-in's. + At eight ranks the built-in one-shot's ingress is ``7 * tokens * 8 KiB`` per rank. Lamport slots. Every slot starts as the sentinel word ``0x80000000``; a 32-bit word that still holds it has not arrived. No payload word is the sentinel: a partial has @@ -122,8 +128,16 @@ MAX_ROWS = 512 #: The largest call the built-in fused finalize serves in one launch. BUILTIN_MAX_TOKENS = 128 -#: The tensor-parallel group sizes the transport is laid out for. -RANK_COUNTS = (4,) +#: The tensor-parallel group sizes the transport is laid out for: four ranks on one tray, or +#: eight over two trays of four in one NVLink domain (fabric memory, as the TP8 built-in). +RANK_COUNTS = (4, 8) +#: The fewest tokens a call takes, by the group's rank count. Four ranks: ``MIN_ROWS``. Eight +#: ranks: 17. The protocol needs at least ``ranks`` tokens per call (every rank owns one), and +#: the built-in TP8 kernel runs this kernel's thread map (one 512-thread CTA per token) from 17 +#: tokens on (``musespark_open_finalize_ar._CLUSTER_TABLES[8]``); at 16 tokens and fewer it +#: splits a token over a cluster of 8 CTAs and adds the per-CTA sums of squares in another order. +#: So a call in the band is bitwise the built-in kernel's at either rank count. +MIN_ROWS_BY_RANKS = {4: MIN_ROWS, 8: 17} MAX_RANKS = 8 SENTINEL = 0x80000000 BUFFERS = 2 @@ -577,21 +591,25 @@ def _pdl_enabled() -> bool: return _transport._pdl_enabled() -def takes(num_tokens, max_rows: int = MAX_ROWS) -> bool: +def takes(num_tokens, max_rows: int = MAX_ROWS, min_rows: int = MIN_ROWS) -> bool: """Whether a call of ``num_tokens`` is in this entry's band (the same on every rank).""" return ( isinstance(num_tokens, int) and not isinstance(num_tokens, bool) - and MIN_ROWS <= num_tokens <= max_rows + and min_rows <= num_tokens <= max_rows ) class FinalizeTwoShotWorkspace: - """This transport's memory and state for one TP4 group, in front of another finalize workspace. + """This transport's memory and state for one TP group, in front of another finalize workspace. + + The group is four ranks on one tray, or eight over two trays of four in one NVLink domain + (``RANK_COUNTS``); across trays the receive memory is a fabric allocation every rank maps. ``finalize`` and ``finalize_chunked`` have the built-in ``FinalizeIpcWorkspace`` signatures. While ``active`` (set by the fast path's switch around a decode forward, - identical on every rank), calls of ``MIN_ROWS..max_rows`` tokens run this module's + identical on every rank), calls of ``min_rows..max_rows`` tokens (``min_rows``: + ``MIN_ROWS_BY_RANKS`` of the group's size) run this module's kernel; every other call goes to ``wrapped`` (the ``finalize_transport`` workspace or the built-in one), whose protocol state is separate, so the two interleave freely. Construction is collective over the group (fabric allocation, sentinel fill, compilation @@ -615,10 +633,14 @@ def __init__( ) if mapping.tp_size not in RANK_COUNTS: - raise ValueError(f"the two-shot finalize serves TP4, got TP{mapping.tp_size}") - if not MIN_ROWS <= max_rows <= GRID_STEP or max_rows % mapping.tp_size: raise ValueError( - f"max_rows must be a multiple of the TP size in {MIN_ROWS}..{GRID_STEP}, " + f"the two-shot finalize serves TP{' or TP'.join(map(str, RANK_COUNTS))}, " + f"got TP{mapping.tp_size}" + ) + self.min_rows = MIN_ROWS_BY_RANKS[mapping.tp_size] + if not self.min_rows <= max_rows <= GRID_STEP or max_rows % mapping.tp_size: + raise ValueError( + f"max_rows must be a multiple of the TP size in {self.min_rows}..{GRID_STEP}, " f"got {max_rows}" ) self.wrapped = wrapped @@ -635,13 +657,16 @@ def __init__( try: if not HAS_CUTE_DSL: raise RuntimeError("the CuTe DSL or the built-in finalize module is unavailable") - if mapping.is_multi_node(): - raise RuntimeError("the two-shot finalize serves one tray") + # Four ranks: one tray (CUDA IPC). Eight: the two trays' fabric (the topology is + # the fast path's to check: ``kernels/finalize_two_shot.install.refusal``). + use_fabric = mapping.is_multi_node() + if use_fabric and self.ranks != 8: + raise RuntimeError("the two-shot finalize serves TP4 on one tray") if torch.cuda.get_device_capability(self.device) != (10, 3): raise RuntimeError("the two-shot finalize runs on SM103 only") self.sms = torch.cuda.get_device_properties(self.device).multi_processor_count self._handle = _make_mnnvl_mcast_buffer( - self._comm, self.workspace_bytes, mapping, False + self._comm, self.workspace_bytes, mapping, use_fabric ) words = self.workspace_bytes // 4 self._uc = [ @@ -712,13 +737,13 @@ def _launch(self, gemm, indices, weights, norm, out, rows, row_stride, tokens, e def takes(self, num_tokens) -> bool: """Whether a call of ``num_tokens`` runs this kernel while active.""" - return takes(num_tokens, self.max_rows) + return takes(num_tokens, self.max_rows, self.min_rows) def _validate(self, gemm2_output, indices, routing_weights, norm_weight, eps, num_tokens): - """The built-in workspace's input contract, for ``MIN_ROWS..max_rows`` tokens.""" + """The built-in workspace's input contract, for ``min_rows..max_rows`` tokens.""" if not self.takes(num_tokens): raise ValueError( - f"num_tokens must be in [{MIN_ROWS}, {self.max_rows}], got {num_tokens!r}" + f"num_tokens must be in [{self.min_rows}, {self.max_rows}], got {num_tokens!r}" ) if not (isinstance(eps, (int, float)) and math.isfinite(eps) and eps > 0): raise ValueError(f"eps must be positive and finite, got {eps!r}") @@ -802,7 +827,7 @@ def finalize( eps: float, num_tokens: int, ) -> torch.Tensor: - """The wrapped ``finalize``, with ``MIN_ROWS..128`` tokens on this kernel while active.""" + """The wrapped ``finalize``, with ``min_rows..128`` tokens on this kernel while active.""" if self.active and self.takes(num_tokens) and num_tokens <= BUILTIN_MAX_TOKENS: return self.launch(gemm2_output, indices, routing_weights, norm_weight, eps, num_tokens) return self.wrapped.finalize( @@ -818,7 +843,7 @@ def finalize_chunked( eps: float, num_tokens: int, ) -> torch.Tensor: - """This kernel for ``MIN_ROWS..max_rows`` tokens while active, else the wrapped ``finalize_chunked``. + """This kernel for ``min_rows..max_rows`` tokens while active, else the wrapped ``finalize_chunked``. Raises: RuntimeError: for a call this kernel does not take when ``wrapped`` has no @@ -859,8 +884,9 @@ class FinalizeTwoShot(OpWrapper): Certified bit for bit against the built-in kernel (``finalize`` up to 128 tokens, ``finalize_chunked`` above) by its own four-rank test (every cell, a delayed rank, back-to-back calls of changing token counts, CUDA-graph replays with changing - inputs) and against `reference` within a band. Not certified: TP8, more than 512 - tokens. + inputs) and against `reference` within a band. At TP8 over two trays (17..512 tokens) + only by the bitwise comparison against the built-in TP8 fabric kernel of a multi-node + qualification run, not by a committed test. Not certified: more than 512 tokens. """ def __call__( @@ -874,7 +900,9 @@ def __call__( workspace: FinalizeTwoShotWorkspace, ) -> torch.Tensor: if not workspace.takes(num_tokens): - raise ValueError(f"finalize_two_shot serves {MIN_ROWS}..{workspace.max_rows} tokens") + raise ValueError( + f"finalize_two_shot serves {workspace.min_rows}..{workspace.max_rows} tokens" + ) return workspace.launch( gemm2_output, indices, routing_weights, norm_weight, eps, num_tokens ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/persistent_experts.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/persistent_experts.py index e02cd67fabb4..46559e3f8e3d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/persistent_experts.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/persistent_experts.py @@ -32,6 +32,33 @@ three-window stack ``[own | b-1 | b-2]``. The kernel's active-expert bit set is over window ids (two 32-bit words up to 64), so the ring's block size does not enter it. An expert's rows do not depend on the window it is read from. + +Wider calls on a window of at most 64 experts: the wide build (9..16 rows, +``wide_persistent_experts``), or with ``batch=True`` the batch builds (9..32 +rows: 16 or 32 token columns, ``BATCH_OPTIONS``), which keep the 8-row build's +six weight stages by reusing the activation's shared-memory stages for the FC2 +operand. Every build computes a routed row with the same instructions in the +same order, so the rows are bitwise the same whichever build serves a call. + +Early start (``early_start=True`` on a batch call; ``tep8_experts_early_start``): +the batch build variant that reads its epoch before its dependency wait, for a launch +whose same-stream predecessor released it right after its own dependency wait +(``ring_buddy_assign(early_trigger=True)``, the call that wrote ``topk_ids``), so the grid +launches and initializes while that predecessor runs; ``ring_buddy`` pairs the two calls. +Same workspace, same rows. + +Interleaved fetch (``tep8_experts_interleaved_fetch``, or ``fetch_kb`` on a +call): the batch builds' two fetch warps move the FC2 operand's k-blocks alternately instead of +in halves (a 32-column build ``BATCH_FETCH_KB`` per polling pass), so the k-blocks an FC2 tile +still waits for are polled by both warps in the MMA's order instead of by one warp in series +(``kernel.py``'s module docstring). Every build and variant of a batch call (early start +included) takes it; same workspace, same rows. + +In-kernel placement (``place_rank`` on a batch call on the 64-expert window; the call site is +``ring_buddy``'s ``tep8_experts_in_kernel_placement``): ``topk_ids`` are the +step's selected global ids and the kernel computes ``ring_buddy_assign``'s balanced tep8 +placement for rank ``place_rank`` itself, so no placement launch precedes it. The default +entry (its epoch after the wait, so any predecessor will do); same workspace, same rows. """ from __future__ import annotations @@ -57,6 +84,27 @@ #: What the wide build changes in the device's ``OPTIONS``: its B operand and compact scale #: stages are twice the 8-row build's, so four weight stages fit in shared memory. WIDE_OPTIONS: Dict[str, int] = {"num_stages": 4} +#: Rows the batch builds serve (``batch=True``; windows of at most 64 experts): sm103 TP8 +#: decode steps of two to four requests at K = 7 on the ring-buddy stack +#: (``tep8_persistent_batch_experts``). +BATCH_MAX_ROWS = 32 +#: The batch builds' token columns: the smallest that holds a call's rows above 8. +BATCH_NTOKS = (16, 32) +#: What a batch build changes in the device's ``OPTIONS`` (``kernel.py``'s ``overlap``): the +#: FC2 operand image reuses the activation's 16 B stages, which keeps the six weight +#: stages at 32 token columns, and the activation lands in groups of k-blocks; the image +#: moves by the tagged hand-off (``OPTIONS``' ``handoff``) on per-k-block barriers; the +#: accumulator stages fit 64 TMEM columns. +BATCH_OPTIONS: Dict[int, Dict[str, int]] = { + 16: {"overlap": 1}, + 32: {"overlap": 1, "acc_stages": 2}, +} +#: Interleaved fetch (``tep8_experts_interleaved_fetch``): k-blocks one fetch warp of a +#: 32-column batch build polls +#: per pass (``kernel.py``'s ``fetch_kb``; the 16-column build keeps its four). One, as before: +#: the interleave alone carries the gain, and two per pass measured no better (four, slower). +BATCH_FETCH_KB = 1 +#: The switch of the batch builds' interleaved fetch: unset, empty or ``0`` is off, ``1`` on. #: Expert windows the kernel is compiled for. WINDOWS = (32, 64, 96) SF_VEC = 16 @@ -180,6 +228,55 @@ def _refuse_under_capture(what: str) -> None: ) +_FETCH_LOGGED: list = [] +#: ``tep8_experts_interleaved_fetch``: the batch builds' ``fetch_kb``, 0 while it is off. +_FETCH_KB = 0 + + +def configure(*, interleaved_fetch: bool = False) -> None: + """Set the batch builds' interleaved fetch. + + ``lever_settings.apply`` sets it from the model's plan before any install; off, the + default, leaves ``fetch_kb`` at 0 and the builds fetch as before. + + Args: + interleaved_fetch: the ``tep8_experts_interleaved_fetch`` lever; on, the builds take + ``BATCH_FETCH_KB`` (``_FETCH_KB``). + """ + global _FETCH_KB + _FETCH_KB = BATCH_FETCH_KB if interleaved_fetch else 0 + + +def batch_fetch_kb() -> int: + """The batch builds' ``fetch_kb``: ``BATCH_FETCH_KB`` under + ``tep8_experts_interleaved_fetch``, else 0.""" + return _FETCH_KB + + +def _log_fetch_once(fetch_kb: int) -> None: + """One INFO line per process, when the batch builds are prepared with the interleaved fetch.""" + if _FETCH_LOGGED: + return + _FETCH_LOGGED.append(True) + from tensorrt_llm.logger import logger + + logger.info( + "[musespark_open] tep8_experts_interleaved_fetch: the batch builds' two fetch warps " + f"move the FC2 operand's k-blocks alternately ({fetch_kb} per polling pass at 32 columns, " + "4 at 16 columns); logged once per process" + ) + + +def batch_ntok(rows: int) -> int: + """The batch build that serves a call of ``rows`` (9..``BATCH_MAX_ROWS``) rows.""" + for ntok in BATCH_NTOKS: + if rows <= ntok: + return ntok + raise ValueError( + f"persistent_experts: the batch builds serve at most {BATCH_MAX_ROWS} rows, got {rows}" + ) + + class _Workspace: """Graph-stable per-(device, grid) scratch shared by every launch on a stream. @@ -195,7 +292,9 @@ class _Workspace: once, outside graph capture, and never re-initialized. """ - def __init__(self, device: torch.device, num_ctas: int, ntok: int = MAX_ROWS): + def __init__( + self, device: torch.device, num_ctas: int, ntok: int = MAX_ROWS, batch: bool = False + ): from ._persistent_experts.kernel import FC1_MT, MAX_CTAS, MAX_SLOTS, TRACE_WORDS, Geometry assert num_ctas <= MAX_CTAS @@ -203,25 +302,32 @@ def __init__(self, device: torch.device, num_ctas: int, ntok: int = MAX_ROWS): self.epoch = torch.zeros(MAX_CTAS, dtype=torch.int32, device=device) self.flags = torch.zeros(MAX_SLOTS * FC1_MT, dtype=torch.int32, device=device) self.images = torch.zeros( - MAX_SLOTS * Geometry(ntok).ws_slot_bytes, dtype=torch.uint8, device=device + MAX_SLOTS * Geometry(ntok, batch).ws_slot_bytes, dtype=torch.uint8, device=device ) self.trace = torch.zeros(num_ctas, TRACE_WORDS, dtype=torch.int64, device=device) def _workspace( - device: torch.device, num_ctas: int, handoff: Optional[int] = None, ntok: int = MAX_ROWS + device: torch.device, + num_ctas: int, + handoff: Optional[int] = None, + ntok: int = MAX_ROWS, + batch: bool = False, ) -> _Workspace: - """The workspace of this device, grid, hand-off protocol and build width (its own epochs).""" + """The workspace of this device, grid, hand-off protocol and build width (its own epochs); + a batch build has its own.""" if handoff is None: handoff = _config(device)[1]["handoff"] key = (_device_index(device), num_ctas, int(handoff)) if ntok != MAX_ROWS: key = key + (int(ntok),) + if batch: + key = key + ("batch",) with _LOCK: ws = _WORKSPACES.get(key) if ws is None: _refuse_under_capture("workspace") - ws = _Workspace(device, num_ctas, ntok) + ws = _Workspace(device, num_ctas, ntok, batch) _WORKSPACES[key] = ws return ws @@ -300,16 +406,40 @@ def _compiled( w2_sf, early=None, ntok: int = MAX_ROWS, + batch: bool = False, + early_start: bool = False, + fetch_kb: int = 0, + place_rank: int = -1, ): """(compiled launcher, grid size, options) for this device's configuration and window. With ``early`` (an ``EarlyExpertIds``), the variant that takes its ids from that buffer's tagged words before its dependency wait (``early_expert_ids``). With ``ntok = WIDE_MAX_ROWS``, the wide build (``WIDE_OPTIONS`` applied; no early ids). + With ``batch``, the batch build of ``ntok`` (``BATCH_NTOKS``) columns + (``BATCH_OPTIONS[ntok]`` applied; no early ids); with ``early_start`` as well, its + early-start variant (see the module docstring); with ``fetch_kb`` (non-zero), its + interleaved fetch; with ``place_rank`` (0..7), its in-kernel placement for that rank. """ num_ctas, options = _config(device) extra = _early_options(early) - if ntok != MAX_ROWS: + assert batch or not early_start, "early start is a batch build variant" + assert batch or not fetch_kb, "the interleaved fetch is a batch build variant" + assert place_rank < 0 or (batch and not early_start and num_local == 64), ( + "the in-kernel placement is a batch build variant on the 64-expert window, default entry" + ) + if batch: + assert early is None, "the batch builds have no early-ids variant" + assert num_local <= 64, "the batch builds serve windows of at most 64 experts" + options = dict(options, **BATCH_OPTIONS[ntok]) + if early_start: + options["early_start"] = 1 + if fetch_kb: + options["fetch_kb"] = int(fetch_kb) + if place_rank >= 0: + options["place_rank"] = int(place_rank) + extra = {"ntok": int(ntok)} + elif ntok != MAX_ROWS: assert early is None, "the wide build has no early-ids variant" options = dict(options, **WIDE_OPTIONS) extra = {"ntok": int(ntok)} @@ -397,7 +527,10 @@ class PersistentExperts(OpWrapper): """NVFP4 experts of one window for <= 8 pre-routed rows, one persistent launch. Calls of 9..16 rows on a window of at most 64 experts run the wide build (MMA N = 16, - its own workspace), with the same contract and the same per-row arithmetic. + its own workspace), with the same contract and the same per-row arithmetic. With + ``batch=True``, calls of 9..``BATCH_MAX_ROWS`` rows on such a window run the batch + builds instead (MMA N = 16 or 32, ``BATCH_OPTIONS``, a workspace each), again with the + same contract and per-row arithmetic. Arguments mirror the pre-routed ``fp4_block_scale_moe_runner`` call the MuseSpark layer makes: NVFP4 rows ``hidden_states`` ``[M, 2048]`` uint8 with @@ -459,6 +592,10 @@ def __call__( output1_scale_scalar: torch.Tensor, output1_scale_gate_scalar: torch.Tensor, output2_scale_scalar: torch.Tensor, + batch: bool = False, + early_start: bool = False, + fetch_kb: Optional[int] = None, + place_rank: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: import cutlass from cuda.bindings import driver as cuda @@ -466,7 +603,23 @@ def __call__( device = hidden_states.device rows_in = hidden_states.shape[0] num_local = gemm1_weights.shape[0] - ntok = MAX_ROWS if rows_in <= MAX_ROWS else WIDE_MAX_ROWS + batch = bool(batch) and rows_in > MAX_ROWS + early_start = bool(early_start) and batch + place = -1 + if place_rank is not None: + # topk_ids are global: only the batch build's in-kernel placement may read them + assert batch and num_local == 64 and 0 <= int(place_rank) < 8, ( + "persistent_experts: place_rank needs a batch call of 9..32 rows on the " + "64-expert window" + ) + place = int(place_rank) + early_start = False + fetch = 0 + if batch: + ntok = batch_ntok(rows_in) + fetch = batch_fetch_kb() if fetch_kb is None else int(fetch_kb) + else: + ntok = MAX_ROWS if rows_in <= MAX_ROWS else WIDE_MAX_ROWS fn, num_ctas, options = _compiled( device, num_local, @@ -475,8 +628,12 @@ def __call__( gemm2_weights, gemm2_weights_scale, ntok=ntok, + batch=batch, + early_start=early_start, + fetch_kb=fetch, + place_rank=place, ) - ws = _workspace(device, num_ctas, options["handoff"], ntok) + ws = _workspace(device, num_ctas, options["handoff"], ntok, batch) if _EARLY and ntok == MAX_ROWS: from . import early_expert_ids as _early_ids @@ -523,15 +680,57 @@ def prepare( gemm2_weights_scale: torch.Tensor, early=None, wide: bool = False, + batch_rows: int = 0, + early_start: bool = False, + fetch_kb: Optional[int] = None, + place_rank: Optional[int] = None, ) -> None: """Compile the kernel for this window and create its workspace, now. Call it eagerly (at install time, before any graph capture): a call under capture that would have to do either refuses. With ``early`` (the workspace's ``EarlyExpertIds``) the early-ids variant - is compiled too; with ``wide``, the wide build and its workspace as well. + is compiled too; with ``wide``, the wide build and its workspace as well; + with ``batch_rows`` above ``MAX_ROWS``, every batch build that serves 9..``batch_rows`` + rows and its workspace, and with ``early_start`` their early-start variants too, all + with the batch calls' fetch (``fetch_kb``, by default ``batch_fetch_kb()``); with + ``place_rank`` (0..7, the 64-expert window) their in-kernel placement variants as well. """ device = gemm1_weights.device + if batch_rows > MAX_ROWS: + fetch = batch_fetch_kb() if fetch_kb is None else int(fetch_kb) + if fetch and fetch_kb is None: + _log_fetch_once(fetch) + for ntok in BATCH_NTOKS: + if ntok > batch_ntok(batch_rows): + break + for early_variant in (False, True) if early_start else (False,): + _, num_ctas, options = _compiled( + device, + gemm1_weights.shape[0], + gemm1_weights, + gemm1_weights_scale, + gemm2_weights, + gemm2_weights_scale, + ntok=ntok, + batch=True, + early_start=early_variant, + fetch_kb=fetch, + ) + if place_rank is not None: + _compiled( + device, + gemm1_weights.shape[0], + gemm1_weights, + gemm1_weights_scale, + gemm2_weights, + gemm2_weights_scale, + ntok=ntok, + batch=True, + fetch_kb=fetch, + place_rank=int(place_rank), + ) + _workspace(device, num_ctas, options["handoff"], ntok, True) if wide: _, num_ctas, options = _compiled( device, @@ -577,15 +776,28 @@ def reference( output1_scale_scalar: torch.Tensor, output1_scale_gate_scalar: torch.Tensor, output2_scale_scalar: torch.Tensor, + batch: bool = False, + early_start: bool = False, + fetch_kb: Optional[int] = None, + place_rank: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """The layer in plain torch, FP64, over the prepared (shuffled) operands. + ``batch``, ``early_start`` and ``fetch_kb`` (which build serves the call, when it loads + its activation, and how it moves the FC2 operand) do not change the result and are + ignored. With ``place_rank``, ``topk_ids`` are global and are placed first by + ``ring_buddy_assign``'s reference (balanced, tep8) for that rank. + Undoes the 32-row shuffle and the gate/up interleave by index, decodes the swizzled scales, and requantizes the SwiGLU output to NVFP4 with exact division (the kernels use approximate reciprocals). Rows of non-local slots are zero here (the op leaves them unwritten); a local expert a token repeats is routed through its first slot only. """ + if place_rank is not None: + from .ring_buddy_assign import ring_buddy_assign as _assign + + topk_ids = _assign.reference(topk_ids, topk_weights.float(), int(place_rank))[0] m = hidden_states.shape[0] num_local = gemm1_weights.shape[0] device = hidden_states.device @@ -633,11 +845,22 @@ def is_valid( output1_scale_scalar: torch.Tensor, output1_scale_gate_scalar: torch.Tensor, output2_scale_scalar: torch.Tensor, + batch: bool = False, + early_start: bool = False, + fetch_kb: Optional[int] = None, + place_rank: Optional[int] = None, ) -> None: """Inputs the kernel would read past or misread without an error.""" m = hidden_states.shape[0] num_local = gemm1_weights.shape[0] - limit = WIDE_MAX_ROWS if num_local <= 64 else MAX_ROWS + if place_rank is not None: + assert batch and m > MAX_ROWS and num_local == 64 and 0 <= place_rank < 8, ( + "place_rank: a batch call of 9..32 rows on the 64-expert window, rank 0..7" + ) + if num_local > 64: + limit = MAX_ROWS + else: + limit = BATCH_MAX_ROWS if batch else WIDE_MAX_ROWS assert 1 <= m <= limit, f"one launch serves 1..{limit} rows on this window, got {m}" assert hidden_states.dtype == torch.uint8 and hidden_states.is_contiguous(), ( "hidden_states must be contiguous packed NVFP4 bytes" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/ring_buddy_assign.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/ring_buddy_assign.py index 5bd6b4e2de49..98bce3fb484e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/ring_buddy_assign.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/ring_buddy_assign.py @@ -46,6 +46,12 @@ alone (``local_num_experts = BLOCK``), the owner-only window ``ring_buddy.install(windows=1)`` gives ``persistent_experts`` and ``bf16_routed_experts`` without replication. + +``early_trigger=True`` launches the program with programmatic dependent launch: it waits +for the kernel that wrote ``selected``, then releases its own dependents at once, so the +expert call behind it (the persistent kernel's early-start variant, +``tep8_experts_early_start``) launches and initializes while this program +still runs. The placement is the same. """ from __future__ import annotations @@ -81,9 +87,51 @@ def block_of(ranks: int) -> int: raise ValueError(f"ring_buddy_assign serves rings of {sorted(RING_SIZES)} ranks, got {ranks}") -#: Largest ``M * TOP_K`` one program handles: 64 decode rows (eight requests' K = 7 verify -#: steps): the same [SLOTS, RANKS] program at a 512-slot tile. -MAX_SLOTS = 512 +#: Largest ``M * TOP_K`` one program handles: 256 decode rows (32 requests' K = 7 verify +#: steps), the same [SLOTS, RANKS] program at a 2048-slot tile. The stacked path's own cap +#: stays at ``ring_buddy.MAX_ROWS``; the tiles above 512 slots serve the tep8 decode steps +#: ``tep8_ring_buddy_wide`` admits. +MAX_SLOTS = 2048 + +#: ``tep8_wide_balance``'s ``MAX_ROWS`` (default off): on tep8 the balanced placement of the +#: NVFP4 layers also serves the decode calls above the stacked path's row cap, up to those rows +#: (129..256, ``8 * B`` for B requests at K = 7), each call of more than 64 rows placed by the +#: wide program. +WIDE_BALANCE_MIN_ROWS = 129 +_WIDE_BALANCE_ROWS: Optional[int] = None + + +def configure(*, wide_balance_rows: Optional[int] = None) -> None: + """Set the band the balanced NVFP4 placement adds above the stacked path's cap. + + ``lever_settings.apply`` sets it from the model's plan (through + ``ring_buddy.configure``) before any install; None, the default, adds no band. + + Args: + wide_balance_rows: ``tep8_wide_balance``'s ``MAX_ROWS`` (``_WIDE_BALANCE_ROWS``), + in ``WIDE_BALANCE_MIN_ROWS``..``MAX_SLOTS // TOP_K``. + + Raises: + ValueError: for a row cap outside that range. + """ + global _WIDE_BALANCE_ROWS + if wide_balance_rows is None: + _WIDE_BALANCE_ROWS = None + return + rows = int(wide_balance_rows) + if not WIDE_BALANCE_MIN_ROWS <= rows <= MAX_SLOTS // TOP_K: + raise ValueError( + f"[musespark_open] tep8_wide_balance MAX_ROWS {rows}: expected the balanced " + f"placement's row cap for the NVFP4 layers' decode calls, a whole number in " + f"{WIDE_BALANCE_MIN_ROWS}..{MAX_SLOTS // TOP_K}" + ) + _WIDE_BALANCE_ROWS = rows + + +def wide_balance_rows() -> Optional[int]: + """The row cap ``configure`` set, or None when the lever is off.""" + return _WIDE_BALANCE_ROWS + #: The half window (``half_window_experts``, tep4): a block's upper half starts at this #: local index, and the stack ``[own block | the predecessor's upper half]`` holds @@ -100,6 +148,39 @@ def block_of(ranks: int) -> int: #: 3.08 / 2.37 / 1.73 / 2.15; a one-kernel BF16 cast of the weights is 0.8). NUM_WARPS = 4 + +def num_warps_for(slots: int) -> int: + """Warps of the one program at a ``slots`` tile: ``NUM_WARPS`` up to 512 slots, then one + warp per 128 slots (8 at 1024, 16 at 2048), so every thread keeps the 32 elements of each + ``[SLOTS, RANKS]`` census tensor it holds at 512 slots instead of spilling them.""" + return NUM_WARPS if slots <= 512 else min(16, slots // 128) + + +#: Warps of the wide program (``wide=True``, ``_ring_buddy_assign_wide_kernel``) at every +#: tile it serves (1024 and 2048 slots). GB300, CUDA graph of 62 back-to-back launches, us per +#: launch at 4 / 8 / 16 warps: 1024 slots 1.95-1.98 / 1.89-1.95 / 2.12, 2048 slots +#: 2.22-2.38 / 2.09-2.25 / 2.15-2.28; the narrow program takes 3.4-3.6 at 1024 slots and +#: 5.8-5.9 at 2048, the BF16 cast it replaces 1.2-1.4. +WIDE_NUM_WARPS = 8 + +#: The wide program serves tiles above this many slots (more than 64 decode rows). +WIDE_MIN_SLOTS = 512 + + +def wide_serves( + slots: int, ranks: int, balance: bool, static_from: Optional[int], half: bool +) -> bool: + """Whether ``wide=True`` runs a call of a ``slots`` tile on the wide program: the + balanced tep8 placement (no half window, no static move) above ``WIDE_MIN_SLOTS``.""" + return ( + slots > WIDE_MIN_SLOTS + and ranks == 8 + and bool(balance) + and not half + and (static_from is None or static_from == block_of(ranks)) + ) + + #: At TP8 a call with ``meta_tile`` also writes the routing #: tables TRTLLM-gen's expert op reads for the two-window ``[own block | predecessor's block]`` #: call (``local_expert_offset=0``, ``local_num_experts=2*BLOCK``): expanded->permuted (``e2p``, @@ -175,9 +256,16 @@ def _ring_buddy_assign_kernel( META: tl.constexpr, TILE: tl.constexpr, CTAS: tl.constexpr, + PDL: tl.constexpr, ): tl.static_assert(BLOCK == 32 or BLOCK == 64, "one 32- or 64-bit presence word per block") tl.static_assert(not HALF or (BLOCK == 64 and BALANCE), "the half window: tep4, balanced") + if PDL: + # ``early_trigger``: every grid before this one has completed once the wait returns + # (``selected`` is final); the dependents may launch from then on, and read this + # program's outputs only after their own wait. + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() WIDE: tl.constexpr = BLOCK == 64 UPPER: tl.constexpr = BLOCK // 2 # HALF_FROM SPAN: tl.constexpr = BLOCK + BLOCK // 2 # HALF_SPAN @@ -292,6 +380,135 @@ def _ring_buddy_assign_kernel( tl.store(stats_ptr + RANKS + ranks, load - x + x_prev) +@triton.jit +def _ring_buddy_assign_wide_kernel( + selected_ptr, # [num_slots] int32 global expert ids + weights_ptr, # [num_slots] float32 routing weights + local_ids_ptr, # [num_slots] int32 out + weights_bf16_ptr, # [num_slots] bfloat16 out + stats_ptr, # [2 * RANKS] int32 out: per-rank loads before, after + rank, + num_slots, + SLOTS: tl.constexpr, + EXPERTS: tl.constexpr, + RANKS: tl.constexpr, + BLOCK: tl.constexpr, + PDL: tl.constexpr, +): + """``_ring_buddy_assign_kernel``'s balanced two-window placement (tep8) with per-expert + instead of ``[SLOTS, RANKS]`` work, for the wide tiles: the census is a histogram of the + step's ids, the move decision is made once per expert, and each slot gathers its expert's + local id. Every output is the same integer function of ``selected`` and ``rank``.""" + tl.static_assert(BLOCK == 32 and RANKS * BLOCK == EXPERTS, "the tep8 ring of 32-expert blocks") + if PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + s = tl.arange(0, SLOTS) + smask = s < num_slots + g = tl.load(selected_ptr + s, mask=smask, other=0) + w = tl.load(weights_ptr + s, mask=smask, other=0.0) + experts = tl.arange(0, EXPERTS) + ranks = tl.arange(0, RANKS) + + # Census: the per-expert slot counts (the SLOTS - num_slots masked-off slots load + # expert 0); a block's presence word sets bit i for each of its selected experts, so a + # repeated expert counts once. + counts = tl.histogram(g, EXPERTS) - tl.where(experts == 0, SLOTS - num_slots, 0) + i_e = experts % BLOCK + bit_e = tl.where( + counts > 0, + tl.full([EXPERTS], 1, tl.uint32) << i_e.to(tl.uint32), + tl.zeros([EXPERTS], dtype=tl.uint32), + ) + word = tl.reduce(tl.reshape(bit_e, [RANKS, BLOCK]), 1, _bit_or) # [RANKS] uint32 + load = _popc(word, False) # [RANKS] + + # The balanced placement's smallest maximum T and smallest flow x (as the narrow program). + span = ranks + 1 + back = (ranks[:, None] - ranks[None, :] + RANKS) % RANKS # [r, k] + inarc = back[:, None, :] < span[None, :, None] # [r, n, k] + arc = tl.sum(tl.where(inarc, load[None, None, :], 0), axis=2) # [r, n] + need = tl.where( + span[None, :] == RANKS, + (arc + RANKS - 1) // RANKS, + (arc + span[None, :]) // (span[None, :] + 1), + ) + t = tl.max(tl.max(need, axis=1), axis=0) + x = tl.max(tl.maximum(arc - span[None, :] * t, 0), axis=1) # [RANKS] + + # Per expert: it moves to its block's successor when fewer than x[b] selected experts of + # its block sit above it; its local id on this rank, or one outside the window. + word_e = tl.reshape(tl.broadcast_to(word[:, None], [RANKS, BLOCK]), [EXPERTS]) + x_e = tl.reshape(tl.broadcast_to(x[:, None], [RANKS, BLOCK]), [EXPERTS]) + blk_e = experts // BLOCK + above_e = _popc((word_e >> i_e.to(tl.uint32)) >> 1, False) + er_e = (blk_e + (above_e < x_e).to(tl.int32)) % RANKS + local_e = tl.where( + er_e == rank, + tl.where(blk_e == rank, i_e, BLOCK + i_e), + 2 * BLOCK + experts % (EXPERTS - 2 * BLOCK), + ) + local = tl.gather(local_e, g, 0) # [SLOTS] + tl.store(local_ids_ptr + s, local, mask=smask) + tl.store(weights_bf16_ptr + s, w.to(tl.bfloat16), mask=smask) + + from_prev = ((ranks[:, None] + RANKS - 1) % RANKS) == ranks[None, :] + x_prev = tl.sum(tl.where(from_prev, x[None, :], 0), axis=1) + tl.store(stats_ptr + ranks, load) + tl.store(stats_ptr + RANKS + ranks, load - x + x_prev) + + +#: Fields of one band of ``accumulate_placement_stats``'s accumulator (int64 sums, in order): +#: placements, their rows, the busiest rank's distinct experts under the owner-only +#: placement and under the computed one, and the distinct experts of the whole ring; the +#: last three are zero. +PLACEMENT_STATS_FIELDS = 8 + + +@triton.jit +def _placement_stats_kernel( + stats_ptr, # [2 * RANKS] int32, ring_buddy_assign's per-rank loads before, after + acc_ptr, # [8] int64 in/out, one band of the accumulator + rows, + RANKS: tl.constexpr, +): + r = tl.arange(0, RANKS) + before = tl.load(stats_ptr + r).to(tl.int64) + after = tl.load(stats_ptr + RANKS + r).to(tl.int64) + f = tl.arange(0, 8) + add = tl.where( + f == 0, + 1, + tl.where( + f == 1, + rows, + tl.where( + f == 2, + tl.max(before, axis=0), + tl.where( + f == 3, tl.max(after, axis=0), tl.where(f == 4, tl.sum(before, axis=0), 0) + ), + ), + ), + ).to(tl.int64) + tl.store(acc_ptr + f, tl.load(acc_ptr + f) + add) + + +def accumulate_placement_stats( + stats: torch.Tensor, acc: torch.Tensor, rows: int, ranks: int = RANKS +) -> None: + """Add one placement's census (``stats``, the entry's third output) to ``acc``. + + ``acc`` is ``PLACEMENT_STATS_FIELDS`` contiguous int64: device memory, or pinned host + memory that a host thread reads without a CUDA call (``ring_buddy``'s + ``ring_buddy_stats``). One program and a plain read-modify-write, so the + calls that share ``acc`` must be ordered by one stream, as every placement of a process + is (the routing method's call, on its main stream). Diagnostics only: the placement and + the expert call do not read ``acc``. + """ + _placement_stats_kernel[(1,)](stats, acc, rows, RANKS=ranks, num_warps=1) + + def balance_ring(load: List[int]) -> Tuple[int, List[int]]: """The kernel's placement in plain Python: ``(max load, x)``. @@ -330,13 +547,19 @@ class RingBuddyAssign(OpWrapper): (tep4, ``half_window_experts``) is the balanced placement onto the successor's copy of each block's upper half only (``ring_windows.place_half``): its ids address the 96-expert ``[own block | predecessor's experts 32..63]`` stack. + ``wide=True`` runs the balanced tep8 placement (``ranks=8``, ``balance``, no + ``half``, no static move) of a tile above ``WIDE_MIN_SLOTS`` on the wide program + (``_ring_buddy_assign_wide_kernel``): the same outputs, from per-expert work instead + of ``[SLOTS, RANKS]`` work; every other call ignores it. A call that also writes the + prerouted tables (``meta_tile``) keeps the ``[SLOTS, RANKS]`` program, which is where + they are built. """ ARCHS = frozenset({Arch.SM_103}) note = """ Pure function of ``selected`` and ``rank``; one program, so it is meant for - decode rows (``M * 8 <= 256``). Duplicate ids within a row are not expected + decode rows (``M * 8 <= MAX_SLOTS``). Duplicate ids within a row are not expected (top-k selection is distinct) and are counted once in the census. The rewritten ids are meaningful only for an expert call with @@ -357,6 +580,8 @@ def __call__( ranks: int = RANKS, half: bool = False, meta_tile: Optional[int] = None, + wide: bool = False, + early_trigger: bool = False, ): block = block_of(ranks) static_from = block if static_from is None else static_from @@ -380,6 +605,23 @@ def __call__( meta["rows"] = rows meta["tile"] = meta_tile ctas = max(16, triton.next_power_of_2(sizes["c2b"])) + if wide and meta is None and wide_serves(slots, ranks, balance, static_from, half): + _ring_buddy_assign_wide_kernel[(1,)]( + selected, + weights, + local_ids, + weights_bf16, + stats, + rank, + num_slots, + SLOTS=slots, + EXPERTS=EXPERTS, + RANKS=ranks, + BLOCK=block, + PDL=False, + num_warps=WIDE_NUM_WARPS, + ) + return local_ids, weights_bf16, stats dummy = stats _ring_buddy_assign_kernel[(1,)]( selected, @@ -405,7 +647,9 @@ def __call__( META=meta is not None, TILE=meta_tile if meta is not None else 8, CTAS=ctas, - num_warps=_META_WARPS if meta is not None else NUM_WARPS, + PDL=bool(early_trigger), + num_warps=_META_WARPS if meta is not None else num_warps_for(slots), + launch_pdl=bool(early_trigger), ) if meta is not None: return local_ids, weights_bf16, stats, meta @@ -421,9 +665,12 @@ def reference( ranks: int = RANKS, half: bool = False, meta_tile: Optional[int] = None, + wide: bool = False, + early_trigger: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """The placement in plain Python over the host copy of ``selected`` - (``ring_windows.place_half`` with ``half``).""" + (``ring_windows.place_half`` with ``half``; ``wide`` changes only which program + computes it, and ``early_trigger`` only when the dependents launch).""" if half: from ...kernels.replicated_experts import windows as _windows @@ -481,6 +728,8 @@ def is_valid( ranks: int = RANKS, half: bool = False, meta_tile: Optional[int] = None, + wide: bool = False, + early_trigger: bool = False, ) -> None: """Inputs the kernel would read past or misplace without an error.""" assert meta_tile is None or (not half and meta_tile in (8, 16, 32, 64)), ( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_gemm.py index 8d861ea010c4..58a68802fecb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_gemm.py @@ -38,6 +38,20 @@ Programmatic dependent launch: the GEMM waits (``griddepcontrol.wait``) before its first load and releases its dependents after its last MMA; the sum kernel waits before reading the partials and releases its dependents. + +``router_staged_gemm`` (default off; set once, when the model is +imported; ``set_staged`` for tests and benches): a call of ``STAGED_MIN_ROWS..STAGED_MAX_ROWS`` +rows -- two 128-row tiles, 38 K splits on GB300, a TEP8 decode step of 17..32 requests at +K=7 -- runs its GEMM as ``ops/moe/router_gemm_staged.Sm100BF16x3RouterGemmStaged`` (the +feature branch's sm107 kernel, #21754): the same tile, pipeline, K tiles per CTA, weight +split, MMAs, accumulators and per-element add, with the finished partial tile written +through shared memory (``bulk``: one 512-byte bulk copy per row; ``v4``: 16-byte vector +stores) instead of one 4-byte store per element, and with ``-prewait`` every weight tile +of the CTA (at most four here) loaded, and the first two split into tensor memory, before +``griddepcontrol.wait``. The partials, hence the logits, the selected experts and their +weights, are this entry's bits exactly; the sum kernel is unchanged. ``1`` selects +``STAGED_DEFAULT``. Every other call, and every call while the switch is off, runs the +tree's kernel as before. """ from __future__ import annotations @@ -93,6 +107,77 @@ def staged_epilogue() -> Optional[str]: return _staged_epilogue +#: The staged GEMM's row band: the calls of two 128-row tiles. +STAGED_MIN_ROWS = ROW_TILE + 1 +STAGED_MAX_ROWS = 2 * ROW_TILE +#: The staged GEMM's builds: (epilogue, weight loads before the dependency wait). Every +#: epilogue is a ``STAGED_MODES`` entry; the pre-wait builds are this target's (the sm_103 +#: ``bf16x3_router_gemm`` has no pre-wait build of its own). +STAGED_BUILDS = { + "bulk": ("bulk", False), + "v4": ("v4", False), + "bulk-prewait": ("bulk", True), + "v4-prewait": ("v4", True), +} +#: The build the ``router_staged_gemm`` lever selects when its ``BUILD`` names none. +STAGED_DEFAULT = "bulk-prewait" + + +def staged_build(value: Optional[str]) -> Optional[str]: + """The ``STAGED_BUILDS`` entry a switch value names: None for unset, ``""`` or ``0``. + + Raises: + ValueError: for any other value outside ``STAGED_BUILDS`` and ``1`` (on every rank alike). + """ + value = "" if value is None else str(value).strip().lower() + if value in ("", "0"): + return None + if value == "1": + return STAGED_DEFAULT + if value not in STAGED_BUILDS: + raise ValueError( + f"[musespark_open] router staged build {value!r}: expected 0, 1 or one of " + f"{sorted(STAGED_BUILDS)}" + ) + return value + + +#: ``router_staged_gemm``: the build calls of the staged band run, None for the tree's kernel. +_staged: Optional[str] = None + + +def configure(*, staged: Optional[str] = None) -> None: + """Set the staged build calls of ``STAGED_MIN_ROWS..STAGED_MAX_ROWS`` rows run. + + ``lever_settings.apply`` sets it from the model's plan before any install; None, the + default, keeps the tree's kernel at every row count. + + Args: + staged: a ``STAGED_BUILDS`` name, or None (``_staged``). + + Raises: + ValueError: as ``staged_build``. + """ + set_staged(staged) + + +def set_staged(value: Optional[str]) -> None: + """Run every later call of ``STAGED_MIN_ROWS..STAGED_MAX_ROWS`` rows with the staged build + ``value`` names (``staged_build``; None or ``0``: the tree's kernel). Before capture.""" + global _staged + _staged = staged_build(value) + + +def staged() -> Optional[str]: + """The staged build calls of ``STAGED_MIN_ROWS..STAGED_MAX_ROWS`` rows run, or None.""" + return _staged + + +def staged_takes(rows: int) -> bool: + """Whether a call of ``rows`` rows runs the staged GEMM.""" + return _staged is not None and STAGED_MIN_ROWS <= rows <= STAGED_MAX_ROWS + + def row_tile(rows: int) -> int: """The GEMM's N tile for a call of ``rows`` rows: the next power of two, 8..128.""" return min(max(triton.next_power_of_2(rows), 8), ROW_TILE) @@ -116,7 +201,13 @@ def splits(rows: int, sms: Optional[int] = None) -> int: return sms // (EXPERT_TILES * triton.cdiv(rows, ROW_TILE)) -def _gemm(tile: int): +def _gemm(tile: int, build: Optional[str] = None): + """The compiled GEMM for ``tile``-row tiles: the tree's kernel, or the staged ``build``.""" + if build is not None: + from .router_gemm_staged import Sm100BF16x3RouterGemmStaged + + epilogue, prewait = STAGED_BUILDS[build] + return Sm100BF16x3RouterGemmStaged.compile(tile, HIDDEN, epilogue, prewait) # Imported at the first compile, as before the move. from .bf16x3_router_gemm import Sm100BF16x3RouterGemm @@ -176,7 +267,8 @@ def __call__( partials = torch.empty( (split_count, rows, EXPERTS), device=hidden_states.device, dtype=torch.float32 ) - _gemm(row_tile(rows))(hidden_states, weight, partials, split_count) + build = _staged if staged_takes(rows) else None + _gemm(row_tile(rows), build)(hidden_states, weight, partials, split_count) output = torch.empty((rows, EXPERTS), device=hidden_states.device, dtype=torch.float32) size = rows * EXPERTS splits_pow2 = triton.next_power_of_2(split_count) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_gemm_staged.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_gemm_staged.py new file mode 100644 index 000000000000..5439c991ee40 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_gemm_staged.py @@ -0,0 +1,428 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The BF16x3 router GEMM of ``ops/moe/router_gemm`` with its partial tile written through shared memory. + +``Sm100BF16x3RouterGemm`` (``cute_dsl_kernels/blackwell/bf16x3_router_gemm``) ends every CTA with +an epilogue in which each of the four epilogue warps writes its accumulator lanes with one 4-byte +store per output element: 128 stores per thread of the ``[split, row, expert]`` partial at a 128-row +tile. Those narrow stores take most of the CTA's life at the decode row counts on R200. This +kernel is the same GEMM with the epilogue +split in two: + +1. every epilogue thread reads its lane of the two accumulators from tensor memory, adds them per + element (the served statements) and stores the FP32 sums into a ``[BN, 128]`` row-major tile in + shared memory (row = the call's row, column = the expert); +2. after a barrier of the epilogue warps, the tile goes to global memory either as 16-byte vector + stores (``EPI_MODE`` ``"v4"``: each warp writes 512 contiguous bytes of one row per instruction) + or as one 512-byte bulk copy per valid row (``"bulk"``, ``cp.async.bulk`` from shared to global, + waited for before the CTA exits). + +Arithmetic: unchanged. The pipeline, the K tiles a CTA owns, the weight split, the MMAs, the +accumulators and the per-element add are the served kernel's; only the path of the finished FP32 +values to the same global addresses changes. So the partials, and the sum kernel's logits, are +bit-identical to ``Sm100BF16x3RouterGemm``'s. ``weight_prologue=True`` also loads the CTA's first +weight tiles before the dependency wait (same bits). + +Memory order and progress (PDL). Reads before the wait touch only the weight. The partials are +written only after the epilogue warps' ``griddepcontrol.wait`` (so a buffer the previous layer's sum +kernel read is written after that kernel completed). The shared tile is written by the generic +proxy and read either by the same warps after a named barrier (``"v4"``) or by the bulk-copy +(async) proxy after ``fence.proxy.async.shared::cta`` and the barrier (``"bulk"``); each issuing +thread waits for its bulk group's reads before the CTA's final barrier, so no shared memory is +released while a copy reads it, and the grid's completion (which the sum kernel's +``griddepcontrol.wait`` observes) includes the copies' writes. No warp waits on a producer that is +itself blocked, so a CTA completes for any number of K tiles. + +Provenance: the feature branch's sm107 target module of the same name +(``muse_spark_open_816b_a42b_nvfp4__sm_107__tep/ops/moe/router_gemm_staged.py``, #21754, 736cfdc7c), +statement for statement; only this paragraph is new. In this target ``ops/moe/router_gemm`` runs it +for the router GEMM calls of ``STAGED_MIN_ROWS..STAGED_MAX_ROWS`` rows under +``router_staged_gemm`` (default off). +""" + +from __future__ import annotations + +from functools import cache + +import cutlass +from cuda.bindings.driver import CUstream +from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync +from cutlass.cutlass_dsl import dsl_user_op + +from tensorrt_llm._torch.cute_dsl_kernels.blackwell import router_tcgen05 as _tcgen05 +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.bf16x3_router_gemm import ( + _decompose_fp32x2_to_3xbf16x2, +) +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.cute_ptx_utils import simple_tma_copy + +#: Epilogue modes: 16-byte vector stores from the shared tile, or one bulk copy per row. +EPI_MODES = ("v4", "bulk") +#: Shared-memory slots for weight tiles when the weight prologue is on (every K tile a CTA owns +#: at up to 384 rows on 212 SMs). +W_SLOTS = 4 + + +@dsl_user_op +def _st_shared_f32(addr: Int32, value: Float32, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [Int32(addr).ir_value(loc=loc, ip=ip), Float32(value).ir_value(loc=loc, ip=ip)], + "st.shared.f32 [$0], $1;", + "r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _copy_shared_to_global_v4(saddr: Int32, gaddr: Int64, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [Int32(saddr).ir_value(loc=loc, ip=ip), Int64(gaddr).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 a, b, c, d;\n\t" + "ld.shared.v4.f32 {a, b, c, d}, [$0];\n\t" + "st.global.v4.f32 [$1], {a, b, c, d};\n\t}", + "r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fence_proxy_async_shared(*, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [], + "fence.proxy.async.shared::cta;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _bulk_shared_to_global_512(saddr: Int32, gaddr: Int64, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [Int64(gaddr).ir_value(loc=loc, ip=ip), Int32(saddr).ir_value(loc=loc, ip=ip)], + "cp.async.bulk.global.shared::cta.bulk_group [$0], [$1], 512;\n\t" + "cp.async.bulk.commit_group;\n\t" + "cp.async.bulk.wait_group.read 0;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +class Sm100BF16x3RouterGemmStaged: + """``Sm100BF16x3RouterGemm`` (same tile, warps, pipeline, arithmetic) with a shared-memory epilogue.""" + + def __init__(self, BN: int = 128, epi_mode: str = "v4", weight_prologue: bool = False) -> None: + if epi_mode not in EPI_MODES: + raise ValueError(f"epi_mode {epi_mode!r} not in {EPI_MODES}") + if BN % 8: + raise ValueError(f"BN {BN} must be a multiple of 8") + self.cta_tile = (BN, 128, 64) + self.num_stages = 2 + self.num_warps = 10 + self.epi_mode = epi_mode + self.weight_prologue = weight_prologue + + @cute.jit + def _make_tma(self, tensor: cute.Tensor, BM: int, BK: int, stages: int): + op = cpasync.CopyBulkTensorTileG2SOp() + swizzle_128B = cute.make_swizzle(3, 4, 3) + elems = 128 * 8 // tensor.element_type.width # 128B + slayout = cute.make_layout( + (BM, (elems, BK // elems), stages), + stride=(elems, (1, BM * elems), BM * BK), + ) + slayout = cute.make_composed_layout(swizzle_128B, 0, slayout) + return cpasync.make_tiled_tma_atom(op, tensor, slayout, (BM, BK)) + + @cute.jit + def __call__( + self, + X: cute.Tensor, + W: cute.Tensor, + out: cute.Tensor, + split_k: Int32, + stream: CUstream, + ): + BN, BM, BK = self.cta_tile + w_stages = W_SLOTS if cutlass.const_expr(self.weight_prologue) else self.num_stages + W_tma = self._make_tma(W, BM, BK, w_stages) + X_tma = self._make_tma(X, BN, BK, self.num_stages) + + grid_m = cute.ceil_div(W.shape[0], BM) + grid_n = cute.ceil_div(X.shape[0], BN) + + self.kernel(X_tma, W_tma, out).launch( + grid=(grid_m, grid_n, split_k), + block=(self.num_warps * 32, 1, 1), + stream=stream, + use_pdl=True, + ) + + @cute.jit + def _load_w( + self, W_tma, gW_tiles, sW, w_full_mbar, w_free_mbar, tile_k, jt, WS: cutlass.Constexpr + ): + BN, BM, BK = self.cta_tile + slot = jt % WS + cute.arch.mbarrier_wait(w_free_mbar + slot, ((jt // WS) & 1) ^ 1) + mbar = w_full_mbar + slot + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar, BM * BK * 4) + simple_tma_copy(W_tma.atom, gW_tiles[None, None, tile_k], sW[None, None, slot], mbar) + + @cute.kernel + def kernel(self, X_tma: cpasync.TmaInfo, W_tma: cpasync.TmaInfo, out: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid_m, bid_n, bid_k = cute.arch.block_idx() + _, _, split_k = cute.arch.grid_dim() + + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + BN, BM, BK = self.cta_tile + STAGES = self.num_stages + PROLOGUE = self.weight_prologue + WS = W_SLOTS if cutlass.const_expr(PROLOGUE) else STAGES + + N, K = X_tma.tma_tensor.shape + M, _ = W_tma.tma_tensor.shape + k_tiles = cute.ceil_div(K, BK) + + smem = cutlass.utils.SmemAllocator() + sX = smem.allocate_tensor( + BFloat16, + X_tma.smem_layout.outer, + byte_alignment=128, + swizzle=X_tma.smem_layout.inner, + ) + sW = smem.allocate_tensor( + Float32, + W_tma.smem_layout.outer, + byte_alignment=128, + swizzle=W_tma.smem_layout.inner, + ) + + w_full_mbar = smem.allocate_array(Int64, WS) + w_free_mbar = smem.allocate_array(Int64, WS) + x_full_mbar = smem.allocate_array(Int64, STAGES) + x_empty_mbar = smem.allocate_array(Int64, STAGES) + d_full_mbar = smem.allocate_array(Int64, STAGES) + d_empty_mbar = smem.allocate_array(Int64, STAGES) + mma_mbar = smem.allocate_array(Int64, 1) + taddr = smem.allocate(Int32, 4) + + BAR_PREP = 2 + BAR_EPI = 3 + BAR_PREP_DONE = 4 + + acc_main = 0 + acc_res = BN + w_tmem_base = BN * 2 + + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(WS): + cute.arch.mbarrier_init(w_full_mbar + i, 1) + cute.arch.mbarrier_init(w_free_mbar + i, 1) + for i in cutlass.range_constexpr(STAGES): + cute.arch.mbarrier_init(x_full_mbar + i, 1) + cute.arch.mbarrier_init(x_empty_mbar + i, 1) + cute.arch.mbarrier_init(d_full_mbar + i, 128) + cute.arch.mbarrier_init(d_empty_mbar + i, 1) + cute.arch.mbarrier_init(mma_mbar, 1) + cute.arch.mbarrier_init_fence() + # All 512 columns, before any warp can write tensor memory. + _tcgen05.alloc(taddr) + elif warp_id == 1: + cpasync.prefetch_descriptor(X_tma.atom) + cpasync.prefetch_descriptor(W_tma.atom) + cute.arch.sync_threads() + + if warp_id == 9: + # TMA warp. + gW_tiles = cute.local_tile(W_tma.tma_tensor, (BM, BK), (bid_m, None)) + gX_tiles = cute.local_tile(X_tma.tma_tensor, (BN, BK), (bid_n, None)) + if cutlass.const_expr(PROLOGUE): + jt = Int32(0) + for tile_k in cutlass.range(bid_k, k_tiles, split_k, unroll=1): + if jt < WS: + mbar = w_full_mbar + jt + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar, BM * BK * 4) + simple_tma_copy( + W_tma.atom, gW_tiles[None, None, tile_k], sW[None, None, jt], mbar + ) + jt = jt + 1 + cute.arch.griddepcontrol_wait() + jt = Int32(0) + for tile_k in cutlass.range(bid_k, k_tiles, split_k, unroll=1): + if cutlass.const_expr(PROLOGUE): + if jt >= WS: + self._load_w(W_tma, gW_tiles, sW, w_full_mbar, w_free_mbar, tile_k, jt, WS) + else: + self._load_w(W_tma, gW_tiles, sW, w_full_mbar, w_free_mbar, tile_k, jt, WS) + stage = jt % STAGES + cute.arch.mbarrier_wait(x_empty_mbar + stage, ((jt // STAGES) & 1) ^ 1) + mbar = x_full_mbar + stage + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar, BN * BK * 2) + simple_tma_copy( + X_tma.atom, gX_tiles[None, None, tile_k], sX[None, None, stage], mbar + ) + jt = jt + 1 + + elif warp_id == 8: + # MMA warp: the served MMAs, in the served order. + idesc = _tcgen05.make_bf16_idesc(BM, BN) + sdesc = _tcgen05.make_sdesc_128B_swizzle(0) + jm = Int32(0) + for tile_k in cutlass.range(bid_k, k_tiles, split_k, unroll=1): + stage = jm % STAGES + parity = (jm // STAGES) & 1 + cute.arch.mbarrier_wait(x_full_mbar + stage, parity) + cute.arch.mbarrier_wait(d_full_mbar + stage, parity) + _tcgen05.fence_after_thread_sync() + w_tmem = w_tmem_base + stage * (BK // 2 * 3) + x_desc = sdesc | (sX[None, None, stage].iterator.toint() >> 4) + for k in cutlass.range_constexpr(BK // 16): + enable_d = (tile_k > bid_k) or (k > 0) + _tcgen05.mma_ts_f16(acc_main, w_tmem + k * 8, x_desc, idesc, enable_d) + _tcgen05.mma_ts_f16(acc_res, w_tmem + 32 + k * 8, x_desc, idesc, enable_d) + _tcgen05.mma_ts_f16(acc_res, w_tmem + 64 + k * 8, x_desc, idesc, True) + x_desc += 32 >> 4 + _tcgen05.commit(x_empty_mbar + stage) + _tcgen05.commit(d_empty_mbar + stage) + jm = jm + 1 + _tcgen05.commit(mma_mbar) + + elif warp_id >= 4: + # Prep warps: split each FP32 weight tile into 3 BF16 terms in tensor memory. + warp_id_ = warp_id % 4 + op = cute.nvgpu.CopyUniversalOp() + cp_atom = cute.make_copy_atom(op, Float32, num_bits_per_copy=128) + row = warp_id_ * 32 + lane_id + sW_view = cute.zipped_divide(sW[row, None, None], (4, 1)) + jp = Int32(0) + for _ in cutlass.range(bid_k, k_tiles, split_k, unroll=1): + slot = jp % WS + stage = jp % STAGES + if warp_id_ == 0: + cute.arch.mbarrier_wait(w_full_mbar + slot, (jp // WS) & 1) + cute.arch.mbarrier_wait(d_empty_mbar + stage, ((jp // STAGES) & 1) ^ 1) + cute.arch.barrier(barrier_id=BAR_PREP, number_of_threads=128) + _tcgen05.fence_after_thread_sync() + w_tmem = w_tmem_base + stage * (BK // 2 * 3) + for kblock in cutlass.range_constexpr(BK // 4): + w0 = cute.make_rmem_tensor(2, Uint32) + w1 = cute.make_rmem_tensor(2, Uint32) + w2 = cute.make_rmem_tensor(2, Uint32) + w_tmp = cute.make_rmem_tensor(4, Float32) + cute.copy(cp_atom, sW_view[None, (kblock, slot)], w_tmp) + w0[0], w1[0], w2[0] = _decompose_fp32x2_to_3xbf16x2(w_tmp[0], w_tmp[1]) + w0[1], w1[1], w2[1] = _decompose_fp32x2_to_3xbf16x2(w_tmp[2], w_tmp[3]) + tcol = kblock * 2 + _tcgen05.st(warp_id_ * 32, w_tmem + 0 + tcol, "32x32b", 2, w0) + _tcgen05.st(warp_id_ * 32, w_tmem + 32 + tcol, "32x32b", 2, w1) + _tcgen05.st(warp_id_ * 32, w_tmem + 64 + tcol, "32x32b", 2, w2) + _tcgen05.wait_st() + _tcgen05.fence_before_thread_sync() + cute.arch.mbarrier_arrive(d_full_mbar + stage) + # Every prep thread has read the slot: it may be refilled. + cute.arch.barrier(barrier_id=BAR_PREP_DONE, number_of_threads=128) + if warp_id_ == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(w_free_mbar + slot) + jp = jp + 1 + + else: + # Epilogue warps. + cute.arch.griddepcontrol_wait() + if warp_id == 0: + cute.arch.mbarrier_wait(mma_mbar, 0) + cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=128) + _tcgen05.fence_after_thread_sync() + cute.arch.griddepcontrol_launch_dependents() + + # The weight slots are free once every MMA completed (the prep warps read every landed + # tile before the MMAs that consumed it): they hold the [BN, 128] FP32 tile (<= 64 KiB). + _fence_proxy_async_shared() + s_base = sW.iterator.toint() + WIDTH = 8 + for i in cutlass.range_constexpr(BN // WIDTH): + tcol = i * WIDTH + main_regs = cute.make_rmem_tensor(WIDTH, Float32) + res_regs = cute.make_rmem_tensor(WIDTH, Float32) + main_regs.store(_tcgen05.ld(warp_id * 32, tcol, "32x32b", WIDTH)) + res_regs.store(_tcgen05.ld(warp_id * 32, BN + tcol, "32x32b", WIDTH)) + _tcgen05.wait_ld() + for j in cutlass.range(WIDTH, vectorize=True): + main_regs[j] += res_regs[j] + for j in cutlass.range_constexpr(WIDTH): + # Shared tile: row = the call's row within the tile, column = the expert lane. + _st_shared_f32(s_base + ((tcol + j) * BM + tid) * 4, main_regs[j]) + + if cutlass.const_expr(self.epi_mode == "bulk"): + _fence_proxy_async_shared() + cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=128) + + g_base = out.iterator.toint() + if cutlass.const_expr(self.epi_mode == "v4"): + # 32 float4 per 128-expert row; 128 threads cover 4 rows per pass. + for p in cutlass.range_constexpr(BN // 4): + r = p * 4 + tid // 32 + c4 = tid % 32 + x_row = bid_n * BN + r + if x_row < N: + g_off = (Int64(bid_k) * N + x_row) * M + bid_m * BM + c4 * 4 + _copy_shared_to_global_v4( + s_base + (r * BM + c4 * 4) * 4, g_base + g_off * 4 + ) + else: + if tid < BN: + x_row = bid_n * BN + tid + if x_row < N: + g_off = (Int64(bid_k) * N + x_row) * M + bid_m * BM + _bulk_shared_to_global_512(s_base + tid * BM * 4, g_base + g_off * 4) + + cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=128) + if warp_id == 0: + _tcgen05.dealloc() + + @cache + @staticmethod + def compile(BN: int = 128, K: int = 8192, epi_mode: str = "v4", weight_prologue: bool = False): + N = cute.sym_int() + M = cute.sym_int() + SPLIT_K = cute.sym_int() + X = cute.runtime.make_fake_compact_tensor( + BFloat16, (N, K), stride_order=(1, 0), assumed_align=16 + ) + W = cute.runtime.make_fake_compact_tensor( + Float32, (M, K), stride_order=(1, 0), assumed_align=16 + ) + out = cute.runtime.make_fake_compact_tensor( + Float32, (SPLIT_K, N, M), stride_order=(2, 1, 0), assumed_align=16 + ) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = Sm100BF16x3RouterGemmStaged(BN, epi_mode, weight_prologue) + return cute.compile(kernel, X, W, out, Int32(1), stream, options="--enable-tvm-ffi") diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_shard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_shard.py index 3d4e7baf83c4..294e9c4b5d02 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_shard.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_shard.py @@ -35,11 +35,14 @@ from __future__ import annotations +from typing import Optional + import torch import triton import triton.language as tl from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger #: The checkpoint's router: 256 experts over the 8192-wide residual. EXPERTS = 256 @@ -543,3 +546,525 @@ def is_valid( router_shard_publish = RouterShardPublish() + + +#: ``router_shard_wide`` (``wide_prologue``): the rows a call may carry, and the fewest. The +#: built-in router takes its tensor-core launch from 16 rows up to 128; the wide prologue +#: serves TP4 decode steps of two to four requests (16..32 rows at K=7). +WIDE_MIN_ROWS = 16 +WIDE_MAX_ROWS = 32 +#: The built-in tensor-core launch's tiles: 16 tokens x 32 experts per program, the +#: 8192-wide dot split into 16 programs of 512, 4 warps (``_router_tensorcore_kernel``). +WIDE_TOKENS_PER_TILE = 16 +WIDE_EXPERTS_PER_TILE = 32 +WIDE_SPLITS = 16 + + +#: ``router_shard_wide`` launches the built-in ``_router_tensorcore_kernel``'s statements over +#: the rank's ``L`` experts: the same tile, loads and ``tl.dot`` (TF32x3) per program, so +#: each program's partial is the built-in program's for the same 32 experts; only the +#: weight rows it is given and the partials' row stride (``L`` instead of 256) differ. +@triton.jit +def _router_shard_wide_kernel( + hidden_states, + weight, + partials, + TOKENS: tl.constexpr, + EXPERTS: tl.constexpr, +): + tokens = tl.program_id(0) * 16 + tl.arange(0, 16) + experts = tl.program_id(1) * 32 + tl.arange(0, 32) + split = tl.program_id(2) + acc = tl.zeros((16, 32), tl.float32) + for block in range(4): + hidden = split * 512 + block * 128 + tl.arange(0, 128) + x = tl.load( + hidden_states + tokens[:, None] * 8192 + hidden[None, :], + tokens[:, None] < TOKENS, + other=0.0, + ).to(tl.float32) + w = tl.load(weight + experts[None, :] * 8192 + hidden[:, None]) + acc = tl.dot(x, w, acc, input_precision="tf32x3") + tl.store( + partials + split * TOKENS * EXPERTS + tokens[:, None] * EXPERTS + experts[None, :], + acc, + tokens[:, None] < TOKENS, + ) + + +#: The built-in ``_router_tensorcore_reduce_kernel``'s statements: the 16 partials of an +#: element summed by one ``tl.sum`` over a ``[16, 128]`` tile at 4 warps, then scaled. +@triton.jit +def _router_shard_wide_reduce_kernel(partials, output, scale, SIZE: tl.constexpr): + offsets = tl.program_id(0) * 128 + tl.arange(0, 128) + splits = tl.arange(0, 16) + values = tl.load( + partials + splits[:, None] * SIZE + offsets[None, :], + offsets[None, :] < SIZE, + other=0.0, + ) + tl.store(output + offsets, tl.sum(values, axis=0) * scale, offsets < SIZE) + + +class RouterShardWide(OpWrapper): + """The rank's ``L`` router logit columns of a 16..32-row step, the built-in tensor-core bits. + + At 16..128 rows the built-in ``MuseSparkOpenMoE.router_logits`` runs + ``_router_tensorcore_kernel`` (TF32x3, ``(cdiv(M, 16), 8, 16)`` programs of 16 + tokens x 32 experts x 512 hidden) and a reduce of its 16 partials, reading the + whole 8 MiB weight on every rank. This entry launches the same two kernels' + statements over this rank's ``L`` experts only (``(cdiv(M, 16), L / 32, 16)`` + programs on ``weight[L r : L r + L]``): every program is the built-in program of + those 32 experts, and every logit's 16 partials are summed by the same reduction, + so each logit is the built-in one bit for bit, from a quarter of the weight at EP4. + + A plain launch (no programmatic dependent launch) into the caller's ``out``, + for a side stream forked from the producer of ``hidden_states`` (``router_fork``). + + Args: + hidden_states: ``[M, 8192]`` BF16, contiguous, ``16 <= M <= 32``. + weight: the full router weight ``[256, 8192]`` FP32, contiguous. + rank: the expert-parallel rank ``r``; columns ``[L r, L r + L)`` are computed. + scale: the metap multiplier, a Python float (applied in FP32 after the sum). + out: ``[M, L]`` FP32, contiguous, on the input's device; overwritten. + ranks: the expert-parallel group size (4 or 8); ``L = local_experts(ranks)``. + + Returns: + ``out``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="tep4: two requests' K=7 verify, 16 rows, 64 columns per rank", + spec=dict(rows=16, ranks=4), + ), + Cell(why="tep4: three requests, 24 rows (two token tiles)", spec=dict(rows=24, ranks=4)), + Cell(why="tep4: four requests, 32 rows (the wide capacity)", spec=dict(rows=32, ranks=4)), + Cell(why="a row count off the K=7 grid (19 rows)", spec=dict(rows=19, ranks=4)), + Cell(why="tep8 layout: 16 rows, 32 columns per rank", spec=dict(rows=16, ranks=8)), + ) + + note = """ + The claim the test proves is equality with the built-in router's columns + ``L r .. L r + L - 1`` at every row count 16..32, for every rank, at EP4 and EP8, + eagerly and through a CUDA-graph replay on a forked side stream; the cells + additionally gate the op against the float64 reference. + """ + + def __call__( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + rank: int, + scale: float, + out: torch.Tensor, + ranks: int = 4, + ) -> torch.Tensor: + tokens = hidden_states.shape[0] + local = local_experts(ranks) + first = rank * local + partials = torch.empty( + (WIDE_SPLITS, tokens, local), device=hidden_states.device, dtype=torch.float32 + ) + _router_shard_wide_kernel[ + (triton.cdiv(tokens, WIDE_TOKENS_PER_TILE), local // WIDE_EXPERTS_PER_TILE, WIDE_SPLITS) + ]( + hidden_states, + weight[first : first + local], + partials, + TOKENS=tokens, + EXPERTS=local, + num_warps=4, + ) + _router_shard_wide_reduce_kernel[(triton.cdiv(tokens * local, 128),)]( + partials, + out, + scale, + SIZE=tokens * local, + num_warps=4, + ) + return out + + def reference( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + rank: int, + scale: float, + out: torch.Tensor, + ranks: int = 4, + ) -> torch.Tensor: + """``RouterShard.reference``: the float64 product, the multiplier on the FP32 sum. + + ``out`` does not change the values.""" + del out + return router_shard.reference(hidden_states, weight, rank, scale, 16, ranks=ranks) + + def is_valid( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + rank: int, + scale: float, + out: torch.Tensor, + ranks: int = 4, + ) -> None: + """Dense row-major operands; 16..32 rows (the built-in's tensor-core launch range).""" + if hidden_states.dtype != torch.bfloat16 or hidden_states.dim() != 2: + raise ValueError(f"hidden_states must be 2-D BF16, got {hidden_states.dtype}") + if not hidden_states.is_contiguous() or hidden_states.shape[1] != HIDDEN: + raise ValueError(f"hidden_states must be contiguous [M, {HIDDEN}]") + rows = hidden_states.shape[0] + if not WIDE_MIN_ROWS <= rows <= WIDE_MAX_ROWS: + raise ValueError(f"{WIDE_MIN_ROWS}..{WIDE_MAX_ROWS} rows, got {rows}") + if weight.dtype != torch.float32 or tuple(weight.shape) != (EXPERTS, HIDDEN): + raise ValueError(f"weight must be FP32 [{EXPERTS}, {HIDDEN}]") + if not weight.is_contiguous(): + raise ValueError("weight must be contiguous") + if ranks not in RANK_COUNTS: + raise ValueError(f"ranks must be one of {RANK_COUNTS}, got {ranks}") + if not 0 <= rank < ranks: + raise ValueError(f"rank must be in [0, {ranks}), got {rank}") + local = local_experts(ranks) + if out.dtype != torch.float32 or tuple(out.shape) != (rows, local): + raise ValueError(f"out must be FP32 [{rows}, {local}]") + if not out.is_contiguous() or out.device != hidden_states.device: + raise ValueError("out must be contiguous, on the input's device") + + +router_shard_wide = RouterShardWide() + + +#: ``router_shard_rows`` (``wide_prologue``) launches ``_router_shard_side_kernel``'s statements +#: with each program looping over ``PER`` consecutive rows of one expert: the weight row is +#: loaded once per program, and every row's logit is the same ``tl.sum`` over the same +#: ``[1, 8192]`` tile layout (the warp count alone fixes it), so each row's bits are those +#: of ``router_shard`` on that row. Placed after every other kernel of this file. +@triton.jit +def _router_shard_rows_kernel( + hidden_states, + weight, + output, + scale, + TOKENS: tl.constexpr, + PER: tl.constexpr, + EXPERTS: tl.constexpr, + HIDDEN: tl.constexpr, + ROWS: tl.constexpr, +): + first = tl.program_id(0) * PER + expert = tl.program_id(1) * ROWS + tl.arange(0, ROWS) + offsets = tl.arange(0, HIDDEN) + w = tl.load( + weight + expert[:, None] * HIDDEN + offsets[None, :], + expert[:, None] < EXPERTS, + other=0.0, + ) + for k in tl.static_range(PER): + token = first + k + if token < TOKENS: + x = tl.load(hidden_states + token * HIDDEN + offsets).to(tl.float32) + logits = tl.sum(w * x[None, :], axis=1) * scale + tl.store(output + token * EXPERTS + expert, logits, expert < EXPERTS) + + +#: Rows a ``router_shard_rows`` call may carry: two to four requests' K=7 verify steps. +ROWS_MAX_ROWS = 32 +#: Rows per program by row count: the fastest of the bitwise-equal launches on one R200 at +#: 16 warps (every entry writes the same bits). +ROWS_PER_PROGRAM = {16: 1, 24: 4, 32: 2} + + +#: ``configure(rows_tiling=...)`` = ``rows:per:experts,...`` (e.g. ``16:2:2,24:4:2,32:4:2``; +#: sm_107 sets ``16:1:2,24:2:2,32:2:2``) sets ``router_shard_rows``' rows per program and experts +#: per program at those row counts; None (or a row count not listed) keeps ``ROWS_PER_PROGRAM`` +#: with one expert per program. In the engine the router runs beside the pool GEMM, whose CTAs +#: hold 128 of the 212 SMs, so fewer, larger programs can finish sooner there than in the +#: isolated launch ``ROWS_PER_PROGRAM`` was chosen on. Neither count is part of a logit's +#: summation order (the ``[experts, 8192]`` tile is a register repetition of the ``[1, 8192]`` +#: row layout, which the warp count alone fixes), so every logit keeps its bits. +ROWS_TILING: dict[int, tuple[int, int]] = {} + + +def _rows_tiling(raw: Optional[str]) -> dict[int, tuple[int, int]]: + """``rows:per:experts,...`` as ``{rows: (per, experts)}``; None or empty: no entry. + + Raises: + ValueError: for a cell outside 1..32 rows, below one row per program, or with an + expert count other than 1, 2, 4 or 8. + """ + out = {} + for cell in (c for c in (raw or "").strip().split(",") if c.strip()): + rows, per, experts = (int(v) for v in cell.split(":")) + if not (1 <= rows <= 32 and per >= 1 and experts in (1, 2, 4, 8)): + raise ValueError(f"[musespark_open] router_shard rows_tiling: bad cell {cell!r}") + out[rows] = (per, experts) + return out + + +def configure(*, rows_tiling: Optional[str] = None) -> None: + """Set the tiling ``router_shard_rows`` launches with (``ROWS_TILING``). + + ``lever_settings.apply`` sets it from the model's plan before any install; the default is + the built-in tiling (``ROWS_PER_PROGRAM``, one expert per program). + + Args: + rows_tiling: ``rows:per:experts,...`` cells (rows per program and experts per program + at a row count), or None for none. + + Raises: + ValueError: for a malformed cell (``_rows_tiling``); the tiling is then unchanged. + """ + global ROWS_TILING + ROWS_TILING = _rows_tiling(rows_tiling) + if ROWS_TILING: + cells = sorted(ROWS_TILING.items()) + logger.info( + "[musespark_open] router_shard_rows: rows per program / experts per program " + + ", ".join(f"{rows} rows {per}/{experts}" for rows, (per, experts) in cells) + ) + + +def rows_per_program(tokens: int) -> int: + """Rows each ``router_shard_rows`` program loops over at ``tokens`` rows (2 off the K=7 grid).""" + if tokens in ROWS_TILING: + return ROWS_TILING[tokens][0] + return ROWS_PER_PROGRAM.get(tokens, 2) + + +def experts_per_program(tokens: int) -> int: + """Expert rows of the weight tile each ``router_shard_rows`` program holds at ``tokens`` rows.""" + return ROWS_TILING[tokens][1] if tokens in ROWS_TILING else TILE_ROWS + + +class RouterShardRows(OpWrapper): + """The rank's ``L`` router logit columns of up to 32 rows, each row's one-request bits. + + ``router_shard_side``'s statements -- the built-in small-batch router's -- over every row + of a 2..4-request decode step: each logit is the one ``router_shard`` computes for that row + at the same warp count, i.e. the bits the built-in router gives the row when its request + decodes alone (``num_warps = builtin_warps(MAX_ROWS, overlaps)``). The built-in router of a + 16..32-row step instead takes its TF32x3 tensor-core launch, which ``router_shard_wide`` + reproduces; this entry trades those bits for the one-request ones in one launch and no + reduce. A plain launch (no programmatic dependent launch) into the caller's ``out``, for + the router fork's side stream. + + With ``enable_pdl`` the same statements run as ``_router_shard_rows_pdl_kernel``, launched + with programmatic dependent launch: the immutable weight tile is loaded before + ``griddepcontrol.wait``, the rows and the stores only after it. On a side stream forked + from the producer of ``hidden_states`` inside a CUDA graph capture, the fork is captured as + a programmatic edge, so the program becomes resident and streams its weight while the + producer finishes (as the pool GEMM beside it does). ``early_trigger`` (with + ``enable_pdl``) also releases this kernel's own dependents at entry; their + ``griddepcontrol.wait`` still returns only after this grid completed. Same bits. + + Args: + hidden_states: ``[M, 8192]`` BF16, contiguous, ``1 <= M <= 32``. + weight: the full router weight ``[256, 8192]`` FP32, contiguous. + rank: the expert-parallel rank ``r``; columns ``[L r, L r + L)`` are computed. + scale: the metap multiplier, a Python float (applied in FP32 after the sum). + num_warps: the warp count of the built-in launch at one request's rows (4 or 16). + out: ``[M, L]`` FP32, contiguous, on the input's device; overwritten. + ranks: the expert-parallel group size (4 or 8); ``L = local_experts(ranks)``. + tiling: ``(rows per program, experts per program)``; None: ``ROWS_TILING`` / + ``ROWS_PER_PROGRAM``. Not part of any logit's summation order. + enable_pdl: launch ``_router_shard_rows_pdl_kernel`` with programmatic dependent launch. + early_trigger: with ``enable_pdl``, trigger the dependents at entry. + + Returns: + ``out``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="tep4: two requests' K=7 verify, 16 rows, 64 columns, the served 16 warps", + spec=dict(rows=16, num_warps=16, ranks=4), + ), + Cell( + why="tep4: three requests, 24 rows (4 rows per program)", + spec=dict(rows=24, num_warps=16, ranks=4), + ), + Cell( + why="tep4: four requests, 32 rows (the wide capacity, 2 rows per program)", + spec=dict(rows=32, num_warps=16, ranks=4), + ), + Cell( + why="a row count off the K=7 grid (19 rows, a partial last program)", + spec=dict(rows=19, num_warps=16, ranks=4), + ), + Cell( + why="without the prepare stream: the 4-warp summation order", + spec=dict(rows=24, num_warps=4, ranks=4), + ), + Cell(why="tep8 layout: 16 rows, 32 columns", spec=dict(rows=16, num_warps=16, ranks=8)), + ) + + note = """ + The claim the test proves is that every row of every rank's output is ``router_shard``'s + logit for that row bit for bit -- the built-in small-batch router's columns + ``L r .. L r + L - 1`` of that row in a one-request (8-row) call -- at every row count + 1..32, both warp counts, EP4 and EP8, eagerly and through a CUDA-graph replay on a forked + side stream; the cells additionally gate the op against the float64 reference. + """ + + def __call__( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + rank: int, + scale: float, + num_warps: int, + out: torch.Tensor, + ranks: int = 4, + tiling: Optional[tuple[int, int]] = None, + enable_pdl: bool = False, + early_trigger: bool = False, + ) -> torch.Tensor: + tokens = hidden_states.shape[0] + local = local_experts(ranks) + first = rank * local + if tiling is None: + per, tile = rows_per_program(tokens), experts_per_program(tokens) + else: + per, tile = (int(v) for v in tiling) + grid = (triton.cdiv(tokens, per), triton.cdiv(local, tile)) + if enable_pdl: + _router_shard_rows_pdl_kernel[grid]( + hidden_states, + weight[first : first + local], + out, + scale, + TOKENS=tokens, + PER=per, + EXPERTS=local, + HIDDEN=HIDDEN, + ROWS=tile, + EARLY_TRIGGER=bool(early_trigger), + num_warps=num_warps, + enable_fp_fusion=False, + launch_pdl=True, + ) + return out + _router_shard_rows_kernel[grid]( + hidden_states, + weight[first : first + local], + out, + scale, + TOKENS=tokens, + PER=per, + EXPERTS=local, + HIDDEN=HIDDEN, + ROWS=tile, + num_warps=num_warps, + enable_fp_fusion=False, + launch_pdl=False, + ) + return out + + def reference( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + rank: int, + scale: float, + num_warps: int, + out: torch.Tensor, + ranks: int = 4, + tiling: Optional[tuple[int, int]] = None, + enable_pdl: bool = False, + early_trigger: bool = False, + ) -> torch.Tensor: + """``RouterShard.reference``; ``num_warps``, ``out``, ``tiling`` and the launch do not change + the values.""" + del out, tiling, enable_pdl, early_trigger + return router_shard.reference(hidden_states, weight, rank, scale, num_warps, ranks=ranks) + + def is_valid( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + rank: int, + scale: float, + num_warps: int, + out: torch.Tensor, + ranks: int = 4, + tiling: Optional[tuple[int, int]] = None, + enable_pdl: bool = False, + early_trigger: bool = False, + ) -> None: + """Dense row-major operands; 1..32 rows; a built-in warp count; a tiling ``_rows_tiling`` + accepts; ``early_trigger`` only with ``enable_pdl``.""" + if tiling is not None: + per, experts = tiling + if not (per >= 1 and experts in (1, 2, 4, 8)): + raise ValueError(f"tiling must be (rows >= 1, experts in 1/2/4/8), got {tiling}") + if early_trigger and not enable_pdl: + raise ValueError("early_trigger needs enable_pdl") + if hidden_states.dtype != torch.bfloat16 or hidden_states.dim() != 2: + raise ValueError(f"hidden_states must be 2-D BF16, got {hidden_states.dtype}") + if not hidden_states.is_contiguous() or hidden_states.shape[1] != HIDDEN: + raise ValueError(f"hidden_states must be contiguous [M, {HIDDEN}]") + rows = hidden_states.shape[0] + if not 1 <= rows <= ROWS_MAX_ROWS: + raise ValueError(f"1..{ROWS_MAX_ROWS} rows, got {rows}") + if weight.dtype != torch.float32 or tuple(weight.shape) != (EXPERTS, HIDDEN): + raise ValueError(f"weight must be FP32 [{EXPERTS}, {HIDDEN}]") + if not weight.is_contiguous(): + raise ValueError("weight must be contiguous") + if ranks not in RANK_COUNTS: + raise ValueError(f"ranks must be one of {RANK_COUNTS}, got {ranks}") + if not 0 <= rank < ranks: + raise ValueError(f"rank must be in [0, {ranks}), got {rank}") + if num_warps not in WARP_COUNTS: + raise ValueError(f"num_warps must be one of {WARP_COUNTS}, got {num_warps}") + local = local_experts(ranks) + if out.dtype != torch.float32 or tuple(out.shape) != (rows, local): + raise ValueError(f"out must be FP32 [{rows}, {local}]") + if not out.is_contiguous() or out.device != hidden_states.device: + raise ValueError("out must be contiguous, on the input's device") + + +router_shard_rows = RouterShardRows() + + +#: ``router_shard_rows(..., enable_pdl=True)`` launches ``_router_shard_rows_kernel``'s statements +#: under a symbol of its own, with programmatic dependent launch: the immutable weight tile before +#: ``griddepcontrol.wait`` (as ``_router_shard_kernel`` with ``PDL``), every row load and every store +#: after it, so each program waits before it exits and every logit is the same ``tl.sum`` over the +#: same tile as ``_router_shard_rows_kernel``'s. With ``EARLY_TRIGGER`` the program first releases +#: its dependents (``griddepcontrol.launch_dependents``). Placed after every other kernel of this +#: file, so their source lines and line info are unchanged. +@triton.jit +def _router_shard_rows_pdl_kernel( + hidden_states, + weight, + output, + scale, + TOKENS: tl.constexpr, + PER: tl.constexpr, + EXPERTS: tl.constexpr, + HIDDEN: tl.constexpr, + ROWS: tl.constexpr, + EARLY_TRIGGER: tl.constexpr, +): + if EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + first = tl.program_id(0) * PER + expert = tl.program_id(1) * ROWS + tl.arange(0, ROWS) + offsets = tl.arange(0, HIDDEN) + w = tl.load( + weight + expert[:, None] * HIDDEN + offsets[None, :], + expert[:, None] < EXPERTS, + other=0.0, + ) + tl.extra.cuda.gdc_wait() + for k in tl.static_range(PER): + token = first + k + if token < TOKENS: + x = tl.load(hidden_states + token * HIDDEN + offsets).to(tl.float32) + logits = tl.sum(w * x[None, :], axis=1) * scale + tl.store(output + token * EXPERTS + expert, logits, expert < EXPERTS) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_step_prep.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_step_prep.py new file mode 100644 index 000000000000..dffdbac5eeac --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_step_prep.py @@ -0,0 +1,504 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode step's drafter bookkeeping on a paged context, in four launches (``drafter_fused_prep``). + +On a decode step the ``context_kv_fork`` side chain decodes the step's draft page table +(``DFlashWorker._refresh_ctx_block_tables``), stores the ``K + 1`` projected target rows of +every request as the drafter's context K/V (``DFlashWorker._store_step_context``: the +columns and positions, ``precompute_context_kv``, then one flashinfer paged append per +drafter layer, each behind two ``.contiguous()`` copies), and its commit advances the +context lengths and reads the lengths the drafter attends over +(``DFlashWorker._commit_context``: ``_advance_ctx_len``, ``_empty_dummy_slot``, +``_drafter_context_lengths`` with its allocation check). Over the KV cache manager's pages +that is about fifty torch launches of one to eight CTAs each, all on the step's critical +path between the target's LM head and the drafter's first layer. + +This module runs the integer and copy parts of that chain as four Triton launches, bit for +bit the built-in's bytes: + +* ``refresh_block_tables``: a fill of the page counts, then one launch over + ``(requests, column tiles)`` that decodes the table and takes each row's allocation extent + (last valid entry + 1) with an atomic max; +* ``context_positions``: the RoPE positions ``ctx_len[slot] + r`` of the stored rows (what + ``precompute_context_kv`` reads); +* ``store_paged_context``: every layer's K and V rows of every stored token into the page + ``table[row, col // page]`` at ``col % page``, with the built-in's column clamp, in one + launch (``precompute_context_kv`` itself, the fused K/V GEMM, its norm and RoPE, is the + built-in's and runs unchanged between the two launches above); +* ``commit_context``: the length advance, the clamp of every slot to ``max_ctx``, the dummy + slot's reset, and the advertised lengths with the allocation check, in one single-CTA + launch: NVIDIA/TensorRT-LLM #21535's ``context_kv_commit_paged`` (aaba1dea on + ``feat/musespark-open-agentx``), which this branch now carries as + ``ops/spec/context_kv_commit_paged.py`` -- this module calls it, it is not backported here. + +The kernels' arithmetic is integer (index math, clamps, extents) and copies; nothing is +rounded. Where the built-in's own result is unspecified (two stored rows of one launch +clamped onto one cache column: a request past its allocation), the fused store is as +unspecified as the built-in's paged append. + +The allocation check (``DFlashWorker._check_block_within_allocation``, a +``torch._assert_async`` that traps on the device) is #21535's ``trap`` in the commit launch on +the same condition: both end the process's CUDA context with a device-side error; the +built-in's message is not printed. + +Upstream: #21535 (merged into ``feat/musespark-open-agentx`` as aaba1dea) fuses the commit +only; this module calls its op, and ``context_kv_fork`` reaches the same op without this +module through its own ``paged_context_commit`` branch (``_commits_paged``), so +``commit_context`` here adds nothing but a second gate. The refresh, the columns and the store have no upstream form +on that branch. NVIDIA/TensorRT-LLM PR #20550 (open, against main: "DFlash / MLA DSpark: fuse +the decode-step drafter bookkeeping into Triton kernels") fuses a block-table refresh and a +single-latent store on main's ``dflash.py``; this image's refresh takes the extent of the valid +entries rather than their count and subtracts the per-layer table offset, and its drafter is +GQA (K and V per KV head), which #20550's fused store leaves on the built-in path, so the +refresh and store kernels here are written for this image's methods (#20550's per-row +structure, with a column-tiled grid and an atomic max for the extent). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm.logger import logger + +from .context_kv_commit_paged import context_kv_commit_paged + +#: ``drafter_fused_prep`` runs the fused bookkeeping on the ``context_kv_fork`` side chain of +#: decode steps of ``MIN_REQUESTS`` .. ``context_kv_commit.MAX_REQUESTS`` requests; off, the +#: built-in methods stand. Resolved once, where ``context_kv_fork`` installs. +#: The fewest generation requests of a step the fused bookkeeping takes (default 2: the +#: multi-request decode steps of 16 rows and more at K = 7); ``1`` adds the one-request step. +DEFAULT_MIN_REQUESTS = 2 + +#: Drafter layers the store's layer offsets are passed for (scalar arguments, no device table). +MAX_LAYERS = 8 +#: Columns of the page table one refresh program decodes. +REFRESH_BLOCK = 1024 +#: Stored tokens one positions program writes. +POSITIONS_BLOCK = 256 + + +# -------------------------------------------------------------------------- +# Kernels +# -------------------------------------------------------------------------- +@triton.jit +def _refresh_kernel( + src_ptr, + src_row, + src_col, + table_ptr, + table_row, + counts_ptr, + width, + offset, + divisor, + BLOCK: tl.constexpr, +): + """``encoded = src - offset``; ``counts[r] = max((encoded >= 0) * (c + 1))``; + ``table[r, c] = int32(max(encoded, 0) // divisor)``. ``counts[r]`` must be 0 on entry.""" + r = tl.program_id(0).to(tl.int64) + c = tl.program_id(1).to(tl.int64) * BLOCK + tl.arange(0, BLOCK).to(tl.int64) + inside = c < width + encoded = tl.load(src_ptr + r * src_row + c * src_col, mask=inside, other=0).to(tl.int64) + encoded = encoded - offset + valid = inside & (encoded >= 0) + extent = tl.max(tl.where(valid, c + 1, 0), axis=0) + tl.atomic_max(counts_ptr + r, extent) + decoded = tl.maximum(encoded, 0) // divisor + tl.store(table_ptr + r * table_row + c, decoded.to(tl.int32), mask=inside) + + +@triton.jit +def _columns_kernel( + ctx_len_ptr, + base_ptr, + slots_ptr, + rows_ptr, + counts_ptr, + pos_ptr, + col_ptr, + tokens, + page_size, + PER_REQUEST: tl.constexpr, + HAS_BASE: tl.constexpr, + BLOCK: tl.constexpr, +): + """Token ``t`` (request ``i = t // PER_REQUEST``, row ``r``): ``pos[t] = ctx_len[slots[i]] + r`` + and ``col[t] = max(min(ctx_len[slots[i]] - base[slots[i]] + r, counts[rows[i]] * page - 1), 0)``.""" + t = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + inside = t < tokens + request = t // PER_REQUEST + r = (t % PER_REQUEST).to(tl.int64) + slot = tl.load(slots_ptr + request, mask=inside, other=0) + row = tl.load(rows_ptr + request, mask=inside, other=0) + length = tl.load(ctx_len_ptr + slot, mask=inside, other=0) + tl.store(pos_ptr + t, length + r, mask=inside) + col = length + r + if HAS_BASE: + col = length - tl.load(base_ptr + slot, mask=inside, other=0) + r + capacity = tl.load(counts_ptr + row, mask=inside, other=0) * page_size + tl.store(col_ptr + t, tl.maximum(tl.minimum(col, capacity - 1), 0), mask=inside) + + +@triton.jit +def _layer_offset( + layer, o0, o1, o2, o3, o4, o5, o6, o7 +): + offset = o0 + offset = tl.where(layer == 1, o1, offset) + offset = tl.where(layer == 2, o2, offset) + offset = tl.where(layer == 3, o3, offset) + offset = tl.where(layer == 4, o4, offset) + offset = tl.where(layer == 5, o5, offset) + offset = tl.where(layer == 6, o6, offset) + offset = tl.where(layer == 7, o7, offset) + return offset + + +@triton.jit +def _store_kernel( + k_ptr, + v_ptr, + k_token, + k_layer, + k_head, + v_token, + v_layer, + v_head, + col_ptr, + rows_ptr, + table_ptr, + table_row, + pool_ptr, + pool_page, + pool_kv, + pool_head, + pool_entry, + page_size, + o0, + o1, + o2, + o3, + o4, + o5, + o6, + o7, + PER_REQUEST: tl.constexpr, + HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEADS_P2: tl.constexpr, + DIM_P2: tl.constexpr, +): + """Token ``t`` (request ``t // PER_REQUEST``), layer ``l``: page ``table[rows[request], col[t] // page]``; + ``pool_l[page, 0 / 1, h, col[t] % page, :] = k / v[t, l, h, :]`` (flashinfer's HND paged append).""" + t = tl.program_id(0).to(tl.int64) + layer = tl.program_id(1) + row = tl.load(rows_ptr + t // PER_REQUEST).to(tl.int64) + col = tl.load(col_ptr + t) + page = tl.load(table_ptr + row * table_row + col // page_size).to(tl.int64) + entry = col % page_size + dest = pool_ptr + _layer_offset(layer, o0, o1, o2, o3, o4, o5, o6, o7) + dest = dest + page * pool_page + entry * pool_entry + h = tl.arange(0, HEADS_P2)[:, None] + d = tl.arange(0, DIM_P2)[None, :] + inside = (h < HEADS) & (d < HEAD_DIM) + k = tl.load(k_ptr + t * k_token + layer * k_layer + h * k_head + d, mask=inside) + v = tl.load(v_ptr + t * v_token + layer * v_layer + h * v_head + d, mask=inside) + tl.store(dest + h * pool_head + d, k, mask=inside) + tl.store(dest + pool_kv + h * pool_head + d, v, mask=inside) + + +# -------------------------------------------------------------------------- +# The step's bookkeeping +# -------------------------------------------------------------------------- +def _p2(n: int) -> int: + return triton.next_power_of_2(max(int(n), 1)) + + +def _pool_geometry(pool) -> Optional[tuple]: + """``(layer 0 view, per-layer element offsets, page / kv / head / entry strides, heads, head dim)`` + of a list of per-layer ``[pages, 2, heads, page, head_dim]`` views, or None. + + Every layer must share the dtype, the shape past the page dimension and the strides, with + unit stride along the head dimension, at a whole number of elements from layer 0 (V2 wraps + each layer's base address in a view of its own; the offsets are the distances between them). + """ + if not isinstance(pool, (list, tuple)) or not 0 < len(pool) <= MAX_LAYERS: + return None + first = pool[0] + if not isinstance(first, torch.Tensor) or first.dim() != 5 or first.shape[1] != 2: + return None + size = first.element_size() + offsets = [] + for view in pool: + if ( + not isinstance(view, torch.Tensor) + or view.dtype != first.dtype + or view.device != first.device + or tuple(view.shape[1:]) != tuple(first.shape[1:]) + or tuple(view.stride()) != tuple(first.stride()) + or view.stride(4) != 1 + or (view.data_ptr() - first.data_ptr()) % size + ): + return None + offsets.append((view.data_ptr() - first.data_ptr()) // size) + offsets += [0] * (MAX_LAYERS - len(offsets)) + return ( + first, + tuple(offsets), + first.stride(0), + first.stride(1), + first.stride(2), + first.stride(3), + first.shape[2], + first.shape[4], + ) + + +class DrafterStepPrep: + """The fused bookkeeping of one DFlash worker's decode steps (``context_kv_fork``'s side chain).""" + + def __init__(self, min_requests: int, max_requests: int) -> None: + self.min_requests = min_requests + self.max_requests = max_requests + #: Steps launched fused, and the first step logged. + self.steps = 0 + self._geometry_key = None + self._geometry = None + self._declined_logged = False + + # -- what the fused path takes --------------------------------------------- + def geometry(self, worker) -> Optional[tuple]: + """The worker's paged context pool geometry (``_pool_geometry``), cached per pool.""" + pool = getattr(worker, "_ctx_kv_buf", None) + if not isinstance(pool, (list, tuple)): + return None + key = tuple((v.data_ptr(), tuple(v.shape), tuple(v.stride()), v.dtype) for v in pool) + if key != self._geometry_key: + self._geometry_key = key + self._geometry = _pool_geometry(pool) + return self._geometry + + def unsupported_reason(self, worker, num_gens: int, per_request: int) -> Optional[str]: + """Why a decode step of ``num_gens`` requests keeps the built-in bookkeeping, or None. + + Host-side reads only (per-capture constants under a CUDA graph), the same on every rank. + """ + if not self.min_requests <= num_gens <= self.max_requests: + return f"{num_gens} requests, outside {self.min_requests}..{self.max_requests}" + if not getattr(worker, "_ctx_paged", False): + return "the drafter context is not paged" + tables = getattr(worker, "_ctx_block_tables", None) + counts = getattr(worker, "_ctx_block_counts", None) + if tables is None or counts is None: + return "the drafter context is not on the KV cache manager's pages" + if tables.dtype != torch.int32 or tables.dim() != 2 or tables.stride(1) != 1: + return "the page table is not a row-major int32 table" + if counts.dtype != torch.int64 or num_gens > counts.numel() or num_gens > tables.shape[0]: + return "the page counts are not int64 per request" + for name in ("_ctx_len", "_batch_to_slot"): + tensor = getattr(worker, name, None) + if tensor is None or tensor.dtype != torch.int64 or not tensor.is_contiguous(): + return f"{name} is not a contiguous int64 tensor" + base = getattr(worker, "_ctx_base", None) + if base is not None and (base.dtype != torch.int64 or not base.is_contiguous()): + return "_ctx_base is not a contiguous int64 tensor" + if per_request != getattr(worker, "_draft_tokens_per_req", None): + return f"{per_request} stored rows per request, not the worker's K + 1" + if self.geometry(worker) is None: + return "the context pool's layers do not share one geometry" + return None + + def takes(self, worker, num_gens: int, per_request: int) -> bool: + reason = self.unsupported_reason(worker, num_gens, per_request) + if reason is not None and num_gens >= self.min_requests and not self._declined_logged: + self._declined_logged = True + logger.warning( + f"[musespark_open] drafter_fused_prep: a decode step keeps the built-in drafter " + f"bookkeeping: {reason} (logged once)" + ) + return reason is None + + # -- the side chain -------------------------------------------------------- + def refresh_block_tables(self, worker, attn_metadata, num_seqs: int) -> bool: + """``DFlashWorker._refresh_ctx_block_tables``, in two launches; the built-in when the + step's offsets are missing or shaped otherwise (it raises or decodes as it does).""" + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + tables = worker._ctx_block_tables + offsets_name = ( + "kv_cache_block_offsets" if worker._ctx_unified_cache else "draft_kv_cache_block_offsets" + ) + src = getattr(attn_metadata, offsets_name, None) + if src is None or num_seqs <= 0 or src.dim() != 4: + return DFlashWorker._refresh_ctx_block_tables(worker, attn_metadata, num_seqs) + view = src[worker._ctx_pool_idx, :num_seqs, 0] + width = view.shape[-1] + if ( + width != tables.shape[1] + or width > worker._ctx_block_cols.numel() + or view.dtype not in (torch.int32, torch.int64) + ): + return DFlashWorker._refresh_ctx_block_tables(worker, attn_metadata, num_seqs) + counts = worker._ctx_block_counts + counts[:num_seqs].zero_() + _refresh_kernel[(num_seqs, triton.cdiv(width, REFRESH_BLOCK))]( + view, + view.stride(0), + view.stride(1), + tables, + tables.stride(0), + counts, + width, + int(worker._ctx_block_offset), + int(worker._ctx_block_divisor), + BLOCK=REFRESH_BLOCK, + num_warps=4, + ) + return True + + def store_step_context(self, worker, draft_model, gen_hs, slots, rows, per_request: int) -> None: + """``DFlashWorker._store_step_context`` without the mask (the fork's side chain) over the + manager's pages: the positions and columns in one launch (from one read of the lengths, + as the built-in's ``ctx_len``), the built-in projection and ``precompute_context_kv``, + then every layer's K/V append in one launch.""" + requests = slots.numel() + tokens = requests * per_request + positions = torch.empty(tokens, dtype=torch.long, device=slots.device) + columns = torch.empty(tokens, dtype=torch.long, device=slots.device) + base = worker._ctx_base + _columns_kernel[(triton.cdiv(tokens, POSITIONS_BLOCK),)]( + worker._ctx_len, + worker._ctx_len if base is None else base, + slots, + rows, + worker._ctx_block_counts, + positions, + columns, + tokens, + int(worker._ctx_page_size), + PER_REQUEST=per_request, + HAS_BASE=base is not None, + BLOCK=POSITIONS_BLOCK, + num_warps=4, + ) + projected = draft_model.project_target_hidden(gen_hs.reshape(-1, gen_hs.shape[-1])) + cache_dtype = worker._ctx_kv_buf[0].dtype + k_new, v_new = draft_model.precompute_context_kv(projected.to(cache_dtype), positions) + geometry = self.geometry(worker) + if ( + v_new is None + or geometry is None + or k_new.dim() != 4 + or tuple(k_new.shape) != tuple(v_new.shape) + or k_new.shape[0] != tokens + or k_new.shape[1] != len(worker._ctx_kv_buf) + or (k_new.shape[2], k_new.shape[3]) != (geometry[6], geometry[7]) + or k_new.dtype != geometry[0].dtype + or v_new.dtype != geometry[0].dtype + or k_new.stride(3) != 1 + or v_new.stride(3) != 1 + ): + # Any other K/V layout: the built-in paged append, on the same columns. + rows_long = rows.unsqueeze(1).expand(requests, per_request).reshape(-1) + worker._store_context_kv_paged(k_new, v_new, rows_long, columns) + return + first, offsets, s_page, s_kv, s_head, s_entry, heads, head_dim = geometry + _store_kernel[(tokens, k_new.shape[1])]( + k_new, + v_new, + k_new.stride(0), + k_new.stride(1), + k_new.stride(2), + v_new.stride(0), + v_new.stride(1), + v_new.stride(2), + columns, + rows, + worker._ctx_block_tables, + worker._ctx_block_tables.stride(0), + first, + s_page, + s_kv, + s_head, + s_entry, + int(worker._ctx_page_size), + *offsets, + PER_REQUEST=per_request, + HEADS=heads, + HEAD_DIM=head_dim, + HEADS_P2=_p2(heads), + DIM_P2=_p2(head_dim), + num_warps=1 if heads * head_dim <= 256 else 4, + ) + + # -- the commit ------------------------------------------------------------ + def commit_context(self, worker, slots, num_accepted, rows, out) -> None: + """``out.copy_(DFlashWorker._commit_context(slots, num_accepted.long(), rows))`` over the + manager's pages, in one launch: NVIDIA/TensorRT-LLM #21535's ``context_kv_commit_paged`` + (``ops/spec/context_kv_commit_paged.py``, backported verbatim; fd07 ships it). The built-in + for an accepted-count layout that op does not take.""" + requests = slots.numel() + if requests == 0: + return + if ( + num_accepted.dtype not in (torch.int32, torch.int64) + or num_accepted.dim() != 1 + or num_accepted.numel() != requests + or not num_accepted.is_contiguous() + or not slots.is_contiguous() + or not rows.is_contiguous() + or worker._dummy_slot is None + ): + out.copy_(worker._commit_context(slots, num_accepted.long(), rows)) + return + context_kv_commit_paged( + worker._ctx_len, + worker._ctx_base, + slots, + num_accepted, + rows, + worker._ctx_block_counts if worker._ctx_block_tables is not None else None, + worker._ctx_page_size, + worker._compute_block_size, + worker._max_ctx, + worker._dummy_slot, + worker._ctx_unified_cache, + out, + ) + self.steps += 1 + if self.steps == 1: + logger.info( + f"[musespark_open] drafter_fused_prep: the first fused decode step ran ({requests} " + "requests): page-table refresh, context positions and columns, every layer's context K/V " + "store and the commit as Triton launches" + ) + + +def from_plan(on: bool, min_requests: int = DEFAULT_MIN_REQUESTS) -> Optional[DrafterStepPrep]: + """The fused bookkeeping when ``drafter_fused_prep`` is on, else None. + + ``min_requests`` is the lever's ``MIN_REQUESTS``; the plan resolved it on every rank alike. + + Raises: + ValueError: for a ``min_requests`` outside ``1..context_kv_commit.MAX_REQUESTS``. + """ + if not on: + return None + from .context_kv_commit import MAX_REQUESTS + + min_requests = int(min_requests) + if not 1 <= min_requests <= MAX_REQUESTS: + raise ValueError( + f"[musespark_open] drafter_fused_prep MIN_REQUESTS {min_requests} is outside " + f"1..{MAX_REQUESTS}" + ) + logger.info( + f"[musespark_open] drafter_fused_prep: decode steps of {min_requests}..{MAX_REQUESTS} " + "requests run the drafter's context-K/V bookkeeping (page-table refresh, store, commit) " + "as fused Triton launches" + ) + return DrafterStepPrep(min_requests, MAX_REQUESTS) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/spec_worker.py index 664ef5cec47a..bacc75f35f8a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/spec_worker.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/spec_worker.py @@ -104,6 +104,43 @@ class still runs those built-ins (no subclass overrides them) and that no #: (prefill chunks, batched decode) run the built-in. TP4 has no row bound. FC_SHARD_MAX_ROWS = 8 +#: Tokens the TP8 projection gather's workspace holds (``musespark_open_projection_gather``). +FC_SHARD_TP8_ROW_CAPACITY = 128 +#: ``tep8_fc_shard_wide``: the TP8 row bound this module installs, ``FC_SHARD_MAX_ROWS`` +#: until ``configure`` raises it. Above the default, decode steps of two to five requests +#: (16..40 rows at K = 7) shard ``fc`` too: per step the rank reads its 84 MB slice instead +#: of the 671 MB weight. +_FC_SHARD_TP8_MAX_ROWS = FC_SHARD_MAX_ROWS + + +def configure(*, fc_shard_tp8_max_rows: int = FC_SHARD_MAX_ROWS) -> None: + """Set the TP8 row bound of the sharded ``fc``. + + ``lever_settings.apply`` sets it from the model's plan before any install; the default + is the stock bound, where only the batch-1 decode step (K + 1 <= 8 rows) shards. + + Args: + fc_shard_tp8_max_rows: the most rows a TP8 call may have and still shard + (``_FC_SHARD_TP8_MAX_ROWS``), ``FC_SHARD_MAX_ROWS``..``FC_SHARD_TP8_ROW_CAPACITY``. + + Raises: + ValueError: for a bound outside that range. + """ + global _FC_SHARD_TP8_MAX_ROWS + rows = int(fc_shard_tp8_max_rows) + if not FC_SHARD_MAX_ROWS <= rows <= FC_SHARD_TP8_ROW_CAPACITY: + raise ValueError( + f"[musespark_open] drafter_fc_shard TP8 row bound {rows}: expected " + f"{FC_SHARD_MAX_ROWS}..{FC_SHARD_TP8_ROW_CAPACITY}" + ) + _FC_SHARD_TP8_MAX_ROWS = rows + + +def fc_shard_tp8_max_rows() -> int: + """The TP8 row bound ``configure`` set, else ``FC_SHARD_MAX_ROWS``.""" + return _FC_SHARD_TP8_MAX_ROWS + + #: The TP widths the sharded ``fc`` is written for: 8 gathers over the fabric #: projection-gather workspace, 4 with the built-in NCCL all-gather. FC_SHARD_WORLDS = (4, 8) @@ -295,8 +332,15 @@ def install_sharded_fc( gather=gather, enable_pdl=enable_pdl, world=world, - max_rows=FC_SHARD_MAX_ROWS if world == 8 else None, + max_rows=fc_shard_tp8_max_rows() if world == 8 else None, ) + if world == 8 and draft_model._fc_shard.max_rows > FC_SHARD_MAX_ROWS: + from tensorrt_llm.logger import logger + + logger.info( + f"[musespark_open] tep8_fc_shard_wide: the TP8 sharded fc takes calls of up to " + f"{draft_model._fc_shard.max_rows} rows" + ) draft_model.project_target_hidden = types.MethodType(sharded_project_target_hidden, draft_model) return None diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py index 561f88b1e4ac..d8ef281ddb87 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py @@ -71,7 +71,7 @@ def install_loaded_features( whichever expert call the paths above installed). """ if plan.enabled("replicated_experts"): - _install_ring_buddy(text_model, mapping, plan) + _install_ring_buddy(text_model, mapping, plan, compute_capability) elif plan.enabled("persistent_experts") or plan.enabled("bf16_routed_experts"): _install_owner_window(text_model, mapping, plan) if plan.enabled("hot_replicas"): @@ -88,14 +88,16 @@ def install_loaded_features( "[musespark_open] pool_into_route is planned but the model shell built no direct pool " "transport" ) - pool_direct.install_sites(text_model, mapping) + pool_direct.install_sites(text_model, mapping, wide=plan.enabled("tep8_wide_direct_pool")) if plan.enabled("next_qkv_prefetch"): from .kernels.next_qkv_prefetch import prefetch as _qkv_prefetch _qkv_prefetch.wrap_expert_forwards(text_model) -def _install_ring_buddy(text_model, mapping, plan: "_dispatch.Plan") -> None: +def _install_ring_buddy( + text_model, mapping, plan: "_dispatch.Plan", compute_capability=None +) -> None: """``install_loaded_features``'s ``replicated_experts`` install.""" from .kernels.replicated_experts import ring_buddy as _ring_buddy @@ -104,6 +106,26 @@ def _install_ring_buddy(text_model, mapping, plan: "_dispatch.Plan") -> None: extra["windows"] = 3 if plan.enabled("bf16_routed_experts"): extra["bf16_unpacked"] = True + max_rows = _ring_buddy_row_cap(_ring_buddy, compute_capability, mapping) + batch_rows = _tep8_persistent_batch_rows( + compute_capability, + mapping, + plan.enabled("persistent_experts"), + max_rows, + plan.param("tep8_persistent_batch_experts", "MAX_ROWS") + if plan.enabled("tep8_persistent_batch_experts") + else 0, + ) + if batch_rows: + extra["persistent_rows"] = batch_rows + extra["persistent_batch"] = True + if _tep8_experts_switch( + "tep8_experts_early_start", + plan.enabled("tep8_experts_early_start"), + batch_rows, + EARLY_START_WHAT, + ): + extra["persistent_early_start"] = True _ring_buddy.install( text_model, mapping, @@ -113,6 +135,116 @@ def _install_ring_buddy(text_model, mapping, plan: "_dispatch.Plan") -> None: ) +def _ring_buddy_row_cap(ring_buddy, compute_capability, mapping) -> int: + """The largest stacked call the balanced ring-buddy path takes on this device. + + ``max_rows_for`` is the cap ``ring_buddy.install`` will itself apply, from + ``tep8_ring_buddy_wide`` (else the module constant ``MAX_ROWS``, 16, so the batch builds + are clamped to 16 rows here and only the 16-column build is reachable). It takes the ring + size, which is what ``install`` passes it, not the compute capability: a tep8 mapping + gives 8, a tep4 one 4, anything else None. Older trees carry neither function, and the cap + is then ``MAX_ROWS``. + """ + resolve = getattr(ring_buddy, "max_rows_for", None) + ring_size = getattr(ring_buddy, "ring_size", None) + if resolve is None or ring_size is None: + return int(ring_buddy.MAX_ROWS) + return int(resolve(ring_size(mapping))) + + +def _tep8_persistent_batch_rows( + compute_capability, mapping, persistent: bool, max_rows: int, rows: int +) -> int: + """``tep8_persistent_batch_experts``' row bound for this install, or 0 (off). + + ``rows`` is the lever's ``MAX_ROWS``: the largest stacked NVFP4 expert call (9..32 rows: + decode steps of two to four requests at K = 7) the persistent expert kernel serves, on its + batch builds (``ops/moe/persistent_experts.BATCH_NTOKS``), instead of TRTLLM-gen's + pre-routed call. Off (0), TRTLLM-gen serves every stacked call above 8 rows. + + Resolved on every rank alike. It takes effect on sm103 at tp 8 with ``persistent_experts`` + planned; elsewhere it warns and does nothing. Calls above the ring-buddy cap (``max_rows``) + never reach the stacked path, so the bound is clamped to it, with a warning. + + Raises: + ValueError: a ``rows`` that is not in ``9 .. BATCH_MAX_ROWS``. + """ + if not rows: + return 0 + from tensorrt_llm.logger import logger + + from .ops.moe import persistent_experts as _persistent + + rows = int(rows) + if not _persistent.MAX_ROWS < rows <= _persistent.BATCH_MAX_ROWS: + raise ValueError( + f"[musespark_open] tep8_persistent_batch_experts MAX_ROWS {rows}: the persistent " + f"expert kernel's batch builds serve stacked calls of {_persistent.MAX_ROWS + 1}.." + f"{_persistent.BATCH_MAX_ROWS} rows; give the largest such call" + ) + tp = getattr(mapping, "tp_size", None) + if _tp_agreement.sm_of(compute_capability) != 103 or tp != 8 or not persistent: + logger.warning( + "[musespark_open] tep8_persistent_batch_experts has no effect here: it serves " + "sm103 tep8 with persistent_experts planned (got sm " + f"{_tp_agreement.sm_of(compute_capability)}, tp {tp}, " + f"persistent_experts {persistent})" + ) + return 0 + if rows > max_rows: + logger.warning( + f"[musespark_open] tep8_persistent_batch_experts MAX_ROWS {rows}: the ring-buddy " + f"stacked path takes calls of at most {max_rows} rows; the batch builds serve " + f"9..{max_rows}" + ) + rows = max_rows + if rows <= _persistent.MAX_ROWS: + return 0 + builds = [n for n in _persistent.BATCH_NTOKS if n < rows] + [_persistent.batch_ntok(rows)] + logger.info( + f"[musespark_open] tep8_persistent_batch_experts: the NVFP4 " + f"layers' stacked calls of {_persistent.MAX_ROWS + 1}..{rows} rows on the 64-expert " + f"ring-buddy window run the persistent expert kernel's batch builds " + f"({', '.join(f'{n} columns' for n in sorted(set(builds)))}) instead of TRTLLM-gen" + ) + return rows + + +#: sm103 tep8, default off; acts only with ``tep8_persistent_batch_experts``: a 9..N-row +#: stacked NVFP4 call launches its ``ring_buddy_assign`` with programmatic dependent launch and +#: an early trigger (after its own dependency wait), and the persistent kernel's batch build +#: behind it launches while the assignment runs and reads its epoch before its own wait. +#: Same rows. +EARLY_START_WHAT = ( + "ring_buddy_assign triggers the persistent expert kernel's batch builds right after its own " + "dependency wait, and they read their epoch before theirs" +) + + +def _tep8_experts_switch(name: str, on: bool, batch_rows: int, what: str) -> bool: + """A lever of the persistent kernel's batch builds, as the plan resolved it. + + Off: off. On where the batch builds serve (``batch_rows``, from + ``tep8_persistent_batch_experts``), with one INFO line; elsewhere a warning and off. + """ + if not on: + return False + from tensorrt_llm.logger import logger + + if not batch_rows: + logger.warning( + f"[musespark_open] {name} has no effect here: it acts on the persistent expert " + "kernel's batch builds, which serve only with tep8_persistent_batch_experts on " + "sm103 tep8 with persistent_experts planned" + ) + return False + logger.info( + f"[musespark_open] {name}: on the NVFP4 layers' stacked calls of " + f"9..{batch_rows} rows, {what}" + ) + return True + + def _install_owner_window(text_model, mapping, plan: "_dispatch.Plan") -> None: """``persistent_experts`` / ``bf16_routed_experts`` without ``replicated_experts``. diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index fe23ada9b761..c525df08e617 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -62,6 +62,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantMode +from . import multi_instance_context as _multi_instance_context from .interface import FmhaPhase from .phased import FmhaParams, PhasedFmha from .utils import ( @@ -1126,29 +1127,48 @@ def run_context( if params.is_cross else AttentionMaskType(fwd.mask_type) == AttentionMaskType.causal ) - flashinfer.prefill.trtllm_batch_context_with_kv_cache( - query=q_processed, - kv_cache=(kv_pool, kv_pool), - workspace_buffer=fmha_workspace, - block_tables=block_tables, - seq_lens=params.sequence_lengths, - max_q_len=max_q_len, - max_kv_len=max_kv_len, - bmm1_scale=ctx_bmm1_scale, - bmm2_scale=ctx_bmm2_scale, - batch_size=params.batch_size, - cum_seq_lens_q=cu_q_seqlens, - cum_seq_lens_kv=cu_kv_seqlens, - window_left=window_left, - out=params.output, - kv_layout=self._layout, - enable_pdl=self._enable_pdl, - sinks=fwd.attention_sinks, - kv_cache_sf=(kv_scale_pool, kv_scale_pool) if kv_scale_pool is not None else None, - uses_shared_paged_kv_idx=self.USE_SHARED_PAGED_KV_IDX, - causal=causal, - multi_ctas_kv_counter_buffer=self._get_multi_ctas_kv_counter_buffer(), + # TRTLLM_MULTI_INSTANCE_CONTEXT_FMHA=1 (off by default): a full-attention context call the + # multi-instance context FMHA serves runs there; every other call runs flashinfer's below. + ran_multi_instance = _multi_instance_context.requested() and ( + _multi_instance_context.maybe_run_context( + params, + q_processed, + kv_pool, + block_tables, + bmm1_scale, + bmm2_scale, + cu_q_seqlens, + window_left, + fp8_context_fmha, + attention_chunk_size, + causal, + self._multi_processor_count, + ) ) + if not ran_multi_instance: + flashinfer.prefill.trtllm_batch_context_with_kv_cache( + query=q_processed, + kv_cache=(kv_pool, kv_pool), + workspace_buffer=fmha_workspace, + block_tables=block_tables, + seq_lens=params.sequence_lengths, + max_q_len=max_q_len, + max_kv_len=max_kv_len, + bmm1_scale=ctx_bmm1_scale, + bmm2_scale=ctx_bmm2_scale, + batch_size=params.batch_size, + cum_seq_lens_q=cu_q_seqlens, + cum_seq_lens_kv=cu_kv_seqlens, + window_left=window_left, + out=params.output, + kv_layout=self._layout, + enable_pdl=self._enable_pdl, + sinks=fwd.attention_sinks, + kv_cache_sf=(kv_scale_pool, kv_scale_pool) if kv_scale_pool is not None else None, + uses_shared_paged_kv_idx=self.USE_SHARED_PAGED_KV_IDX, + causal=causal, + multi_ctas_kv_counter_buffer=self._get_multi_ctas_kv_counter_buffer(), + ) if params.is_cross: return diff --git a/tensorrt_llm/_torch/attention/backends/fmha/multi_instance_context.py b/tensorrt_llm/_torch/attention/backends/fmha/multi_instance_context.py new file mode 100644 index 000000000000..2713619084a5 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fmha/multi_instance_context.py @@ -0,0 +1,817 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Host side of the sm_103 multi-instance paged context FMHA (``multi_instance_context_kernel``). + +``plan_context_fmha`` turns the step's per-request query and key lengths into the kernel's work +units and per-CTA unit ranges; ``paged_context_fmha_multi_instance`` runs the kernel on a plan and, when the +plan cut a unit's key range, merges the partial results. + +Planning. An item is 128 query rows of one request times ``num_inst`` heads of one KV group over +the request's key tiles up to the item's last row. Its cost is its tile count plus a per-item +overhead (``ITEM_OVERHEAD``; a piece of a cut item ``PIECE_OVERHEAD``). With ``split=True`` (the +default) the items, longest first, are laid end to end and cut +into ``num_ctas`` equal shares of the total cost (stream-K): a CTA's share is a list of whole +items plus at most a head and a tail piece of items it shares with its neighbours; a cut is +placed only before an item's first masked tile, so the masked tiles of an item are always in its +last piece. Pieces write ``O / l`` and ``(m, l)`` to a workspace; ``merge`` combines the pieces +of every cut item. A share is at least two minimum pieces plus the overhead (a small call uses +fewer CTAs), and the split is kept only when its busiest CTA has less work than with whole items +assigned longest first to the CTA with the least work (LPT), which ``split=False`` always uses. + +Numerics are the kernel's (see ``multi_instance_context_kernel``); the merge weights every piece by +``l * 2^(m - max m)`` in fp32 and rounds the merged row to BF16 once. +""" + +from __future__ import annotations + +import functools +import heapq +import math +import os +import time +from dataclasses import dataclass +from typing import Optional, Sequence + +import numpy as np +import torch + +from tensorrt_llm._utils import maybe_pin_memory + +#: Head width, tokens per page and query rows per instance the kernel is built for. +HEAD_DIM = 64 +PAGE = 32 +M_TILE = 128 +UNIT_FIELDS = 8 +#: Planner cost of a unit beyond its key tiles, in key tiles: a whole item (BF16 epilogue, +#: diagonal steps, the next unit's Q wait) and a piece of a cut item (fp32 partial epilogue). +#: Fitted to a per-unit trace of the kernel on another Blackwell-family part (short units cost about +#: 7.6 and 13 tiles beyond their tile count; deep units amortise either). +ITEM_OVERHEAD = 8.0 +PIECE_OVERHEAD = 13.0 + + +@dataclass +class MultiInstancePlan: + """Work units on the device plus what the merge needs.""" + + units: torch.Tensor # int32 [num_units, 8] + cta_off: torch.Tensor # int32 [num_ctas + 1] + num_units: int + num_ctas: int + num_slots: int # workspace slots (pieces of cut items) + merge_meta: Optional[torch.Tensor] # int32 [num_merges, 4]: b, m, hg, num_pieces + merge_slots: Optional[torch.Tensor] # int32 [num_merges, max_pieces] + max_pieces: int + num_inst: int + tile_n: int + makespan: float # planned cost of the busiest CTA + ideal: float # total cost / num_ctas + + +def _items(q_lens, kv_lens, num_heads, num_kv_heads, num_inst, tile_n): + groups = num_heads // num_inst + items = [] + for b, (ql, kl) in enumerate(zip(q_lens, kv_lens)): + if ql <= 0 or kl <= 0: + continue + off = kl - ql + last_tile_req = (kl - 1) // tile_n + for m in range(-(-ql // M_TILE)): + t_end = min(last_tile_req, (off + m * M_TILE + M_TILE - 1) // tile_n) + 1 + first_masked = (off + m * M_TILE + 1) // tile_n + for hg in range(groups): + items.append((b, m, hg, t_end, min(first_masked, t_end))) + return items + + +def _lpt(items, num_ctas, item_overhead): + """Whole items, longest first, each to the CTA with the least work.""" + per_cta = [[] for _ in range(num_ctas)] + loads = [0.0] * num_ctas + heap = [(0.0, k) for k in range(num_ctas)] + for it in items: + load, k = heapq.heappop(heap) + per_cta[k].append((it[0], it[1], it[2], 0, it[3], 0, 0)) + load += it[3] + item_overhead + loads[k] = load + heapq.heappush(heap, (load, k)) + return per_cta, loads, {}, 0 + + +def _stream_k(items, num_ctas, share, item_overhead, piece_overhead, min_piece): + """Items laid end to end and cut into CTA shares of ``share`` (see the module docstring). + + A whole item costs its tiles plus ``item_overhead``, a piece of a cut item its tiles plus + ``piece_overhead`` (its partial epilogue writes fp32 rows and statistics). + """ + per_cta = [[] for _ in range(num_ctas)] + loads = [0.0] * num_ctas + pieces = {} # item index -> list of slots + num_slots = 0 + k = 0 + filled = 0.0 + for idx, (b, m, hg, t_end, fm) in enumerate(items): + a = 0 + while a < t_end: + cap = share * (k + 1) - filled + need = (t_end - a) + (item_overhead if a == 0 else piece_overhead) + if need <= cap + 0.5 or k == num_ctas - 1: + if a == 0: + per_cta[k].append((b, m, hg, 0, t_end, 0, 0)) + else: + per_cta[k].append((b, m, hg, a, t_end, 1, num_slots)) + pieces[idx].append(num_slots) + num_slots += 1 + filled += need + loads[k] += need + break + x = int(math.floor(cap - piece_overhead)) + x = min(x, fm - a) + if x < min_piece or t_end - (a + x) < min_piece: + if cap < (t_end - a) / 2 + piece_overhead and k < num_ctas - 1: + # Too little room left here for a useful piece: close this CTA. + filled = share * (k + 1) + k += 1 + continue + # The piece cannot be cut usefully: the CTA takes the rest whole. + if a == 0: + per_cta[k].append((b, m, hg, 0, t_end, 0, 0)) + else: + per_cta[k].append((b, m, hg, a, t_end, 1, num_slots)) + pieces[idx].append(num_slots) + num_slots += 1 + filled += need + loads[k] += need + break + pieces.setdefault(idx, []).append(num_slots) + per_cta[k].append((b, m, hg, a, a + x, 1, num_slots)) + num_slots += 1 + filled += x + piece_overhead + loads[k] += x + piece_overhead + a += x + k += 1 + return per_cta, loads, pieces, num_slots + + +def plan_context_fmha( + q_lens: Sequence[int], + kv_lens: Sequence[int], + *, + num_ctas: int, + num_heads: int = 32, + num_kv_heads: int = 4, + num_inst: int = 4, + tile_n: int = 64, + split: bool = True, + item_overhead: float = ITEM_OVERHEAD, + piece_overhead: float = PIECE_OVERHEAD, + min_piece: int = 16, + device: Optional[torch.device] = None, +) -> MultiInstancePlan: + """The work units of one context call; see the module docstring. + + The plan's int32 arrays reach ``device`` in one copy: from pinned host memory without a + host wait when ``device`` is a CUDA device (the copy is ordered on the current stream, so a + kernel launched after this call on that stream reads it), synchronously otherwise. + """ + items = _items(q_lens, kv_lens, num_heads, num_kv_heads, num_inst, tile_n) + items.sort(key=lambda it: -it[3]) + per_cta, loads, pieces, num_slots = _lpt(items, num_ctas, item_overhead) + total = sum(it[3] + item_overhead for it in items) + ideal = total / max(num_ctas, 1) + if split and items: + # A cut turns one whole item into two pieces; a stream-K split makes at most + # num_ctas - 1 cuts. A share below two minimum pieces would close every CTA without a + # cut (the last CTA would take the rest), so small calls use fewer CTAs. + per_cut = max(2.0 * piece_overhead - item_overhead, 0.0) + share = (total + per_cut * (num_ctas - 1)) / max(num_ctas, 1) + share = max(share, 2.0 * min_piece + piece_overhead) + sk = _stream_k(items, num_ctas, share, item_overhead, piece_overhead, min_piece) + if max(sk[1]) < max(loads): + per_cta, loads, pieces, num_slots = sk + units = [] + cta_off = [0] + for k in range(num_ctas): + units.extend(per_cta[k]) + cta_off.append(len(units)) + num_units = max(len(units), 1) + merges = [(items[i][0], items[i][1], items[i][2], s) for i, s in pieces.items()] + max_pieces = max((len(s) for *_, s in merges), default=0) + # One int32 buffer: units [num_units, 8], cta_off [num_ctas + 1], then (when items were + # cut) merge_meta [merges, 4] and merge_slots [merges, max_pieces]; every section starts at + # a 16-byte boundary. + lengths = [num_units * UNIT_FIELDS, num_ctas + 1, 4 * len(merges), max_pieces * len(merges)] + sizes = [-(-n // 4) * 4 for n in lengths] + starts = np.cumsum([0] + sizes[:-1]).tolist() + host = np.zeros(sum(sizes), dtype=np.int32) + + def section(i): + return host[starts[i] : starts[i] + lengths[i]] + + if units: + section(0)[: len(units) * UNIT_FIELDS].reshape(len(units), UNIT_FIELDS)[:, :7] = np.asarray( + units, dtype=np.int32 + ) + section(1)[:] = cta_off + if merges: + mm = section(2).reshape(len(merges), 4) + ms = section(3).reshape(len(merges), max_pieces) + ms.fill(-1) + for i, (b, m, hg, s) in enumerate(merges): + mm[i] = (b, m, hg, len(s)) + ms[i, : len(s)] = s + dev = device if device is not None else torch.device("cuda") + buf = torch.from_numpy(host) + if dev.type == "cuda": + buf = maybe_pin_memory(buf).to(dev, non_blocking=True) + else: + buf = buf.to(dev) + views = [buf[s : s + n] for s, n in zip(starts, lengths)] + return MultiInstancePlan( + units=views[0].view(num_units, UNIT_FIELDS), + cta_off=views[1], + num_units=len(units), + num_ctas=num_ctas, + num_slots=num_slots, + merge_meta=views[2].view(len(merges), 4) if merges else None, + merge_slots=views[3].view(len(merges), max_pieces) if merges else None, + max_pieces=max_pieces, + num_inst=num_inst, + tile_n=tile_n, + makespan=max(loads) if loads else 0.0, + ideal=ideal, + ) + + +@functools.lru_cache(maxsize=32) +def _compiled_kernel( + device_index: int, + num_heads: int, + num_kv_heads: int, + num_inst: int, + tile_n: int, + write_stats: bool, + ablate_mufu: bool, + kv_stage: int, + trace: bool = False, + s_stages: int = 1, + p_sep: bool = False, + rescale_tau: float = 0.0, + ex2_emu_every: int = 0, + ablate_p_in_o: bool = False, + ablate_rescale: bool = False, + rescale_cols: int = 16, +): + """The compiled kernel of one build on one device (every size is a runtime argument).""" + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import make_ptr + + from . import multi_instance_context_kernel as fk + + kernel = fk.MultiInstancePagedContextFmha( + num_heads, + num_kv_heads, + num_inst=num_inst, + tile_n=tile_n, + kv_stage=kv_stage, + s_stages=s_stages, + write_stats=write_stats, + ablate_mufu=ablate_mufu, + trace=trace, + p_sep=p_sep, + rescale_tau=rescale_tau, + ex2_emu_every=ex2_emu_every, + ablate_p_in_o=ablate_p_in_o, + ablate_rescale=ablate_rescale, + rescale_cols=rescale_cols, + ) + gmem = cute.AddressSpace.gmem + ptrs = [ + make_ptr(cutlass.Float8E4M3FN, 16, gmem, assumed_align=16), # q + make_ptr(cutlass.Float8E4M3FN, 16, gmem, assumed_align=16), # pool + make_ptr(cutlass.Int32, 16, gmem, assumed_align=4), # table + make_ptr(cutlass.Int32, 16, gmem, assumed_align=4), # cu_q + make_ptr(cutlass.Int32, 16, gmem, assumed_align=4), # kv_lens + make_ptr(cutlass.BFloat16, 16, gmem, assumed_align=16), # out + make_ptr(cutlass.Int32, 16, gmem, assumed_align=16), # units + make_ptr(cutlass.Int32, 16, gmem, assumed_align=4), # cta_off + make_ptr(cutlass.Float32, 16, gmem, assumed_align=16), # part_o + make_ptr(cutlass.Float32, 16, gmem, assumed_align=8), # part_ml + make_ptr(cutlass.Float32, 16, gmem, assumed_align=8), # stats + make_ptr(cutlass.Float32, 16, gmem, assumed_align=4), # scale_log2 + make_ptr(cutlass.Float32, 16, gmem, assumed_align=4), # scale_out + ] + with torch.cuda.device(device_index): + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile( + kernel, + *ptrs, + *[cutlass.Int32(1)] * 7, + stream, + # Opt-level 1: level 3 (the DSL default) made the GB300 prototype 25-30 % slower. + options="--enable-tvm-ffi --opt-level=1", + ) + + +_MERGE_KERNEL = None + + +def _merge_kernel(): + global _MERGE_KERNEL + if _MERGE_KERNEL is not None: + return _MERGE_KERNEL + import triton + import triton.language as tl + + # ``maxp`` (the row stride of the slot table) is a runtime value, unspecialised, so one + # compiled kernel serves every plan. + @triton.jit(do_not_specialize=["maxp"]) + def _merge( + out_ptr, + stats_ptr, + part_o_ptr, + part_ml_ptr, + meta_ptr, + slots_ptr, + cu_q_ptr, + maxp, + H: tl.constexpr, + NI: tl.constexpr, + DH: tl.constexpr, + MT: tl.constexpr, + WRITE_STATS: tl.constexpr, + ): + e = tl.program_id(0) + j = tl.program_id(1) + b = tl.load(meta_ptr + e * 4 + 0) + m = tl.load(meta_ptr + e * 4 + 1) + hg = tl.load(meta_ptr + e * 4 + 2) + npc = tl.load(meta_ptr + e * 4 + 3) + q0 = tl.load(cu_q_ptr + b) + q_len = tl.load(cu_q_ptr + b + 1) - q0 + rows = tl.arange(0, MT) + d = tl.arange(0, DH) + mx = tl.full([MT], float("-inf"), tl.float32) + for p in range(0, npc): + s = tl.load(slots_ptr + e * maxp + p).to(tl.int64) + base = (s * NI + j) * MT + rows + ml = tl.load(part_ml_ptr + base * 2) + lp = tl.load(part_ml_ptr + base * 2 + 1) + mx = tl.where(lp > 0, tl.maximum(mx, ml), mx) + acc = tl.zeros([MT, DH], tl.float32) + den = tl.zeros([MT], tl.float32) + for p in range(0, npc): + s = tl.load(slots_ptr + e * maxp + p).to(tl.int64) + base = (s * NI + j) * MT + rows + ml = tl.load(part_ml_ptr + base * 2) + lp = tl.load(part_ml_ptr + base * 2 + 1) + w = tl.where(lp > 0, lp * tl.exp2(ml - mx), 0.0) + o = tl.load(part_o_ptr + base[:, None] * DH + d[None, :]) + acc += w[:, None] * o + den += w + res = acc / tl.where(den > 0, den, 1.0)[:, None] + r = m * MT + rows + ok = r < q_len + head = hg * NI + j + dst = (q0 + r).to(tl.int64) * H + head + tl.store(out_ptr + dst[:, None] * DH + d[None, :], res.to(tl.bfloat16), mask=ok[:, None]) + if WRITE_STATS: + # natural-log max of the scaled scores and the row sum without the 2^8 pre-scale + tl.store(stats_ptr + dst * 2, mx * 0.6931471805599453, mask=ok) + tl.store(stats_ptr + dst * 2 + 1, den * 0.00390625, mask=ok) + + _MERGE_KERNEL = _merge + return _MERGE_KERNEL + + +def max_slots(num_ctas: int) -> int: + """The most workspace slots a split plan over ``num_ctas`` CTAs uses: each of its at most + ``num_ctas - 1`` cuts adds at most two pieces.""" + return 2 * max(num_ctas - 1, 1) + + +class MultiInstanceWorkspace: + """Partial-result buffers. ``reserve`` sizes them up front; ``get`` grows them on demand + otherwise (never inside a CUDA-graph capture, which a reserved workspace never needs).""" + + def __init__(self): + self.part_o = None + self.part_ml = None + + def reserve(self, slots: int, num_inst: int, device) -> None: + self.get(slots, num_inst, device) + + def get(self, slots: int, num_inst: int, device): + need = max(slots, 1) * num_inst * M_TILE + if self.part_o is None or self.part_o.numel() < need * HEAD_DIM: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "MultiInstanceWorkspace: the workspace must be reserved before capture" + ) + self.part_o = torch.empty(need * HEAD_DIM, dtype=torch.float32, device=device) + self.part_ml = torch.empty(need * 2, dtype=torch.float32, device=device) + return self.part_o, self.part_ml + + @property + def nbytes(self) -> int: + if self.part_o is None: + return 0 + return (self.part_o.numel() + self.part_ml.numel()) * 4 + + +_WORKSPACE = MultiInstanceWorkspace() + + +def paged_context_fmha_multi_instance( + q: torch.Tensor, + kv_pool: torch.Tensor, + table: torch.Tensor, + cu_q: torch.Tensor, + kv_lens: torch.Tensor, + plan: MultiInstancePlan, + scale_log2: torch.Tensor, + scale_out: torch.Tensor, + out: torch.Tensor, + *, + stats: Optional[torch.Tensor] = None, + ablate_mufu: bool = False, + kv_stage: int = 12, + workspace: Optional[MultiInstanceWorkspace] = None, + trace: Optional[torch.Tensor] = None, + s_stages: int = 1, + p_sep: bool = False, + rescale_tau: float = 0.0, + ex2_emu_every: int = 0, + ablate_p_in_o: bool = False, + ablate_rescale: bool = False, + rescale_cols: int = 16, +) -> torch.Tensor: + """Causal (bottom-right) attention of the packed new rows ``q`` over their paged contexts. + + ``q`` E4M3 ``[total_q, heads, 64]``; ``kv_pool`` E4M3 ``[blocks, kv_heads, 32, 64]``; + ``table`` int32 ``[batch, 2, width]`` (K and V block of each page); ``cu_q`` int32 + ``[batch + 1]``; ``kv_lens`` int32 ``[batch]``; ``scale_log2`` and ``scale_out`` fp32 device + scalars (any tensor whose first element is the value) = the softmax scale times log2(e) and + the output scale; ``out`` BF16 ``[total_q, heads, 64]``; ``stats`` optional fp32 + ``[total_q, heads, 2]`` (natural-log row max, row sum). ``plan`` from ``plan_context_fmha`` on + the same lengths. ``trace`` (diagnostic): an int64 buffer of ``plan.num_units * + TRACE_FIELDS + TRACE_STEP_WORDS`` words (``multi_instance_context_kernel``) selects the trace build. + ``s_stages`` 2 (with ``plan.tile_n`` 32) double-buffers S per instance; ``p_sep``, + ``rescale_tau``, ``ex2_emu_every``, ``rescale_cols`` and the ``ablate_*`` diagnostics select + the kernel build (``multi_instance_context_kernel.MultiInstancePagedContextFmha``). No host + synchronisation: every input is read on the device, in stream order. + """ + import cuda.bindings.driver as cuda + + heads, kv_heads = q.shape[1], kv_pool.shape[1] + dev = q.device + assert stats is None or trace is None + compiled = _compiled_kernel( + dev.index if dev.index is not None else torch.cuda.current_device(), + heads, + kv_heads, + plan.num_inst, + plan.tile_n, + stats is not None, + bool(ablate_mufu), + int(kv_stage), + trace is not None, + int(s_stages), + bool(p_sep), + float(rescale_tau), + int(ex2_emu_every), + bool(ablate_p_in_o), + bool(ablate_rescale), + int(rescale_cols), + ) + ws = workspace if workspace is not None else _WORKSPACE + part_o, part_ml = ws.get(plan.num_slots, plan.num_inst, dev) + batch = kv_lens.shape[0] + stats_ptr = stats.data_ptr() if stats is not None else part_ml.data_ptr() + if trace is not None: + stats_ptr = trace.data_ptr() + compiled( + q.data_ptr(), + kv_pool.data_ptr(), + table.data_ptr(), + cu_q.data_ptr(), + kv_lens.data_ptr(), + out.data_ptr(), + plan.units.data_ptr(), + plan.cta_off.data_ptr(), + part_o.data_ptr(), + part_ml.data_ptr(), + stats_ptr, + scale_log2.data_ptr(), + scale_out.data_ptr(), + q.shape[0], + batch, + kv_pool.shape[0], + table.stride(0), + table.stride(1), + max(plan.num_units, 1), + plan.num_ctas, + cuda.CUstream(torch.cuda.current_stream(dev).cuda_stream), + ) + if plan.merge_meta is not None: + kern = _merge_kernel() + kern[(plan.merge_meta.shape[0], plan.num_inst)]( + out, + stats if stats is not None else part_ml, + part_o, + part_ml, + plan.merge_meta, + plan.merge_slots, + cu_q, + plan.max_pieces, + H=heads, + NI=plan.num_inst, + DH=HEAD_DIM, + MT=M_TILE, + WRITE_STATS=stats is not None, + num_warps=4, + ) + return out + + +# ----------------------------------------------------------------------------- serving integration +# ``FlashInferTrtllmGenFmha.run_context`` calls ``maybe_run_context`` after the context preprocess +# (RoPE, K / V appended to the paged cache, Q quantized to E4M3, the BMM scales on the device) and +# runs flashinfer's trtllm-gen context kernel only when it returns False. + +#: The switch: unset / ``0`` off (flashinfer's trtllm-gen context kernel, as before), ``1`` on. +ENV = "TRTLLM_MULTI_INSTANCE_CONTEXT_FMHA" +#: Validation: ``1`` reads the device KV lengths back each call (a host sync) and declines a call +#: whose device lengths differ from the host lengths it planned with. +CHECK_ENV = "TRTLLM_MULTI_INSTANCE_CONTEXT_FMHA_CHECK" +#: The SM versions the kernel was built and measured on. +SERVED_SMS = (103,) +#: Query instances per CTA and key tile of the build this path runs. +NUM_INST = 4 +TILE_N = 64 + + +def _env_flag(name: str) -> bool: + value = os.environ.get(name, "").strip().lower() + if value in ("", "0", "off", "false"): + return False + if value in ("1", "on", "true"): + return True + raise ValueError(f"{name}={value!r}: expected 0 or 1") + + +@functools.lru_cache(maxsize=1) +def requested() -> bool: + """Whether ``ENV`` asks for this path (read once per process).""" + return _env_flag(ENV) + + +@functools.lru_cache(maxsize=1) +def _check_requested() -> bool: + return _env_flag(CHECK_ENV) + + +class _StepPlans: + """The plan of the latest lengths: a step's full-attention layers run one after another with + the same lengths, so the first plans and the others reuse its device buffer (read in stream + order). A step with other lengths replaces it.""" + + def __init__(self): + self._key = None + self._plan: Optional[MultiInstancePlan] = None + self.planned = 0 + self.plan_seconds = 0.0 + + def get(self, q_lens, kv_lens, num_heads, num_kv_heads, num_ctas, device) -> MultiInstancePlan: + key = (q_lens, kv_lens, num_heads, num_kv_heads, num_ctas, device) + if key != self._key: + start = time.perf_counter() + self._plan = plan_context_fmha( + q_lens, + kv_lens, + num_ctas=num_ctas, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + num_inst=NUM_INST, + tile_n=TILE_N, + device=device, + ) + self.plan_seconds += time.perf_counter() - start + self.planned += 1 + self._key = key + return self._plan + + +_PLANS = _StepPlans() +_RESERVED: dict = {} +#: Context calls this path ran and declined in this process. +calls = {"ran": 0, "declined": 0} + + +def _reserved_workspace(device, num_ctas: int) -> MultiInstanceWorkspace: + """The partial-result workspace, sized once for any plan over ``num_ctas`` CTAs.""" + key = (device, num_ctas) + if key not in _RESERVED: + _WORKSPACE.reserve(max_slots(num_ctas), NUM_INST, device) + _RESERVED[key] = True + return _WORKSPACE + + +def _decline_reason( + params, + q, + kv_pool, + block_tables, + bmm1_scale, + bmm2_scale, + cu_q, + window_left, + fp8_context_fmha, + attention_chunk_size, + causal, + num_ctas, +): + """None when this path can run the context call, else why not (host-side checks only).""" + from tensorrt_llm._utils import get_sm_version + from tensorrt_llm.quantization.mode import QuantMode + + attn, meta, fwd = params.attn, params.meta, params.fwd + sm = get_sm_version() + if sm not in SERVED_SMS: + return f"the kernel is built for sm{SERVED_SMS}, not sm{sm}" + if torch.cuda.is_current_stream_capturing(): + return "a CUDA graph capture is in progress" + if params.is_cross or not causal: + return "a mask other than causal self-attention" + if window_left >= 0: + return "a sliding window shorter than the context" + if attention_chunk_size != 0: + return "chunked attention" + if fwd.attention_sinks is not None: + return "attention sinks" + quant = QuantMode(attn.quant_mode) + if not quant.has_fp8_kv_cache() or quant.has_fp4_kv_cache(): + return "a KV cache that is not FP8" + if not fp8_context_fmha or q.dtype != torch.float8_e4m3fn: + return "a context FMHA that is not FP8" + if attn.head_dim != HEAD_DIM: + return f"head dim {attn.head_dim}" + heads, kv_heads = attn.num_heads, attn.num_kv_heads + if heads % kv_heads or (heads // kv_heads) % NUM_INST: + return f"{heads} query heads over {kv_heads} KV heads" + if params.tokens_per_block != PAGE: + return f"a page size of {params.tokens_per_block}" + batch, num_tokens = params.batch_size, params.num_tokens + out = params.output + if ( + out is None + or out.dtype != torch.bfloat16 + or not out.is_contiguous() + or out.numel() != num_tokens * heads * HEAD_DIM + ): + return "an output that is not contiguous BF16 of the context rows" + if tuple(q.shape) != (num_tokens, heads, HEAD_DIM): + return "a query layout other than [tokens, heads, 64]" + if ( + kv_pool.dim() != 4 + or tuple(kv_pool.shape[1:]) != (kv_heads, PAGE, HEAD_DIM) + or kv_pool.element_size() != 1 + or not kv_pool.is_contiguous() + ): + return f"a KV pool view of shape {tuple(kv_pool.shape)}" + if kv_pool.shape[0] * kv_heads >= 2**31: + return "a KV pool of 2^31 pages or more" + if ( + block_tables.dim() != 3 + or block_tables.shape[0] < batch + or block_tables.shape[1] != 2 + or block_tables.stride(-1) != 1 + or block_tables.dtype != torch.int32 + ): + return f"a block table of shape {tuple(block_tables.shape)}" + if cu_q.dtype != torch.int32 or cu_q.numel() < batch + 1: + return "cumulative query lengths that are not int32 [batch + 1]" + if params.sequence_lengths is None or params.sequence_lengths.dtype != torch.int32: + return "device KV lengths that are not int32" + if bmm1_scale is None or bmm1_scale.numel() < 2 or bmm2_scale is None: + return "no device BMM scales" + # One-model speculative decoding (DFlash: max_draft_len - 1 extra slots) reserves extra KV + # tokens per sequence. Neither length this path reads counts them: ``prepare`` binds the + # pre-extra ``kv_lens`` as the host ``kv_lens_runtime`` the plan is built from, and the device + # ``kv_lens_cuda`` the kernel masks with (and flashinfer's context kernel attends) excludes them + # too; only the ``meta.kv_lens`` buffer adds them. A path that bound that buffer would plan + # past the attended keys, so it is declined. + extra = getattr(getattr(meta, "kv_cache_params", None), "num_extra_kv_tokens", 0) or 0 + host_kv = getattr(meta, "kv_lens_runtime", None) + with_extra = getattr(meta, "kv_lens", None) + if ( + extra + and isinstance(host_kv, torch.Tensor) + and isinstance(with_extra, torch.Tensor) + and host_kv.untyped_storage().data_ptr() == with_extra.untyped_storage().data_ptr() + ): + return "host KV lengths that include the extra KV tokens" + if num_ctas is None or num_ctas <= 0: + return "an unknown SM count" + return None + + +def maybe_run_context( + params, + q, + kv_pool, + block_tables, + bmm1_scale, + bmm2_scale, + cu_q, + window_left, + fp8_context_fmha, + attention_chunk_size, + causal, + num_ctas, +) -> bool: + """Run the context call of ``params`` on the multi-instance kernel when ``ENV`` is on and the + call is one it serves; return whether it ran. A call it declines (logged once per reason) is + left to flashinfer's trtllm-gen context kernel.""" + from tensorrt_llm.logger import logger + + if not requested(): + return False + reason = _decline_reason( + params, + q, + kv_pool, + block_tables, + bmm1_scale, + bmm2_scale, + cu_q, + window_left, + fp8_context_fmha, + attention_chunk_size, + causal, + num_ctas, + ) + lengths = None + if reason is None: + meta, batch, s0 = params.meta, params.batch_size, params.seq_offset + q_lens = tuple(int(x) for x in meta.prompt_lens_cpu_runtime[s0 : s0 + batch].tolist()) + kv_lens = tuple(int(x) for x in meta.kv_lens_runtime[s0 : s0 + batch].tolist()) + if ( + len(q_lens) != batch + or len(kv_lens) != batch + or sum(q_lens) != params.num_tokens + or any(ql <= 0 or kl < ql for ql, kl in zip(q_lens, kv_lens)) + ): + reason = "host lengths that disagree with the call" + elif _check_requested(): + device_kv = tuple(int(x) for x in params.sequence_lengths[:batch].tolist()) + if device_kv != kv_lens: + reason = "device KV lengths that differ from the host's" + lengths = (q_lens, kv_lens) + if reason is not None: + calls["declined"] += 1 + logger.info_once( + f"[{ENV}] a context call ran flashinfer's trtllm-gen kernel (layer " + f"{getattr(params.attn, 'layer_idx', '?')}): {reason}", + key=f"{ENV}.decline.{reason}", + ) + return False + attn = params.attn + q_lens, kv_lens = lengths + device = q.device + plan = _PLANS.get(q_lens, kv_lens, attn.num_heads, attn.num_kv_heads, num_ctas, device) + workspace = _reserved_workspace(device, num_ctas) + paged_context_fmha_multi_instance( + q, + kv_pool, + block_tables, + cu_q, + params.sequence_lengths[: params.batch_size], + plan, + bmm1_scale.reshape(-1)[1:2], + bmm2_scale.reshape(-1)[0:1], + params.output.view(params.num_tokens, attn.num_heads, HEAD_DIM), + workspace=workspace, + ) + calls["ran"] += 1 + logger.info_once( + f"[{ENV}] the multi-instance context FMHA ran a context call (layer " + f"{getattr(attn, 'layer_idx', '?')}: {params.batch_size} requests, {params.num_tokens} rows, " + f"{num_ctas} CTAs, plan {plan.num_units} units / {plan.num_slots} partials)", + key=f"{ENV}.ran", + ) + return True diff --git a/tensorrt_llm/_torch/attention/backends/fmha/multi_instance_context_kernel.py b/tensorrt_llm/_torch/attention/backends/fmha/multi_instance_context_kernel.py new file mode 100644 index 000000000000..245e3131903c --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fmha/multi_instance_context_kernel.py @@ -0,0 +1,1312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Portions of this file are derived from the CUTLASS CuTe DSL Blackwell FMHA example +# (examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py, CUTLASS v4.8.0): +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +"""Paged FP8 context FMHA for sm_103 (Blackwell Ultra) with several independent softmax streams per SMSP. + +One rank's full-attention context call: ``total_q`` new query rows of ``batch`` requests packed +back to back (``cu_q``), ``num_heads`` query heads and ``num_kv_heads`` KV heads of width 64, E4M3 +Q ``[total_q, heads, 64]``, E4M3 K / V in the paged pool ``[blocks, kv_heads, 32, 64]`` addressed +through a ``[batch, 2, width]`` block-index table (``table[b, 0, p]`` is the block holding keys +``32 p .. 32 p + 31`` of request ``b``, ``table[b, 1, p]`` the block of their values), BF16 O +``[total_q, heads, 64]``. Request ``b`` has ``kv_lens[b]`` keys; its query rows are its last +``q_len(b)`` positions, so row ``r`` sees keys ``0 .. kv_lens[b] - q_len(b) + r``. + +Organisation. One CTA per SM runs ``num_inst`` query instances at once. Instance ``j`` is 128 +query rows of one head with its own S (``tile_n`` columns) and O (64 columns) in TMEM and its own +softmax warpgroup; the instances of a CTA are the heads ``num_inst * hg + j`` of one KV group, so +every K / V tile in shared memory feeds all of them. With four instances, every SMSP holds four +independent softmax warps (one per instance), which keeps the MUFU pipe fed while any one of them +waits for its next S tile or for a TMEM load. + +Work. The kernel executes an explicit list of work units, eight int32 each, ``(b, m, hg, t0, t1, +mode, slot, 0)``: request ``b``, query rows ``128 m .. 128 m + 127`` of the request, heads +``num_inst * hg + j``, key tiles ``[t0, t1)`` of ``tile_n`` keys. CTA ``k`` runs units +``cta_off[k] .. cta_off[k + 1] - 1`` in order (a host planner balances them and may cut a unit's +key range into several units). ``mode`` 0 writes BF16 O (and, if built with ``write_stats``, the +natural-log row max and the row sum, as the served kernel's softmax statistics); ``mode`` 1 writes +``scale_output * O / l`` in fp32 and ``(m * scale_log2, l)`` of every row to workspace slot ``slot`` for a merge +kernel. A cut never falls inside a unit's masked (diagonal) tiles, so every partial but the last +is mask-free. + +Roles (warp ids): warps ``4 j .. 4 j + 3`` softmax of instance ``j`` (thread ``t`` of the +warpgroup owns TMEM lane / query row ``t``); warp ``4 num_inst`` MMA issue and TMEM allocation; +warp ``4 num_inst + 1`` TMA load; the last two warps idle (they complete the warpgroup that +lowers its register budget). + +Ownership and progress (cuFlow record, per CTA; all waits are on this CTA's own barriers, every +role walks the same unit list in the same order with the same per-unit counts): + +* Q ring (``3 num_inst`` slots of 128 x 64 B, ``q_full`` / ``q_empty``): the load warp fills the + ``num_inst`` slots of a unit after ``q_empty``, the first unit's before its loop and unit + ``u + 1``'s at the start of unit ``u``; the MMA warp reads them in every QK of the unit and + releases them with a ``tcgen05.commit`` after the unit's last QK. Unit ``u + 1`` reuses the + slots of unit ``u - 2``, which the MMA warp released before it took unit ``u - 1``'s K / V + (the load warp's earlier copies), so the wait cannot block on the MMA warp's current unit. +* K / V ring (``kv_stage`` slots of ``tile_n`` x 64 B, ``kv_full`` / ``kv_empty``): per unit the + load warp fills K(t0), then per tile t: V(t), K(t + 1) with one S stage, or K(t + 1), V(t) with + two; each tile is ``tile_n / 32`` page boxes of 32 tokens (2 KiB) under one ``expect_tx``; a + page past the request's last page gets a box index past the pool and is zero-filled by the TMA + unit. The MMA warp consumes in the same order (K(t0) at the unit prologue). K(t) is released + after the QK of every instance, V(t) after the P.V of every instance (commits). +* V rows past ``kv_lens[b]`` in a request's last tile may hold any bytes (an E4M3 NaN times P = 0 + is NaN in O): the MMA warp zeroes them in shared memory after that slot's ``kv_full`` and + before its first P.V, then ``fence.proxy.async.shared::cta`` and a warp sync; the slot cannot be + refilled before its ``kv_empty`` commit, which follows those P.Vs. +* TMEM: instance ``j`` owns ``s_stages * tile_n + 64`` columns from ``j * inst_cols``: S buffers + ``0 .. s_stages - 1`` of ``tile_n`` columns, then O_j (64). Every S tile of a CTA has a running + index n (the MMA warp and the softmax warps count the same tiles of the same units); tile n goes + to buffer ``n % s_stages`` and its P_j (E4M3, ``tile_n / 4`` columns) is stored over the start + of that buffer. The tcgen05 MMAs of one issuing thread execute in issue order (the in-place P of + the CUTLASS example relies on the same). QK_j(n + s_stages) overwrites the buffer of P_j(n); it + is issued after P.V_j(n) (one stage: right after it; two stages: in the next step, which the MMA + warp reaches only after P.V_j(n)), so no S write overlaps a pending P read; P.V_j(n) is issued + only after ``p_full[j]`` for P_j(n), which softmax ``j`` arrives after ``tcgen05.wait::st``. + With two stages QK_j(n + 1) is issued before P.V_j(n) and computes while softmax ``j`` works + on S_j(n). +* ``s_full[j, b]`` (MMA -> softmax ``j``, by ``tcgen05.commit``) / ``s_empty[j, b]`` (128 + arrivals): once per S tile of buffer b on both sides, parities from the running index. Softmax + ``j`` arrives on ``s_empty`` once its TMEM loads of the tile completed; the MMA warp waits for it + before the next QK into that buffer. +* ``p_full[j]`` (128 arrivals) / ``p_empty[j]`` (commit after P.V_j): ``t1 - t0`` uses per unit + on both sides. Softmax ``j`` waits ``p_empty[j]`` at every step but a unit's first, for + P.V_j of the previous step, before it may rescale O_j (P.V_j of this step is not issued before + this step's ``p_full``), and once after a unit's last step (the drain) before its epilogue reads + O_j. The next unit's first P.V_j (accumulate off) overwrites O_j only after ``p_full[j]`` of that + unit's first P, which follows the epilogue's ``tcgen05.wait::ld``. +* TMEM (512 columns) is allocated by the MMA warp and freed by it after a named barrier that every + softmax warp reaches after its last epilogue read; the last MMAs are P.Vs whose completion the + softmax warps observed in their drains. +* Progress. The load warp waits only on empties released by the MMA warp; the MMA warp waits on + ``q_full`` / ``kv_full`` (load warp), ``p_full[j]`` and ``s_empty[j, b]`` (softmax ``j``); + softmax ``j`` waits on ``s_full[j, b]`` and ``p_empty[j]`` (MMA warp). The MMA warp's wait for + P_j(n) is satisfiable: S_j(n) was committed before it (prologue or an earlier QK), and the + softmax's ``p_empty`` waits are for P.Vs issued before (n - 1, or the unit's last). The MMA + warp's ``s_empty`` wait before QK_j(n + s_stages) is for the load of S_j(n), which softmax ``j`` + performs before it arrives P_j(n). No cross-CTA wait exists (partials are merged by a separate + launch). +* Programmatic dependent launch. Every warp executes ``griddepcontrol.wait`` before it reads any + input; the MMA warp issues ``griddepcontrol.launch_dependents`` after its unit loop (a dependent + grid's own wait returns only after this grid completed and its memory is visible). + +Numerics: S = Q K^T in fp32 on the tensor core; online softmax in log2 units with the running +maximum raised on every increase; P = 2^(s * scale_log2 - max * scale_log2 + 8) rounded to E4M3; +the row sum of the fp32 P in registers; O accumulated in fp32 on the tensor core and multiplied by +``scale_output / row_sum``. The scales are read at run time from two device scalars, +``scale_log2`` and ``scale_output`` (the context preprocess writes them: its BMM1 scale's second +entry and its BMM2 scale), so a CUDA-graph replay sees their current values. +""" + +import math + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode +from cutlass.cute.typing import Float32, Int32, Int64 +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +F8 = cutlass.Float8E4M3FN +BF16 = cutlass.BFloat16 + +#: Tokens per KV page, head width, query rows per instance. +PAGE = 32 +D = 64 +M_TILE = 128 +#: int32 fields per work unit. +UNIT_FIELDS = 8 +#: Page ids one load-warp lane holds per batch (one per lane). +IDS_PER_BATCH = 32 +#: P = 2^(s - max + P_PRESCALE_LOG2) uses E4M3's range (max 448) with the maximum raised on every +#: increase: floor(log2(448)). +P_PRESCALE_LOG2 = 8.0 +#: Diagnostic trace (``trace=True`` builds): int64 timestamps per work unit, then per-step cycle +#: counts of the first unit of CTA 0 (see ``MultiInstancePagedContextFmha``). +TRACE_FIELDS = 24 +TRACE_STEPS = 64 +TRACE_STEP_WORDS = 16 * TRACE_STEPS + + +@dsl_user_op +def _cvt_f32x4_to_e4m3x4(a, b, c, d, *, loc=None, ip=None) -> Int32: + """Four fp32 values rounded (satfinite) to E4M3 and packed into one 32-bit word, a lowest.""" + return Int32( + llvm.inline_asm( + T.i32(), + [Float32(x).ir_value(loc=loc, ip=ip) for x in (a, b, c, d)], + "{\n\t.reg .b16 lo, hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\t" + "mov.b32 $0, {lo, hi};\n\t}\n", + "=r,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _ex2_poly_f32x2(a, b, *, loc=None, ip=None): + """2^a and 2^b on the FMA / ALU pipes: 2^floor(x) times a degree-3 polynomial of the fraction + (relative error below 1e-4); inputs are clamped at -127.""" + res = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Float32(x).ir_value(loc=loc, ip=ip) for x in (a, b)], + "{\n\t.reg .f32 xa, xb;\n\t.reg .b64 x, r, w, f, p, k, c1, c2, c3, c0;\n\t" + ".reg .b32 ra, rb, pa, pb, ea, eb;\n\t" + "max.ftz.f32 xa, $2, 0fC2FE0000;\n\tmax.ftz.f32 xb, $3, 0fC2FE0000;\n\t" + "mov.b64 x, {xa, xb};\n\tmov.b64 k, 0x4B4000004B400000;\n\t" + "mov.b64 c3, 0x3D9DF09D3D9DF09D;\n\tmov.b64 c2, 0x3E6906A43E6906A4;\n\t" + "mov.b64 c1, 0x3F31F5193F31F519;\n\tmov.b64 c0, 0x3F8000003F800000;\n\t" + "add.rm.ftz.f32x2 r, x, k;\n\tsub.rn.ftz.f32x2 w, r, k;\n\tsub.rn.ftz.f32x2 f, x, w;\n\t" + "fma.rn.ftz.f32x2 p, c3, f, c2;\n\tfma.rn.ftz.f32x2 p, p, f, c1;\n\t" + "fma.rn.ftz.f32x2 p, p, f, c0;\n\t" + "mov.b64 {ra, rb}, r;\n\tmov.b64 {pa, pb}, p;\n\t" + "shl.b32 ea, ra, 23;\n\tshl.b32 eb, rb, 23;\n\t" + "add.s32 pa, pa, ea;\n\tadd.s32 pb, pb, eb;\n\t" + "mov.b32 $0, pa;\n\tmov.b32 $1, pb;\n\t}\n", + "=f,=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return ( + Float32(llvm.extractvalue(T.f32(), res, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), res, [1], loc=loc, ip=ip)), + ) + + +class MultiInstancePagedContextFmha: + """Compiled per (heads, KV heads, instances, key tile, options); sizes and pointers are runtime + arguments. + + :param num_heads: query heads of the rank. + :param num_kv_heads: KV heads of the rank. + :param num_inst: query instances per CTA (heads of one KV group run together); 2 or 4. + :param tile_n: keys per tile (32, 64 or 128); ``num_inst * (s_stages * tile_n + 64) <= 512`` + TMEM columns. + :param s_stages: S buffers per instance (1 or 2). With 2, QK_j(t + 1) is issued before P.V_j(t) + and computes while softmax ``j`` works on S_j(t). + :param kv_stage: K / V ring slots (one key tile of K or V each). + :param write_stats: mode-0 units also write the natural-log row max and the row sum. + :param p_sep: P gets its own TMEM columns after O (``tile_n + 64 + tile_n / 4`` columns per + instance, rounded up to 32, times ``num_inst`` must fit 512), so QK_j(t + 1) is issued + as soon as softmax ``j`` has read S_j(t), before P.V_j(t); one S buffer. + :param rescale_tau: the running maximum of a row is raised only when the tile's maximum + exceeds it by more than ``rescale_tau`` (log2 units of the scaled score); P then stays + below 2^(8 + tau) <= 448 (E4M3), so ``rescale_tau`` <= 0.8. 0 raises on every increase. + :param ex2_emu_every: with k > 0, every k-th group of four scores takes its exponentials on + the FMA / ALU pipes (``_ex2_poly_f32x2``) instead of MUFU. + :param rescale_cols: columns of O per TMEM load / store in the O rescale (16, 32 or 64; the + values are the same, a wider load takes fewer round trips). + :param ablate_p_in_o: diagnostic build, wrong values: P is stored over the instance's O + columns and QK_j(t + 1) is issued before P.V_j(t) (the timing of a P buffer outside S + at no TMEM cost). + :param ablate_rescale: diagnostic build, wrong values: O is never rescaled. + :param ablate_mufu: diagnostic build, wrong values: exp2 becomes a multiply (no MUFU), to time + everything else. + :param trace: diagnostic build: ``stats_ptr`` is an int64 buffer of ``num_units * + TRACE_FIELDS + TRACE_STEP_WORDS`` words. Unit ``u`` gets ``%globaltimer`` values at + ``u * TRACE_FIELDS + f``: f 0-3 MMA warp (unit start, Q ready, prologue QKs issued, unit + loop done), f 4-7 load warp (unit start, Q copies issued, page ids read, last K / V copy + issued), f 8 + j / 12 + j / 16 + j / 20 + j softmax ``j`` row 0 (unit start, unmasked + steps done, masked steps done, epilogue done). The first unit of CTA 0 also gets + ``%clock64`` per step ``s < TRACE_STEPS`` at ``num_units * TRACE_FIELDS + (4 g + j) * + TRACE_STEPS + s``: g 0 softmax ``j`` S ready, g 1 softmax ``j`` P arrived, g 2 MMA saw + P_j, g 3 MMA committed S_j. + """ + + arch_str = "sm_103" + + def __init__( + self, + num_heads: int = 32, + num_kv_heads: int = 4, + *, + num_inst: int = 4, + tile_n: int = 64, + kv_stage: int = 12, + s_stages: int = 1, + write_stats: bool = False, + ablate_mufu: bool = False, + trace: bool = False, + p_sep: bool = False, + rescale_tau: float = 0.0, + ex2_emu_every: int = 0, + ablate_p_in_o: bool = False, + ablate_rescale: bool = False, + rescale_cols: int = 16, + ): + assert num_heads % num_kv_heads == 0 + assert (num_heads // num_kv_heads) % num_inst == 0, "instances must share one KV head" + assert num_inst in (1, 2, 3, 4) and tile_n in (32, 64, 128) and s_stages in (1, 2) + assert not (trace and write_stats), "the trace build uses the stats buffer" + assert not (p_sep and ablate_p_in_o) + assert s_stages == 1 or not (p_sep or ablate_p_in_o) + assert 0.0 <= rescale_tau <= 0.8 and ex2_emu_every >= 0 + assert rescale_cols in (16, 32, 64) + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.h_r = num_heads // num_kv_heads + self.ni = num_inst + self.tile_n = tile_n + self.pages_per_tile = tile_n // PAGE + self.tiles_per_batch = IDS_PER_BATCH // self.pages_per_tile + self.slot_words = tile_n * D // 4 + self.row_words = D // 4 + self.kv_stage = kv_stage + self.s_stages = s_stages + self.q_stage = 3 * num_inst + self.write_stats = write_stats + self.ablate_mufu = ablate_mufu + self.trace = trace + self.rescale_tau = float(rescale_tau) + self.ex2_emu_every = int(ex2_emu_every) + self.ablate_rescale = ablate_rescale + self.rescale_cols = rescale_cols + self.qk_mma_tiler = (M_TILE, tile_n, D) + self.pv_mma_tiler = (M_TILE, D, tile_n) + # TMEM columns of an instance: S buffers from 0, O from o_off, P from p_off (P over the + # start of its S buffer unless p_sep). + if p_sep: + self.o_off = tile_n + self.p_off = tile_n + D + self.inst_cols = -(-(tile_n + D + tile_n // 4) // 32) * 32 + elif ablate_p_in_o: + self.o_off = tile_n + self.p_off = tile_n + self.inst_cols = tile_n + D + else: + self.o_off = s_stages * tile_n + self.p_off = 0 + self.inst_cols = s_stages * tile_n + D + #: QK_j(t + 1) is issued before P.V_j(t). + self.p_early = p_sep or ablate_p_in_o + assert num_inst * self.inst_cols <= 512, "TMEM columns" + + self.mma_warp_id = 4 * num_inst + self.load_warp_id = 4 * num_inst + 1 + self.num_warps = 4 * num_inst + 4 + self.threads_per_cta = 32 * self.num_warps + launch_regs = min(256, (65536 // self.threads_per_cta) // 8 * 8) + self.num_regs_other = 32 + softmax_regs = (launch_regs * self.threads_per_cta - 128 * self.num_regs_other) // ( + 128 * num_inst + ) + self.num_regs_softmax = min(248, softmax_regs // 8 * 8) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=1, num_threads=32 * (4 * num_inst + 1) + ) + self.tmem_dealloc_barrier = pipeline.NamedBarrier( + barrier_id=2, num_threads=32 * (4 * num_inst + 1) + ) + self.buffer_align_bytes = 1024 + + # ------------------------------------------------------------------ host side + @cute.jit + def __call__( + self, + q_ptr: cute.Pointer, + pool_ptr: cute.Pointer, + table_ptr: cute.Pointer, + cu_q_ptr: cute.Pointer, + kv_lens_ptr: cute.Pointer, + o_ptr: cute.Pointer, + units_ptr: cute.Pointer, + cta_off_ptr: cute.Pointer, + part_o_ptr: cute.Pointer, + part_ml_ptr: cute.Pointer, + stats_ptr: cute.Pointer, + scale_qk_ptr: cute.Pointer, + scale_o_ptr: cute.Pointer, + total_q: Int32, + batch: Int32, + num_blocks: Int32, + table_stride_b: Int32, + table_stride_kv: Int32, + num_units: Int32, + num_ctas: Int32, + stream, + ): + h_q, h_k = self.num_heads, self.num_kv_heads + q = cute.make_tensor(q_ptr, cute.make_layout((total_q, D, h_q), stride=(h_q * D, 1, D))) + o = cute.make_tensor(o_ptr, cute.make_layout((total_q, D, h_q), stride=(h_q * D, 1, D))) + # The pool as boxes of one page of one KV head: (token 32, dim 64, box). + kv_boxes = cute.make_tensor( + pool_ptr, cute.make_layout((PAGE, D, h_k * num_blocks), stride=(D, 1, PAGE * D)) + ) + table = cute.make_tensor(table_ptr, cute.make_layout(batch * table_stride_b)) + cu_q = cute.make_tensor(cu_q_ptr, cute.make_layout(batch + 1)) + kv_lens = cute.make_tensor(kv_lens_ptr, cute.make_layout(batch)) + units = cute.make_tensor(units_ptr, cute.make_layout(num_units * UNIT_FIELDS)) + cta_off = cute.make_tensor(cta_off_ptr, cute.make_layout(num_ctas + 1)) + scale_qk = cute.make_tensor(scale_qk_ptr, cute.make_layout(1)) + scale_o = cute.make_tensor(scale_o_ptr, cute.make_layout(1)) + + cta_group = tcgen05.CtaGroup.ONE + K = OperandMajorMode.K + MN = OperandMajorMode.MN + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + F8, F8, K, K, Float32, cta_group, self.qk_mma_tiler[:2] + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + F8, F8, K, MN, Float32, cta_group, self.pv_mma_tiler[:2], tcgen05.OperandSource.TMEM + ) + # The K operand of one page (N = 32): the shape of one TMA box in shared memory. + qk_box_mma = sm100_utils.make_trivial_tiled_mma( + F8, F8, K, K, Float32, cta_group, (M_TILE, PAGE) + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((1, 1, 1)), (qk_tiled_mma.thr_id.shape,) + ) + + q_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, self.qk_mma_tiler, F8, self.q_stage + ) + k_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, self.qk_mma_tiler, F8, self.kv_stage + ) + p_tmem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, self.pv_mma_tiler, F8, 1 + ) + v_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, self.pv_mma_tiler, F8, self.kv_stage + ) + box_shape = cute.dice((M_TILE, PAGE, D), (None, 1, 1)) + box_shape = qk_box_mma.partition_shape_B(box_shape) + box_atom = sm100_utils.make_smem_layout_atom(sm100_utils.SmemLayoutAtomKind.K_SW64, F8) + box_smem_layout_staged = tcgen05.tile_to_mma_shape( + box_atom, cute.append(box_shape, self.pages_per_tile * self.kv_stage), order=(1, 2, 3) + ) + + tma_load_op = cpasync.CopyBulkTensorTileG2SOp(cta_group) + q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, q, q_smem_layout, self.qk_mma_tiler, qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) # fmt: skip + box_smem_layout = cute.select(box_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_kv, tma_tensor_kv = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, kv_boxes, box_smem_layout, (M_TILE, PAGE, D), qk_box_mma, + self.cluster_layout_vmnk.shape, + ) # fmt: skip + self.q_bytes = cute.size_in_bytes(F8, q_smem_layout) + self.kv_bytes = self.pages_per_tile * cute.size_in_bytes(F8, box_smem_layout) + + ni = self.ni + + @cute.struct + class SharedStorage: + q_full: cute.struct.MemRange[Int64, self.q_stage] + q_empty: cute.struct.MemRange[Int64, self.q_stage] + kv_full: cute.struct.MemRange[Int64, self.kv_stage] + kv_empty: cute.struct.MemRange[Int64, self.kv_stage] + s_full: cute.struct.MemRange[Int64, ni * self.s_stages] + s_empty: cute.struct.MemRange[Int64, ni * self.s_stages] + p_full: cute.struct.MemRange[Int64, ni] + p_empty: cute.struct.MemRange[Int64, ni] + tmem_holding_buf: Int32 + sQ: cute.struct.Align[ + cute.struct.MemRange[F8, cute.cosize(q_smem_layout_staged)], + self.buffer_align_bytes, + ] + sK: cute.struct.Align[ + cute.struct.MemRange[F8, cute.cosize(k_smem_layout_staged)], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + self.kernel( + qk_tiled_mma, + pv_tiled_mma, + qk_box_mma, + tma_atom_q, + tma_tensor_q, + tma_atom_kv, + tma_tensor_kv, + o, + table, + cu_q, + kv_lens, + units, + cta_off, + part_o_ptr, + part_ml_ptr, + stats_ptr, + scale_qk, + scale_o, + table_stride_b, + table_stride_kv, + q_smem_layout_staged, + k_smem_layout_staged, + p_tmem_layout_staged, + v_smem_layout_staged, + box_smem_layout_staged, + ).launch( + grid=(num_ctas, 1, 1), + block=[self.threads_per_cta, 1, 1], + cluster=(1, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + # ------------------------------------------------------------------ device side + @cute.kernel + def kernel( + self, + qk_tiled_mma: cute.TiledMma, + pv_tiled_mma: cute.TiledMma, + qk_box_mma: cute.TiledMma, + tma_atom_q: cute.CopyAtom, + mQ: cute.Tensor, + tma_atom_kv: cute.CopyAtom, + mKV: cute.Tensor, + mO: cute.Tensor, + table: cute.Tensor, + cu_q: cute.Tensor, + kv_lens: cute.Tensor, + units: cute.Tensor, + cta_off: cute.Tensor, + part_o_ptr: cute.Pointer, + part_ml_ptr: cute.Pointer, + stats_ptr: cute.Pointer, + scale_qk: cute.Tensor, + scale_o: cute.Tensor, + table_stride_b: Int32, + table_stride_kv: Int32, + q_smem_layout_staged: cute.ComposedLayout, + k_smem_layout_staged: cute.ComposedLayout, + p_tmem_layout_staged: cute.ComposedLayout, + v_smem_layout_staged: cute.ComposedLayout, + box_smem_layout_staged: cute.ComposedLayout, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + ni = self.ni + + if warp_idx == self.load_warp_id: + cpasync.prefetch_descriptor(tma_atom_q) + cpasync.prefetch_descriptor(tma_atom_kv) + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + q_full = storage.q_full.data_ptr() + q_empty = storage.q_empty.data_ptr() + kv_full = storage.kv_full.data_ptr() + kv_empty = storage.kv_empty.data_ptr() + s_full = storage.s_full.data_ptr() + s_empty = storage.s_empty.data_ptr() + p_full = storage.p_full.data_ptr() + p_empty = storage.p_empty.data_ptr() + if warp_idx == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(self.q_stage): + cute.arch.mbarrier_init(q_full + i, 1) + cute.arch.mbarrier_init(q_empty + i, 1) + for i in cutlass.range_constexpr(self.kv_stage): + cute.arch.mbarrier_init(kv_full + i, 1) + cute.arch.mbarrier_init(kv_empty + i, 1) + for i in cutlass.range_constexpr(ni * self.s_stages): + cute.arch.mbarrier_init(s_full + i, 1) + cute.arch.mbarrier_init(s_empty + i, 128) + for i in cutlass.range_constexpr(ni): + cute.arch.mbarrier_init(p_full + i, 128) + cute.arch.mbarrier_init(p_empty + i, 1) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.mma_warp_id, + arch=self.arch_str, + ) + pipeline_init_arrive(is_relaxed=True) + + sQ = storage.sQ.get_tensor(q_smem_layout_staged.outer, swizzle=q_smem_layout_staged.inner) + sK = storage.sK.get_tensor(k_smem_layout_staged.outer, swizzle=k_smem_layout_staged.inner) + # V and the page boxes are views of the same ring bytes (see the module docstring). + sV = storage.sK.get_tensor(v_smem_layout_staged.outer, swizzle=v_smem_layout_staged.inner) + sBox = storage.sK.get_tensor( + box_smem_layout_staged.outer, swizzle=box_smem_layout_staged.inner + ) + ring_words = cute.make_tensor( + cute.recast_ptr(storage.sK.data_ptr(), dtype=Int32), + cute.make_layout(self.kv_stage * self.slot_words), + ) + + qk_thr_mma = qk_tiled_mma.get_slice(0) + pv_thr_mma = pv_tiled_mma.get_slice(0) + tSrQ = qk_thr_mma.make_fragment_A(sQ) + tSrK = qk_thr_mma.make_fragment_B(sK) + tOrV = pv_thr_mma.make_fragment_B(sV) + gBox = cute.flat_divide(mKV, (PAGE, D)) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_kv, + 0, + cute.make_layout(1), + cute.group_modes(sBox, 0, 3), + cute.group_modes(qk_box_mma.get_slice(0).partition_B(gBox), 0, 3), + ) + + # Every role reads the unit list, cu_q / kv_lens (and the load warp Q, the table and the + # pool) from here on, so every warp waits for the grids this launch depends on. + cute.arch.griddepcontrol_wait() + pipeline_init_wait() + + u_begin = cta_off[bidx] + u_end = cta_off[bidx + 1] + trace = None + step_base = Int32(0) + if cutlass.const_expr(self.trace): + step_base = cute.size(units) // UNIT_FIELDS * TRACE_FIELDS + trace = cute.make_tensor( + cute.recast_ptr(stats_ptr, dtype=Int64), + cute.make_layout(step_base + TRACE_STEP_WORDS), + ) + + # //////////////////////////////////////////////////////////////// IDLE + if warp_idx >= self.load_warp_id + 1: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + # //////////////////////////////////////////////////////////////// LOAD + if warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + lane = cute.arch.lane_idx() + oob_box = cute.size(mKV, mode=[2]) + tpb = self.tiles_per_batch + q_base = Int32(0) + q_ph = Int32(1) + kv_idx = Int32(0) + kv_ph = Int32(1) + q_ring = (tma_atom_q, mQ, sQ, qk_thr_mma, units, cu_q, q_full, q_empty) + if u_begin < u_end: + q_base, q_ph = self.load_q(q_ring, u_begin, q_base, q_ph) + for u in cutlass.range(u_begin, u_end, 1, unroll=1): + if cutlass.const_expr(self.trace): + if lane == 0: + trace[u * TRACE_FIELDS + 4] = cute.arch.globaltimer() + b = units[u * UNIT_FIELDS + 0] + hg = units[u * UNIT_FIELDS + 2] + t0 = units[u * UNIT_FIELDS + 3] + t1 = units[u * UNIT_FIELDS + 4] + kv_len = kv_lens[b] + num_pages = (kv_len + PAGE - 1) // PAGE + kvh = (hg * ni) // self.h_r + # Q of the next unit (this unit's was issued one unit earlier). + if u + 1 < u_end: + q_base, q_ph = self.load_q(q_ring, u + 1, q_base, q_ph) + if cutlass.const_expr(self.trace): + if lane == 0: + trace[u * TRACE_FIELDS + 5] = cute.arch.globaltimer() + # Page ids of key tiles [tpb * bj, tpb * (bj + 1)): one K id and one V id per lane; + # the next batch is read one batch ahead of its use. + bj = t0 // tpb + kid, vid = self.load_page_ids( + table, b, bj, num_pages, lane, table_stride_b, table_stride_kv + ) + kid_n, vid_n = self.load_page_ids( + table, b, bj + 1, num_pages, lane, table_stride_b, table_stride_kv + ) + if cutlass.const_expr(self.trace): + if lane == 0: + # (the ids enter the value so the stamp follows the loads) + trace[u * TRACE_FIELDS + 6] = ( + Int64(kid + kid_n) * Int64(0) + cute.arch.globaltimer() + ) + kv_ring = (tma_atom_kv, tBgB, tBsB, kvh, oob_box, kv_full, kv_empty) + # K(t0), then per tile t: V(t), K(t + 1) (one S stage) or K(t + 1), V(t) (two). + kv_idx, kv_ph = self.load_tile(kv_ring, kid, t0, kv_idx, kv_ph) + for t in cutlass.range(t0, t1, 1, unroll=1): + if t // tpb != bj: + bj = t // tpb + kid = kid_n + vid = vid_n + kid_n, vid_n = self.load_page_ids( + table, b, bj + 1, num_pages, lane, table_stride_b, table_stride_kv + ) + kid_next = cutlass.select_((t + 1) // tpb == bj, kid, kid_n) + if cutlass.const_expr(self.s_stages == 2): + if t + 1 < t1: + kv_idx, kv_ph = self.load_tile(kv_ring, kid_next, t + 1, kv_idx, kv_ph) + kv_idx, kv_ph = self.load_tile(kv_ring, vid, t, kv_idx, kv_ph) + else: + kv_idx, kv_ph = self.load_tile(kv_ring, vid, t, kv_idx, kv_ph) + if t + 1 < t1: + kv_idx, kv_ph = self.load_tile(kv_ring, kid_next, t + 1, kv_idx, kv_ph) + if cutlass.const_expr(self.trace): + if lane == 0: + trace[u * TRACE_FIELDS + 7] = cute.arch.globaltimer() + + # //////////////////////////////////////////////////////////////// MMA + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + tmem.allocate(512) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(Float32) + tStS, tOtO, tOrP = self.make_tmem_tensors( + qk_thr_mma, pv_thr_mma, p_tmem_layout_staged, tmem_ptr + ) + lane = cute.arch.lane_idx() + q_base = Int32(0) + q_ph = Int32(0) + kv_idx = Int32(0) + kv_ph = Int32(0) + p_ph = Int32(0) + # S tiles of this CTA's earlier units (the softmax warps count the same). + n0 = Int32(0) + for u in cutlass.range(u_begin, u_end, 1, unroll=1): + rec = (bidx == 0) & (u == u_begin) + if cutlass.const_expr(self.trace): + if lane == 0: + trace[u * TRACE_FIELDS + 0] = cute.arch.globaltimer() + b = units[u * UNIT_FIELDS + 0] + t0 = units[u * UNIT_FIELDS + 3] + t1 = units[u * UNIT_FIELDS + 4] + kv_len = kv_lens[b] + last_tile = (kv_len - 1) // self.tile_n + for j in cutlass.range_constexpr(ni): + cute.arch.mbarrier_wait(q_full + q_base + j, q_ph) + if cutlass.const_expr(self.trace): + if lane == 0: + trace[u * TRACE_FIELDS + 1] = cute.arch.globaltimer() + # Prologue: S_j(t0) for every instance. + cute.arch.mbarrier_wait(kv_full + kv_idx, kv_ph) + k_slot = kv_idx + kv_idx, kv_ph = self.ring_advance(kv_idx, kv_ph, self.kv_stage) + for j in cutlass.range_constexpr(ni): + self.issue_qk(qk_tiled_mma, tStS[j], tSrQ[None, None, None, q_base + j], + tSrK[None, None, None, k_slot], s_full, s_empty, j, n0) # fmt: skip + self.commit(kv_empty + k_slot) + if cutlass.const_expr(self.trace): + if lane == 0: + trace[u * TRACE_FIELDS + 2] = cute.arch.globaltimer() + for t in cutlass.range(t0, t1, 1, unroll=1): + n = n0 + t - t0 + has_next = t + 1 < t1 + k_slot = kv_idx + v_slot = kv_idx + if cutlass.const_expr(self.s_stages == 2): + if has_next: + cute.arch.mbarrier_wait(kv_full + kv_idx, kv_ph) + kv_idx, kv_ph = self.ring_advance(kv_idx, kv_ph, self.kv_stage) + v_slot = kv_idx + cute.arch.mbarrier_wait(kv_full + kv_idx, kv_ph) + kv_idx, kv_ph = self.ring_advance(kv_idx, kv_ph, self.kv_stage) + else: + cute.arch.mbarrier_wait(kv_full + kv_idx, kv_ph) + kv_idx, kv_ph = self.ring_advance(kv_idx, kv_ph, self.kv_stage) + k_slot = kv_idx + if has_next: + cute.arch.mbarrier_wait(kv_full + kv_idx, kv_ph) + kv_idx, kv_ph = self.ring_advance(kv_idx, kv_ph, self.kv_stage) + if t == last_tile: + # Zero the V rows past the request's last key (any bytes, see the record). + rows_valid = kv_len - t * self.tile_n + if rows_valid < self.tile_n: + n_words = (self.tile_n - rows_valid) * self.row_words + base = v_slot * self.slot_words + rows_valid * self.row_words + for w in cutlass.range(lane, n_words, 32, unroll=1): + ring_words[base + w] = Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + if cutlass.const_expr(self.p_early): + # P has its own buffer: S_j(t + 1) as soon as softmax j has read S_j(t). + if has_next: + for j in cutlass.range_constexpr(ni): + self.issue_qk(qk_tiled_mma, tStS[j], + tSrQ[None, None, None, q_base + j], + tSrK[None, None, None, k_slot], + s_full, s_empty, j, n + 1) # fmt: skip + for j in cutlass.range_constexpr(ni): + if cutlass.const_expr(self.s_stages == 2): + if has_next: + self.issue_qk(qk_tiled_mma, tStS[j], + tSrQ[None, None, None, q_base + j], + tSrK[None, None, None, k_slot], + s_full, s_empty, j, n + 1) # fmt: skip + cute.arch.mbarrier_wait(p_full + j, p_ph) + if cutlass.const_expr(self.trace): + if rec & (lane == 0) & (t - t0 < TRACE_STEPS): + trace[step_base + (8 + j) * TRACE_STEPS + t - t0] = ( + cute.arch.clock64() + ) + self.issue_pv(pv_tiled_mma, tOtO[j], tOrP[j], + tOrV[None, None, None, v_slot], t > t0, n) # fmt: skip + self.commit(p_empty + j) + if cutlass.const_expr(self.s_stages == 1 and not self.p_early): + if has_next: + self.issue_qk(qk_tiled_mma, tStS[j], + tSrQ[None, None, None, q_base + j], + tSrK[None, None, None, k_slot], + s_full, s_empty, j, n + 1) # fmt: skip + if cutlass.const_expr(self.trace): + if rec & (lane == 0) & (t - t0 < TRACE_STEPS): + trace[step_base + (12 + j) * TRACE_STEPS + t - t0] = ( + cute.arch.clock64() + ) + p_ph ^= 1 + self.commit(kv_empty + v_slot) + if has_next: + self.commit(kv_empty + k_slot) + n0 += t1 - t0 + if cutlass.const_expr(self.trace): + if lane == 0: + trace[u * TRACE_FIELDS + 3] = cute.arch.globaltimer() + for j in cutlass.range_constexpr(ni): + self.commit(q_empty + q_base + j) + q_base += ni + if q_base == self.q_stage: + q_base = Int32(0) + q_ph ^= 1 + cute.arch.griddepcontrol_launch_dependents() + tmem.relinquish_alloc_permit() + self.tmem_dealloc_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + + # //////////////////////////////////////////////////////////////// SOFTMAX + if warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_increase(self.num_regs_softmax) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(Float32) + self.softmax( + warp_idx // 4, + tidx % 128, + tmem_ptr, + qk_thr_mma, + pv_thr_mma, + (s_full, s_empty, p_full, p_empty), + (units, u_begin, u_end, cu_q, kv_lens, mO, part_o_ptr, part_ml_ptr, stats_ptr), + (scale_qk, scale_o), + (trace, step_base, bidx), + ) + self.tmem_dealloc_barrier.arrive() + return + + # ------------------------------------------------------------------ helpers + @cute.jit + def ring_advance(self, idx, ph, stages: int): + idx += 1 + if idx == stages: + idx = Int32(0) + ph ^= 1 + return idx, ph + + @cute.jit + def commit(self, mbar_ptr): + with cute.arch.elect_one(): + tcgen05.commit(mbar_ptr) + + @cute.jit + def load_q(self, q_ring, u, q_base, q_ph): + """Copy the Q tiles of unit ``u`` (one per instance) into Q slots ``q_base ..``.""" + tma_atom_q, mQ, sQ, qk_thr_mma, units, cu_q, q_full, q_empty = q_ring + ni = self.ni + b = units[u * UNIT_FIELDS + 0] + m = units[u * UNIT_FIELDS + 1] + hg = units[u * UNIT_FIELDS + 2] + q0 = cu_q[b] + mQ_ = cute.domain_offset((q0 + m * M_TILE, 0, 0), mQ) + gQ = cute.flat_divide(mQ_, (M_TILE, D)) + tSgQ = qk_thr_mma.partition_A(gQ) + tQsQ, tQgQ = cpasync.tma_partition( + tma_atom_q, + 0, + cute.make_layout(1), + cute.group_modes(sQ, 0, 3), + cute.group_modes(tSgQ, 0, 3), + ) + for j in cutlass.range_constexpr(ni): + cute.arch.mbarrier_wait(q_empty + q_base + j, q_ph) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(q_full + q_base + j, self.q_bytes) + cute.copy( + tma_atom_q, + tQgQ[None, 0, 0, hg * ni + j], + tQsQ[None, q_base + j], + tma_bar_ptr=q_full + q_base + j, + ) + q_base += ni + if q_base == self.q_stage: + q_base = Int32(0) + q_ph ^= 1 + return q_base, q_ph + + @cute.jit + def load_tile(self, kv_ring, ids, x, kv_idx, kv_ph): + """Copy key tile ``x`` (K or V, by ``ids``) of the unit into ring slot ``kv_idx``.""" + tma_atom_kv, tBgB, tBsB, kvh, oob_box, kv_full, kv_empty = kv_ring + cute.arch.mbarrier_wait(kv_empty + kv_idx, kv_ph) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(kv_full + kv_idx, self.kv_bytes) + for p in cutlass.range_constexpr(self.pages_per_tile): + pid = cute.arch.shuffle_sync(ids, (x % self.tiles_per_batch) * self.pages_per_tile + p) + box = cutlass.select_(pid >= 0, pid * self.num_kv_heads + kvh, oob_box) + cute.copy( + tma_atom_kv, + tBgB[None, 0, 0, box], + tBsB[None, kv_idx * self.pages_per_tile + p], + tma_bar_ptr=kv_full + kv_idx, + ) + return self.ring_advance(kv_idx, kv_ph, self.kv_stage) + + @cute.jit + def issue_qk(self, tiled_mma, tStS_j, tSrQ, tSrK, s_full, s_empty, j: cutlass.Constexpr, n): + """S tile ``n`` of instance ``j`` into its S buffer ``n % s_stages``, once softmax ``j`` + has read that buffer's previous tile; commits ``s_full``.""" + sb = n % self.s_stages + idx = j * self.s_stages + sb + cute.arch.mbarrier_wait(s_empty + idx, ((n // self.s_stages) & 1) ^ 1) + tS = cute.make_tensor(tStS_j.iterator + sb * self.tile_n, tStS_j.layout) + self.gemm_qk(tiled_mma, tS, tSrQ, tSrK) + self.commit(s_full + idx) + + @cute.jit + def issue_pv(self, tiled_mma, tOtO_j, tOrP_j, tOrV, accumulate, n): + """O_j (+)= P_j(n) V, P_j(n) in S buffer ``n % s_stages``.""" + if cutlass.const_expr(self.s_stages == 2): + if n % 2 == 0: + self.gemm_pv(tiled_mma, tOtO_j, tOrP_j[0], tOrV, accumulate) + else: + self.gemm_pv(tiled_mma, tOtO_j, tOrP_j[1], tOrV, accumulate) + else: + self.gemm_pv(tiled_mma, tOtO_j, tOrP_j[0], tOrV, accumulate) + + @cute.jit + def load_page_ids(self, table, b, bj, num_pages, lane, stride_b, stride_kv): + """Lane ``l``: the K and V blocks of page ``32 bj + l`` of request ``b``, -1 past its last + page.""" + page = bj * IDS_PER_BATCH + lane + kid = Int32(-1) + vid = Int32(-1) + if page < num_pages: + kid = table[b * stride_b + page] + vid = table[b * stride_b + stride_kv + page] + return kid, vid + + def make_tmem_tensors(self, qk_thr_mma, pv_thr_mma, p_tmem_layout_staged, tmem_ptr): + """Per instance: S buffer 0 (fp32, tile_n columns; buffer 1 follows it), O (fp32, 64 + columns, from ``o_off``), P per S buffer (E4M3 from ``p_off`` + the buffer's start).""" + tStS_fake = qk_thr_mma.make_fragment_C( + qk_thr_mma.partition_shape_C((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) + ) + tOtO_fake = pv_thr_mma.make_fragment_C( + pv_thr_mma.partition_shape_C((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) + ) + tP = cute.make_tensor(tmem_ptr, p_tmem_layout_staged.outer) + tOrP_fake = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] + tStS, tOtO, tOrP = [], [], [] + for j in range(self.ni): + base = j * self.inst_cols + tStS.append(cute.make_tensor(tmem_ptr + base, tStS_fake.layout)) + tOtO.append(cute.make_tensor(tmem_ptr + base + self.o_off, tOtO_fake.layout)) + tOrP.append( + [ + cute.make_tensor( + cute.recast_ptr( + tmem_ptr + base + self.p_off + sb * self.tile_n, dtype=tOrP_fake.dtype + ), + tOrP_fake.layout, + ) + for sb in range(self.s_stages) + ] + ) + return tStS, tOtO, tOrP + + def fresh_qk_mma(self): + """A new QK tiled MMA: its runtime fields (accumulate) start from the atom's value in the + region where it is built, so no update from another region is used.""" + return sm100_utils.make_trivial_tiled_mma( + F8, F8, OperandMajorMode.K, OperandMajorMode.K, Float32, tcgen05.CtaGroup.ONE, + self.qk_mma_tiler[:2], + ) # fmt: skip + + def fresh_pv_mma(self): + return sm100_utils.make_trivial_tiled_mma( + F8, F8, OperandMajorMode.K, OperandMajorMode.MN, Float32, tcgen05.CtaGroup.ONE, + self.pv_mma_tiler[:2], tcgen05.OperandSource.TMEM, + ) # fmt: skip + + @cute.jit + def gemm_qk(self, tiled_mma, tStS, tSrQ, tSrK): + tiled_mma = self.fresh_qk_mma() + for k in cutlass.range_constexpr(cute.size(tSrQ, mode=[2])): + tiled_mma.set(tcgen05.Field.ACCUMULATE, k != 0) + cute.gemm(tiled_mma, tStS, tSrQ[None, None, k], tSrK[None, None, k], tStS) + + @cute.jit + def gemm_pv(self, tiled_mma, tOtO, tOrP, tOrV, accumulate): + tiled_mma = self.fresh_pv_mma() + for k in cutlass.range_constexpr(cute.size(tOrP, mode=[2])): + tiled_mma.set(tcgen05.Field.ACCUMULATE, accumulate if k == 0 else True) + cute.gemm(tiled_mma, tOtO, tOrP[None, None, k], tOrV[None, None, k], tOtO) + + @cute.jit + def ex2(self, x): + if cutlass.const_expr(self.ablate_mufu): + return x * Float32(0.5) + return cute.math.exp2(x, fastmath=True) + + # ------------------------------------------------------------------ softmax + @cute.jit + def softmax(self, inst, row, tmem_ptr, qk_thr_mma, pv_thr_mma, bars, work, scales, tr): + """Softmax warpgroup of instance ``inst``; thread ``row`` owns TMEM lane / query row + ``row`` of every unit.""" + s_full, s_empty, p_full, p_empty = bars + units, u_begin, u_end, cu_q, kv_lens, mO, part_o_ptr, part_ml_ptr, stats_ptr = work + s_full_j = s_full + inst * self.s_stages + s_empty_j = s_empty + inst * self.s_stages + p_full_j = p_full + inst + p_empty_j = p_empty + inst + n = self.tile_n + base_col = inst * self.inst_cols + tStS_fake = qk_thr_mma.make_fragment_C( + qk_thr_mma.partition_shape_C((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) + ) + tStS = cute.make_tensor(tmem_ptr + base_col, tStS_fake.layout) + tOtO_fake = pv_thr_mma.make_fragment_C( + pv_thr_mma.partition_shape_C((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) + ) + tOtO = cute.make_tensor(tmem_ptr + base_col + self.o_off, tOtO_fake.layout) + # S loads with the per-32-column max (tcgen05.ld.red.max). + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.LdRed32x32bOp( + tcgen05.copy.Repetition(32), redOp=tcgen05.TmemLoadRedOp.MAX + ), + Float32, + ) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS) + thr_tmem_load = tiled_tmem_load.get_slice(row) + tTMEM_LOADtS = thr_tmem_load.partition_S(tStS) + cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) + tScS = qk_thr_mma.partition_C(cS) + tTMEM_LOADcS = thr_tmem_load.partition_D(tScS) + # P (E4M3, tile_n / 4 columns) from p_off (the start of S unless p_sep). + p_cols = n // 4 + tStS_P = cute.make_tensor( + tStS.iterator + self.p_off, + cute.composition(tStS.layout, cute.make_layout((M_TILE, p_cols))), + ) + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(p_cols)), Float32 + ) + tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) + thr_tmem_store = tiled_tmem_store.get_slice(row) + tTMEM_STOREtP = thr_tmem_store.partition_D(tStS_P) + tScS_P = cute.make_tensor( + tScS.iterator, cute.composition(tScS.layout, cute.make_layout((M_TILE, p_cols))) + ) + tTMEM_STOREcP = thr_tmem_store.partition_S(tScS_P) + n_chunks = cute.size(tTMEM_LOADtS, mode=[1]) + + scale = scales[0][0] + scale_out = scales[1][0] + trace, step_base, bidx = tr + # S tiles of this CTA's earlier units (the MMA warp counts the same); P.V completions + # waited so far (parity). + n0 = Int32(0) + pe_ph = Int32(0) + for u in cutlass.range(u_begin, u_end, 1, unroll=1): + rec = (bidx == 0) & (u == u_begin) & (row == 0) + if cutlass.const_expr(self.trace): + if row == 0: + trace[u * TRACE_FIELDS + 8 + inst] = cute.arch.globaltimer() + b = units[u * UNIT_FIELDS + 0] + m = units[u * UNIT_FIELDS + 1] + hg = units[u * UNIT_FIELDS + 2] + t0 = units[u * UNIT_FIELDS + 3] + t1 = units[u * UNIT_FIELDS + 4] + mode = units[u * UNIT_FIELDS + 5] + slot = units[u * UNIT_FIELDS + 6] + q0 = cu_q[b] + q_len = cu_q[b + 1] - q0 + kv_len = kv_lens[b] + off = kv_len - q_len + r_req = m * M_TILE + row + # Row r keeps keys k <= min(off + r, kv_len - 1); tiles from first_masked on need it. + key_limit = cutlass.min(off + r_req, kv_len - 1) + first_masked = (off + m * M_TILE + 1) // n + t_mid = cutlass.max(t0, cutlass.min(t1, first_masked)) + row_max = -Float32.inf + row_sum = Float32(0.0) + for t in cutlass.range(t0, t_mid, 1, unroll=1): + row_max, row_sum, pe_ph = self.softmax_step( + False, t == t0, t, key_limit, row_max, row_sum, n0 + t - t0, pe_ph, scale, + (s_full_j, s_empty_j, p_full_j, p_empty_j), + (tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADcS, tiled_tmem_store, + tTMEM_STOREtP, tTMEM_STOREcP, n_chunks, tOtO, pv_thr_mma, row), + (trace, step_base, rec, t - t0, inst), + ) # fmt: skip + if cutlass.const_expr(self.trace): + if row == 0: + trace[u * TRACE_FIELDS + 12 + inst] = cute.arch.globaltimer() + for t in cutlass.range(t_mid, t1, 1, unroll=1): + row_max, row_sum, pe_ph = self.softmax_step( + True, t == t0, t, key_limit, row_max, row_sum, n0 + t - t0, pe_ph, scale, + (s_full_j, s_empty_j, p_full_j, p_empty_j), + (tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADcS, tiled_tmem_store, + tTMEM_STOREtP, tTMEM_STOREcP, n_chunks, tOtO, pv_thr_mma, row), + (trace, step_base, rec, t - t0, inst), + ) # fmt: skip + if cutlass.const_expr(self.trace): + if row == 0: + trace[u * TRACE_FIELDS + 16 + inst] = cute.arch.globaltimer() + # Drain: the unit's last P.V of this instance has completed. + cute.arch.mbarrier_wait(p_empty_j, pe_ph) + pe_ph ^= 1 + n0 += t1 - t0 + head = hg * self.ni + inst + r_local = row + row_ok = m * M_TILE + r_local < q_len + inv = Float32(0.0) + if row_sum > Float32(0.0): + inv = Float32(1.0) / row_sum + if mode == 0: + o_row = mO.iterator + ((q0 + m * M_TILE + r_local) * self.num_heads + head) * D + self.store_o(pv_thr_mma, tOtO, o_row, scale_out * inv, row_ok, row) + if cutlass.const_expr(self.write_stats): + if row_ok: + st = cute.make_tensor( + stats_ptr + ((q0 + m * M_TILE + r_local) * self.num_heads + head) * 2, + cute.make_layout(2), + ) + st[0] = row_max * scale * Float32(math.log(2.0)) + st[1] = row_sum * Float32(2.0**-P_PRESCALE_LOG2) + else: + prow = (slot * self.ni + inst) * M_TILE + r_local + self.store_partial(pv_thr_mma, tOtO, part_o_ptr + prow * D, scale_out * inv, row) + ml = cute.make_tensor(part_ml_ptr + prow * 2, cute.make_layout(2)) + ml[0] = row_max * scale + ml[1] = row_sum + if cutlass.const_expr(self.trace): + if row == 0: + trace[u * TRACE_FIELDS + 20 + inst] = cute.arch.globaltimer() + + @cute.jit + def softmax_step( + self, + masked: cutlass.Constexpr, + is_first, + t, + key_limit, + row_max, + row_sum, + n_step, + pe_ph, + scale, + bars, + atoms, + tr, + ): + s_full_j, s_empty_j, p_full_j, p_empty_j = bars + trace, step_base, rec, step, inst = tr + (tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADcS, tiled_tmem_store, tTMEM_STOREtP, + tTMEM_STOREcP, n_chunks, tOtO, pv_thr_mma, row) = atoms # fmt: skip + n = self.tile_n + sb = n_step % self.s_stages + cute.arch.mbarrier_wait(s_full_j + sb, (n_step // self.s_stages) & 1) + if cutlass.const_expr(self.trace): + if rec & (step < TRACE_STEPS): + trace[step_base + inst * TRACE_STEPS + step] = cute.arch.clock64() + rS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, Float32) + rMax = cute.make_rmem_tensor(cute.make_layout((1, n_chunks)), Float32) + tS = cute.make_tensor(tTMEM_LOADtS.iterator + sb * n, tTMEM_LOADtS.layout) + for c in cutlass.range_constexpr(n_chunks): + cute.copy_atom_call( + tiled_tmem_load, tS[None, c, 0, 0], (rS[None, c, 0, 0], rMax[None, c]) + ) + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(s_empty_j + sb) + tile_max = -Float32.inf + if cutlass.const_expr(masked): + # Register i of a 32x32b TMEM load holds column i of this thread's row. + lim = key_limit - t * n + for i in cutlass.range_constexpr(cute.size(rS)): + rS[i] = cutlass.select_(lim < i, -Float32.inf, rS[i]) + tile_max = rS.load().reduce(cute.ReductionOp.MAX, -Float32.inf, 0) + else: + tile_max = rMax.load().reduce(cute.ReductionOp.MAX, -Float32.inf, 0) + old_max = row_max + new_max = cute.arch.fmax(tile_max, row_max) + if new_max == -Float32.inf: + new_max = Float32(0.0) + if cutlass.const_expr(self.rescale_tau > 0.0): + # Keep the running maximum while the tile exceeds it by at most tau (P <= 2^(8 + tau)). + if (new_max - old_max) * scale <= Float32(self.rescale_tau): + new_max = old_max + acc_log2 = scale * (old_max - new_max) + neg = (Float32(0.0) - new_max) * scale + Float32(P_PRESCALE_LOG2) + half = cute.math.exp2(acc_log2, fastmath=True) * Float32(0.5) + sum2 = (row_sum * half, row_sum * half) + rP = cute.make_rmem_tensor(tTMEM_STOREcP.shape, Float32) + rP_i32 = cute.recast_tensor(rP, Int32) + for i in cutlass.range_constexpr(0, cute.size(rS), 4): + a0, a1 = cute.arch.fma_packed_f32x2((rS[i], rS[i + 1]), (scale, scale), (neg, neg)) + a2, a3 = cute.arch.fma_packed_f32x2((rS[i + 2], rS[i + 3]), (scale, scale), (neg, neg)) + k_emu = self.ex2_emu_every + if cutlass.const_expr(k_emu > 0 and (i // 4) % k_emu == k_emu - 1): + e0, e1 = _ex2_poly_f32x2(a0, a1) + e2, e3 = _ex2_poly_f32x2(a2, a3) + else: + e0 = self.ex2(a0) + e1 = self.ex2(a1) + e2 = self.ex2(a2) + e3 = self.ex2(a3) + sum2 = cute.arch.add_packed_f32x2(sum2, (e0, e1)) + sum2 = cute.arch.add_packed_f32x2(sum2, (e2, e3)) + rP_i32[i // 4] = _cvt_f32x4_to_e4m3x4(e0, e1, e2, e3) + if not is_first: + # P.V of the previous step has completed: O is stable until this step's P arrives. + cute.arch.mbarrier_wait(p_empty_j, pe_ph) + pe_ph ^= 1 + if cutlass.const_expr(not self.ablate_rescale): + if cute.arch.vote_ballot_sync(old_max != new_max) != 0: + self.rescale_o(pv_thr_mma, tOtO, cute.math.exp2(acc_log2, fastmath=True), row) + tP = cute.make_tensor(tTMEM_STOREtP.iterator + sb * n, tTMEM_STOREtP.layout) + cute.copy(tiled_tmem_store, rP, tP) + cute.arch.fence_view_async_tmem_store() + cute.arch.mbarrier_arrive(p_full_j) + if cutlass.const_expr(self.trace): + if rec & (step < TRACE_STEPS): + trace[step_base + (4 + inst) * TRACE_STEPS + step] = cute.arch.clock64() + return new_max, sum2[0] + sum2[1], pe_ph + + @cute.jit + def rescale_o(self, thr_mma, tOtO, scale, row): + """O *= scale for this thread's row (``rescale_cols`` columns at a time through registers).""" + corr = self.rescale_cols + ld_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr)), Float32 + ) + st_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr)), Float32 + ) + cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) + tOcO = thr_mma.partition_C(cO) + tOtO_i = cute.make_tensor( + tOtO.iterator, cute.composition(tOtO.layout, cute.make_layout((M_TILE, corr))) + ) + tOcO_i = cute.make_tensor( + tOcO.iterator, cute.composition(tOcO.layout, cute.make_layout((M_TILE, corr))) + ) + tiled_ld = tcgen05.make_tmem_copy(ld_atom, tOtO_i) + tiled_st = tcgen05.make_tmem_copy(st_atom, tOtO_i) + thr_ld = tiled_ld.get_slice(row) + thr_st = tiled_st.get_slice(row) + tLDtO = thr_ld.partition_S(tOtO_i) + tLDcO = thr_ld.partition_D(tOcO_i) + tSTtO = thr_st.partition_D(tOtO_i) + for i in cutlass.range_constexpr(0, D // corr): + r = cute.make_rmem_tensor(tLDcO.shape, Float32) + cute.copy(tiled_ld, cute.make_tensor(tLDtO.iterator + i * corr, tLDtO.layout), r) + for k in cutlass.range_constexpr(0, cute.size(r), 2): + r[k], r[k + 1] = cute.arch.mul_packed_f32x2((r[k], r[k + 1]), (scale, scale)) + cute.copy(tiled_st, r, cute.make_tensor(tSTtO.iterator + i * corr, tSTtO.layout)) + + @cute.jit + def store_o(self, thr_mma, tOtO, dst_ptr, scale, row_ok, row): + """O * scale to BF16 into this thread's 64-element output row at ``dst_ptr`` (16 columns + per TMEM load, 32-byte row pieces).""" + corr = 16 + ld_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr)), Float32 + ) + cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) + tOcO = thr_mma.partition_C(cO) + tOtO_i = cute.make_tensor( + tOtO.iterator, cute.composition(tOtO.layout, cute.make_layout((M_TILE, corr))) + ) + tOcO_i = cute.make_tensor( + tOcO.iterator, cute.composition(tOcO.layout, cute.make_layout((M_TILE, corr))) + ) + tiled_ld = tcgen05.make_tmem_copy(ld_atom, tOtO_i) + thr_ld = tiled_ld.get_slice(row) + tLDtO = thr_ld.partition_S(tOtO_i) + tLDcO = thr_ld.partition_D(tOcO_i) + for i in cutlass.range_constexpr(0, D // corr): + r = cute.make_rmem_tensor(tLDcO.shape, Float32) + cute.copy(tiled_ld, cute.make_tensor(tLDtO.iterator + i * corr, tLDtO.layout), r) + for k in cutlass.range_constexpr(0, cute.size(r), 2): + r[k], r[k + 1] = cute.arch.mul_packed_f32x2((r[k], r[k + 1]), (scale, scale)) + rb = cute.make_rmem_tensor(cute.make_layout(corr), BF16) + rb.store(r.load().to(BF16)) + if row_ok: + # A row starts at a multiple of 64 elements (128 B), a piece at 32 B. + dst = cute.make_tensor((dst_ptr + i * corr).align(32), cute.make_layout(corr)) + cute.autovec_copy(rb, dst) + # Every O read of this unit is complete before this thread's next P commit. + cute.arch.fence_view_async_tmem_load() + + @cute.jit + def store_partial(self, thr_mma, tOtO, dst_ptr, scale, row): + """O * scale in fp32 to this thread's 64-float workspace row at ``dst_ptr``.""" + corr = 16 + ld_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr)), Float32 + ) + cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) + tOcO = thr_mma.partition_C(cO) + tOtO_i = cute.make_tensor( + tOtO.iterator, cute.composition(tOtO.layout, cute.make_layout((M_TILE, corr))) + ) + tOcO_i = cute.make_tensor( + tOcO.iterator, cute.composition(tOcO.layout, cute.make_layout((M_TILE, corr))) + ) + tiled_ld = tcgen05.make_tmem_copy(ld_atom, tOtO_i) + thr_ld = tiled_ld.get_slice(row) + tLDtO = thr_ld.partition_S(tOtO_i) + tLDcO = thr_ld.partition_D(tOcO_i) + for i in cutlass.range_constexpr(0, D // corr): + r = cute.make_rmem_tensor(tLDcO.shape, Float32) + cute.copy(tiled_ld, cute.make_tensor(tLDtO.iterator + i * corr, tLDtO.layout), r) + for k in cutlass.range_constexpr(0, cute.size(r), 2): + r[k], r[k + 1] = cute.arch.mul_packed_f32x2((r[k], r[k + 1]), (scale, scale)) + dst = cute.make_tensor(dst_ptr + i * corr, cute.make_layout(r.shape)) + cute.autovec_copy(r, dst) + cute.arch.fence_view_async_tmem_load() From 0db4fd5e090158e271ef6d72651eecf3bdad39a6 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian Date: Mon, 9 Nov 2026 21:08:36 -0800 Subject: [PATCH 2/9] musespark gb300: the ported decode stack on by default at TP8 Thirteen levers of the sm103 decode stack move from knob-only to on at TP8, via a _tep8_only rule. At TP4 and every other size they stay knob-only: those compositions either do not serve these paths or have not been tuned for them. What carries the default is acceptance, measured on the SPEED-Bench 80-row coding set (chat-templated, use_rejection_sampling, DFlash d7, concurrency 1, GB300 TEP8): with all thirteen on, the run is token-identical to the all-off stack -- AL 4.6303, the same 3623/6986 accepted/drafted, the same output. Every arm of that gate repeats bitwise, so the result is exact rather than inside a band, and token identity makes it independent of the baseline's absolute level. Two levers were deliberately left knob-only: * drafter_fused_prep resolved off by dependency in BOTH measured arms, so the bitwise-clean result does not cover it. * ring_buddy_stats is measured clean but is observability only, and it is mutually exclusive with tep8_experts_in_kernel_placement. TRTLLM_MULTI_INSTANCE_CONTEXT_FMHA is not a lever and is not enabled by this change. On the same gate it costs 3.79 % AL -- all of the loss the full stack showed. The decode step measurement quoted for this stack (12.23 -> 11.04 ms at bs4) was taken with that switch on, so the levers' own share of it is not separated and is not claimed here. Signed-off-by: Iman Tabrizian --- .../levers.py | 90 +++++++++++++++---- 1 file changed, 73 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py index 0443bc39e71f..80d0f9b9069e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py @@ -444,6 +444,34 @@ def _knob_only(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bo return False, "flagged: on only by its knob, until it is gated on this target" +def _tep8_only(what: str) -> RuleFn: + """On at TP8, the size this path was written and measured for; off elsewhere by its knob only. + + The ported decode stack was built for the sm103 TP8 composition over two trays. At TP4 these + paths are either unserved by the composition or untuned, so they stay knob-only there. + + What carries this default is the acceptance evidence: on the SPEED-Bench 80-row coding set + (chat-templated, rejection sampling, DFlash d7, concurrency 1) every lever listed below runs + **token-identical** to the stack with all of them off -- 4.6303 AL, the same 3623/6986 + accepted/drafted, on GB300 TEP8. Every arm there repeats bitwise, so that is an exact result, + not one inside a noise band; and because the emitted tokens are identical, it does not depend + on the baseline's absolute level. A lever is default-on here only if it was in a measured arm. + + The decode-step speedup measured for this stack (12.23 -> 11.04 ms at bs4) was taken with + ``TRTLLM_MULTI_INSTANCE_CONTEXT_FMHA=1`` also set. That switch is **not** a lever and is not + enabled here: it is the one component that moves acceptance (-3.79 % AL on the same gate, all + of the loss). So the levers' own share of that speedup is not yet separated, and this + docstring does not claim it. + """ + + def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + if d.tp == 8: + return True, what + return False, f"tep{d.tp}: on only by its knob (written for the tep8 composition)" + + return rule + + def _tep4_only(what: str) -> RuleFn: """On at TP4, where the composition carries the path; off at any other size by its knob only. @@ -460,6 +488,11 @@ def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, st return rule +def _ported_tep8(kind: Kind, what: str, params: Tuple[Param, ...] = ()) -> Rule: + """A ported switch that is ON by default at TP8 -- the parallelism it was written for.""" + return Rule("LL", kind, _tep8_only(what), params=params) + + def _ported(kind: Kind, params: Tuple[Param, ...] = ()) -> Rule: """A switch of the sm103 decode stack, ported from an environment variable: default off. @@ -606,37 +639,56 @@ def _vocab_parallel_head(d: Deployment, t: Resolving, p: Mapping[str, object]) - params=(Param("TP8_MAX_ROWS", FINALIZE_TWO_SHOT_TP8_ROWS, 24, 152),), ), # The sm103 decode-stack switches, ported from environment variables: default off. - "tep8_fc_shard_wide": _ported( - Kind.FLAGGED, (Param("MAX_ROWS", FC_SHARD_TP8_ROWS, 9, 128),) + "tep8_fc_shard_wide": _ported_tep8( + Kind.FLAGGED, + "the TP8 sharded drafter fc over the tray's NVLink", + (Param("MAX_ROWS", FC_SHARD_TP8_ROWS, 9, 128),), + ), + "tep8_persistent_batch_experts": _ported_tep8( + Kind.FLAGGED, + "the persistent expert kernel's batch builds at TP8", + (Param("MAX_ROWS", PERSISTENT_BATCH_ROWS, 9, 32),), + ), + "tep8_experts_early_start": _ported_tep8( + Kind.FLAGGED, "the stacked expert call launched with its placement's early trigger" + ), + "tep8_router_shard": _ported_tep8(Kind.FLAGGED, "the TP8 wide-prologue router shard"), + "tep8_router_shard_wide": _ported_tep8( + Kind.FLAGGED, "router_shard_wide's built-in bits on the TP8 wide prologue" ), - "tep8_persistent_batch_experts": _ported( - Kind.FLAGGED, (Param("MAX_ROWS", PERSISTENT_BATCH_ROWS, 9, 32),) + "tep8_wide_direct_pool": _ported_tep8( + Kind.FLAGGED, "pool_into_route's TP8 wide direct sites" ), - "tep8_experts_early_start": _ported(Kind.FLAGGED), - "tep8_router_shard": _ported(Kind.FLAGGED), - "tep8_router_shard_wide": _ported(Kind.FLAGGED), - "tep8_wide_direct_pool": _ported(Kind.FLAGGED), "tep8_wide_direct_transports": _ported( Kind.FLAGGED, (Param("SITES", WIDE_DIRECT_SITES, choices=WIDE_DIRECT_SITES),) ), - "tep8_batch_splitk": _ported(Kind.BITWISE), - "finalize_entry_trigger": _ported(Kind.BITWISE), - "tep8_finalize_transport_wide": _ported( - Kind.FLAGGED, (Param("MAX_TOKENS", FINALIZE_TRANSPORT_TOKENS, 17, 40),) + "tep8_batch_splitk": _ported_tep8(Kind.BITWISE, "the TP8 batch split-K projection kernels"), + "finalize_entry_trigger": _ported_tep8( + Kind.BITWISE, "the finalize transport's entry trigger" + ), + "tep8_finalize_transport_wide": _ported_tep8( + Kind.FLAGGED, + "the widened TP8 finalize transport", + (Param("MAX_TOKENS", FINALIZE_TRANSPORT_TOKENS, 17, 40),), ), "multi_request_combine_gate": _ported(Kind.BITWISE), "router_staged_gemm": _ported( Kind.FLAGGED, (Param("BUILD", (ROUTER_STAGED_BUILD,), choices=ROUTER_STAGED_BUILDS),), ), + # NOT default-on: it resolved off (dependency) in BOTH measured arms, so the + # bitwise-clean SPEED-Bench result does not cover it. Knob-only until it does. "drafter_fused_prep": _ported( Kind.FLAGGED, (Param("MIN_REQUESTS", DRAFTER_PREP_MIN_REQUESTS, 1, 64),) ), - "tep8_ring_buddy_wide": _ported( - Kind.FLAGGED, (Param("MAX_ROWS", RING_BUDDY_MAX_ROWS, 1, RING_BUDDY_ROW_CEILING),) + "tep8_ring_buddy_wide": _ported_tep8( + Kind.FLAGGED, + "the TP8 ring-buddy placement's wide row cap", + (Param("MAX_ROWS", RING_BUDDY_MAX_ROWS, 1, RING_BUDDY_ROW_CEILING),), ), - "tep8_wide_balance": _ported( + "tep8_wide_balance": _ported_tep8( Kind.FLAGGED, + "the 256-row wide balance", (Param("MAX_ROWS", WIDE_BALANCE_ROWS, WIDE_BALANCE_MIN_ROWS, RING_BUDDY_ROW_CEILING),), ), # Diagnostics: one extra assign launch per placement, and the in-kernel placement @@ -646,8 +698,12 @@ def _vocab_parallel_head(d: Deployment, t: Resolving, p: Mapping[str, object]) - Kind.FLAGGED, (Param("INTERVAL_S", RING_BUDDY_STATS_INTERVAL_S, 1, 86400),), ), - "tep8_experts_in_kernel_placement": _ported(Kind.FLAGGED), - "tep8_experts_interleaved_fetch": _ported(Kind.FLAGGED), + "tep8_experts_in_kernel_placement": _ported_tep8( + Kind.FLAGGED, "placement computed inside the expert kernel" + ), + "tep8_experts_interleaved_fetch": _ported_tep8( + Kind.FLAGGED, "the expert kernel's interleaved weight fetch" + ), **{name: _tep4_path_rule() for name in TEP4_PATHS}, } ) From f6a8a2e4dc150db7ae54b8bbc377b3c4c3c705fa Mon Sep 17 00:00:00 2001 From: Iman Tabrizian Date: Mon, 9 Nov 2026 23:02:27 -0800 Subject: [PATCH 3/9] musespark gb300: satisfy the repo's ruff-format and pinned-memory hooks CI runs pre-commit over the whole tree while a local run only sees staged files, so two repo hooks were failing on this branch from the port commit on. * ruff-format reformats six files. Verified format-only: each one's AST is identical before and after, so nothing changed meaning. * The pinned-memory policy disallows raw pinning in runtime code. The ring-buddy placement-stats accumulator built its host buffer with torch.zeros(..., pin_memory=True); it now goes through tensorrt_llm._utils.maybe_pin_memory(), which is what the policy asks for and is idempotent. That buffer belongs to ring_buddy_stats, which is knob-only, so no default path is affected. Signed-off-by: Iman Tabrizian --- .../kernels/projection_publish/publish.py | 4 +- .../kernels/replicated_experts/ring_buddy.py | 3 +- .../wide_direct_transports/transports.py | 23 +- .../kernels/wide_prologue/prologue.py | 4 +- .../ops/attention/swa_block.py | 8 +- .../ops/attention/swa_cluster_block.py | 998 ++++++++++++++---- .../ops/spec/drafter_step_prep.py | 12 +- 7 files changed, 833 insertions(+), 219 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_publish/publish.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_publish/publish.py index faec2688e963..3f6d813b1a2c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_publish/publish.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/projection_publish/publish.py @@ -66,9 +66,7 @@ MIN_ROWS, MAX_ROWS = _publish.MIN_ROWS, 448 -def _pool_args( - moe, max_rows: int = MAX_ROWS -) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: +def _pool_args(moe, max_rows: int = MAX_ROWS) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: """(gain, global scale) exactly when the built-in ``prepare_expert_input`` would quantize a BF16 ``[1..max_rows, 4096]`` pool output with flashinfer's fused ``rmsnorm_fp4quant``; else None. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py index 0bbfee1005d3..a357e0676020 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py @@ -108,6 +108,7 @@ from tensorrt_llm._torch.distributed import allgather from tensorrt_llm._torch.utils import Fp4QuantizedTensor +from tensorrt_llm._utils import maybe_pin_memory from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -1515,7 +1516,7 @@ def __init__(self, interval: float, split: int, ranks: int, rank: int): self.rank = rank self.layers = 0 self.failed = False - self.acc = torch.zeros(2 * PLACEMENT_STATS_FIELDS, dtype=torch.int64, pin_memory=True) + self.acc = maybe_pin_memory(torch.zeros(2 * PLACEMENT_STATS_FIELDS, dtype=torch.int64)) def add(self, stats: torch.Tensor, rows: int) -> None: if self.failed: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/transports.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/transports.py index 8794cd93ab37..140e02909959 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/transports.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_direct_transports/transports.py @@ -109,7 +109,9 @@ def _cells(plan, site: str) -> List[str]: return [table[rows - WIDE_MIN_ROWS] for rows in ROWS] -def refusal(model, plan, wide_projections, unpool_direct, sites: FrozenSet[str], capability) -> Optional[str]: +def refusal( + model, plan, wide_projections, unpool_direct, sites: FrozenSet[str], capability +) -> Optional[str]: """Why the path cannot run here, or None. Reads construction-time state only (rank-agnostic).""" layers = list(model.layers) mapping = layers[0].mapping @@ -129,7 +131,9 @@ def refusal(model, plan, wide_projections, unpool_direct, sites: FrozenSet[str], f"the {rows}-row {site} cell {cell} is not a flashinfer kernel, which alone stores " "into a transport slot" ) - if "unpool" in sites and (unpool_direct is None or len(getattr(unpool_direct, "sites", ())) < 2): + if "unpool" in sites and ( + unpool_direct is None or len(getattr(unpool_direct, "sites", ())) < 2 + ): return "unpool_into_boundary (whose joins, sites and buffers the unpool site rides) is not installed" return None @@ -222,7 +226,10 @@ def _build_pool(self) -> None: chosen: List[int] = [] for index, layer in enumerate(self._layers): moe = layer.feed_forward - if isinstance(moe, MuseSparkOpenTepGatherMoE) and _pool_args(moe, WIDE_MAX_ROWS) is not None: + if ( + isinstance(moe, MuseSparkOpenTepGatherMoE) + and _pool_args(moe, WIDE_MAX_ROWS) is not None + ): chosen.append(index) local = tuple(chosen) everyone = Distributed.get(self.mapping).tp_allgather(local, small_payload=True) @@ -235,7 +242,9 @@ def _build_pool(self) -> None: reason = pool_owner_conflict(self._layers[index] for index in chosen) if reason is None: for index in chosen: - dispatcher = getattr(self._layers[index].feed_forward.pooler, "_low_m_gemm_dispatcher", None) + dispatcher = getattr( + self._layers[index].feed_forward.pooler, "_low_m_gemm_dispatcher", None + ) held = getattr(dispatcher, "publish_policy", None) if dispatcher is None or not hasattr(dispatcher, "publish"): reason = f"layer {index}'s pool projection is not on the dense GEMM dispatcher" @@ -317,7 +326,11 @@ def project_pool(self, moe, module, hidden_states: torch.Tensor) -> Optional[tor if buffer is None or module is not moe.pooler: return None rows = hidden_states.shape[0] if hidden_states.dim() == 2 else -1 - if rows != self._rows or hidden_states.dtype != torch.bfloat16 or not hidden_states.is_contiguous(): + if ( + rows != self._rows + or hidden_states.dtype != torch.bfloat16 + or not hidden_states.is_contiguous() + ): raise RuntimeError( f"[musespark_open] wide_direct_transports: armed for {self._rows} rows, the pool " f"projection got {tuple(hidden_states.shape)} {hidden_states.dtype}" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py index c541d665db18..8fbd14b36b2b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py @@ -237,7 +237,9 @@ def tep8_refusal(model: MuseSparkOpenModel) -> Optional[str]: f"{layer.layer_idx} has {type(moe).__name__} without its inbox" ) if moe._router_stream is None: - return f"the TEP8 router shard runs on the router fork; layer {layer.layer_idx} has none" + return ( + f"the TEP8 router shard runs on the router fork; layer {layer.layer_idx} has none" + ) if getattr(moe, "_wide", None) is None or moe._wide_pool_inbox is None: return ( "the TEP8 router shard takes the pool slices wide_projection_gathers defers; layer " diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py index 20a92629ea36..185cb6278eb3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py @@ -351,7 +351,9 @@ def _swa_split_kernel( mask=new[:, None], other=0.0, ) - e_wk = _prep_rows_load(qkv_ptr, tok * qkv_stride, tl.zeros([8], tl.int32) + k_col, write_ok, 8) + e_wk = _prep_rows_load( + qkv_ptr, tok * qkv_stride, tl.zeros([8], tl.int32) + k_col, write_ok, 8 + ) e_wv = tl.load( qkv_ptr + tok[:, None] * qkv_stride + v_col + d[None, :], mask=write_ok[:, None], @@ -1113,7 +1115,9 @@ def _swa_combine_kernel( if EARLY_LOADS: g = g_early.to(tl.float32) else: - g = tl.load(gate_ptr + t * gate_token_stride + head * gate_head_stride + d).to(tl.float32) + g = tl.load(gate_ptr + t * gate_token_stride + head * gate_head_stride + d).to( + tl.float32 + ) sigmoid = tl.sigmoid(g).to(tl.bfloat16) result = normed.to(tl.float32) * sigmoid.to(tl.float32) tl.store(out_ptr + t * out_stride + head * D + d, result) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_cluster_block.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_cluster_block.py index 5d9019ba20d4..39109d885dae 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_cluster_block.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_cluster_block.py @@ -110,8 +110,17 @@ def owned_splits(rank: int) -> list: return [rank + CLUSTER * k for k in range(NB_MAX) if rank + CLUSTER * k < NUM_SPLITS] -def fits(rows: int, reqs: int, window: int, num_splits: int, fused_prep: bool, q_fp8: bool, has_gate: bool, - keys: int, tile_rows: int) -> Optional[str]: +def fits( + rows: int, + reqs: int, + window: int, + num_splits: int, + fused_prep: bool, + q_fp8: bool, + has_gate: bool, + keys: int, + tile_rows: int, +) -> Optional[str]: """``None`` when a multi-request call is one this block computes bit for bit, else why not.""" if keys != BLOCK or tile_rows != ROWS: return f"tile {keys} keys x {tile_rows} rows (needs {BLOCK} x {ROWS})" @@ -128,198 +137,362 @@ def fits(rows: int, reqs: int, window: int, num_splits: int, fused_prep: bool, q AD = llvm.AsmDialect.AD_ATT + def _v(x, loc, ip): return x.ir_value(loc=loc, ip=ip) + def _asm(ret, args, text, cons, side): - return llvm.inline_asm(ret, args, text, cons, has_side_effects=side, is_align_stack=False, asm_dialect=AD) + return llvm.inline_asm( + ret, args, text, cons, has_side_effects=side, is_align_stack=False, asm_dialect=AD + ) + # ------------------------------------------------------------------------------------------ PTX helpers @dsl_user_op def f1(text, a, *, loc=None, ip=None): return Float32(_asm(T.f32(), [_v(Float32(a), loc, ip)], text, "=f,f", False)) + @dsl_user_op def f2(text, a, b, *, loc=None, ip=None): - return Float32(_asm(T.f32(), [_v(Float32(a), loc, ip), _v(Float32(b), loc, ip)], text, "=f,f,f", False)) + return Float32( + _asm(T.f32(), [_v(Float32(a), loc, ip), _v(Float32(b), loc, ip)], text, "=f,f,f", False) + ) + @dsl_user_op def f3(text, a, b, c, *, loc=None, ip=None): - return Float32(_asm(T.f32(), [_v(Float32(a), loc, ip), _v(Float32(b), loc, ip), _v(Float32(c), loc, ip)], - text, "=f,f,f,f", False)) + return Float32( + _asm( + T.f32(), + [_v(Float32(a), loc, ip), _v(Float32(b), loc, ip), _v(Float32(c), loc, ip)], + text, + "=f,f,f,f", + False, + ) + ) + def add(a, b): return f2("add.rn.f32 $0, $1, $2;", a, b) + def add_ftz(a, b): return f2("add.rn.ftz.f32 $0, $1, $2;", a, b) + def sub(a, b): return f2("sub.rn.f32 $0, $1, $2;", a, b) + def mul(a, b): return f2("mul.rn.f32 $0, $1, $2;", a, b) + def mul_ftz(a, b): return f2("mul.rn.ftz.f32 $0, $1, $2;", a, b) + def fma(a, b, c): return f3("fma.rn.f32 $0, $1, $2, $3;", a, b, c) + def fma_ftz(a, b, c): return f3("fma.rn.ftz.f32 $0, $1, $2, $3;", a, b, c) + def div_full(a, b): return f2("div.full.f32 $0, $1, $2;", a, b) + def fmax(a, b): return f2("max.f32 $0, $1, $2;", a, b) + def rsqrt_ftz(a): return f1("rsqrt.approx.ftz.f32 $0, $1;", a) + def ex2_ftz(a): return f1("ex2.approx.ftz.f32 $0, $1;", a) + def neg(a): return f1("neg.f32 $0, $1;", a) + def bf16r(a): """fp32 -> BF16 (round to nearest even) -> fp32.""" return f1("{\n\t.reg .b16 h;\n\tcvt.rn.bf16.f32 h, $1;\n\tcvt.f32.bf16 $0, h;\n\t}", a) + def bf16_mul(a, b): """``mul.bf16`` of two BF16-exact fp32 values, as fp32 (the served ``(x * s).to(bf16)``).""" - return f2("{\n\t.reg .b16 x, y, z;\n\tcvt.rn.bf16.f32 x, $1;\n\tcvt.rn.bf16.f32 y, $2;\n\t" - "mul.rn.bf16 z, x, y;\n\tcvt.f32.bf16 $0, z;\n\t}", a, b) + return f2( + "{\n\t.reg .b16 x, y, z;\n\tcvt.rn.bf16.f32 x, $1;\n\tcvt.rn.bf16.f32 y, $2;\n\t" + "mul.rn.bf16 z, x, y;\n\tcvt.f32.bf16 $0, z;\n\t}", + a, + b, + ) + def sigmoid_bf16(g): """The combine's ``tl.sigmoid(g).to(bf16)``: ``1 / (1 + ex2.approx(g * -log2 e))``, ``div.full``.""" - return f1("{\n\t.reg .f32 x, e;\n\tmul.rn.f32 x, $1, 0fBFB8AA3B;\n\tex2.approx.f32 e, x;\n\t" - "add.rn.f32 e, e, 0f3F800000;\n\tdiv.full.f32 x, 0f3F800000, e;\n\t.reg .b16 h;\n\t" - "cvt.rn.bf16.f32 h, x;\n\tcvt.f32.bf16 $0, h;\n\t}", g) + return f1( + "{\n\t.reg .f32 x, e;\n\tmul.rn.f32 x, $1, 0fBFB8AA3B;\n\tex2.approx.f32 e, x;\n\t" + "add.rn.f32 e, e, 0f3F800000;\n\tdiv.full.f32 x, 0f3F800000, e;\n\t.reg .b16 h;\n\t" + "cvt.rn.bf16.f32 h, x;\n\tcvt.f32.bf16 $0, h;\n\t}", + g, + ) + @dsl_user_op def e4m3x2(hi, lo, *, loc=None, ip=None): """``cvt.rn.satfinite.e4m3x2.f32`` (``hi`` in the upper byte) as the low 16 bits of an Int32.""" - return Int32(_asm(T.i32(), [_v(Float32(hi), loc, ip), _v(Float32(lo), loc, ip)], - "{\n\t.reg .b16 q;\n\tcvt.rn.satfinite.e4m3x2.f32 q, $1, $2;\n\tcvt.u32.u16 $0, q;\n\t}", - "=r,f,f", False)) + return Int32( + _asm( + T.i32(), + [_v(Float32(hi), loc, ip), _v(Float32(lo), loc, ip)], + "{\n\t.reg .b16 q;\n\tcvt.rn.satfinite.e4m3x2.f32 q, $1, $2;\n\tcvt.u32.u16 $0, q;\n\t}", + "=r,f,f", + False, + ) + ) + @dsl_user_op def e4m3x2_to_bf16x2(q, *, loc=None, ip=None): """Two E4M3 (the low 16 bits of ``q``, element 0 in the low byte) -> BF16x2 (exact).""" - return Int32(_asm(T.i32(), [_v(Int32(q), loc, ip)], - "{\n\t.reg .b16 q, a, b;\n\t.reg .b32 h;\n\t.reg .f32 x, y;\n\tcvt.u16.u32 q, $1;\n\t" - "cvt.rn.f16x2.e4m3x2 h, q;\n\tmov.b32 {a, b}, h;\n\tcvt.f32.f16 x, a;\n\t" - "cvt.f32.f16 y, b;\n\tcvt.rn.bf16x2.f32 $0, y, x;\n\t}", "=r,r", False)) + return Int32( + _asm( + T.i32(), + [_v(Int32(q), loc, ip)], + "{\n\t.reg .b16 q, a, b;\n\t.reg .b32 h;\n\t.reg .f32 x, y;\n\tcvt.u16.u32 q, $1;\n\t" + "cvt.rn.f16x2.e4m3x2 h, q;\n\tmov.b32 {a, b}, h;\n\tcvt.f32.f16 x, a;\n\t" + "cvt.f32.f16 y, b;\n\tcvt.rn.bf16x2.f32 $0, y, x;\n\t}", + "=r,r", + False, + ) + ) + @dsl_user_op def lo_bf16(w, *, loc=None, ip=None): return Float32(_asm(T.f32(), [_v(Int32(w), loc, ip)], "shl.b32 $0, $1, 16;", "=f,r", False)) + @dsl_user_op def hi_bf16(w, *, loc=None, ip=None): - return Float32(_asm(T.f32(), [_v(Int32(w), loc, ip)], "and.b32 $0, $1, 0xFFFF0000;", "=f,r", False)) + return Float32( + _asm(T.f32(), [_v(Int32(w), loc, ip)], "and.b32 $0, $1, 0xFFFF0000;", "=f,r", False) + ) + @dsl_user_op def pack_bf16x2(hi, lo, *, loc=None, ip=None): """``cvt.rn.bf16x2.f32`` (``lo`` in the low half).""" - return Int32(_asm(T.i32(), [_v(Float32(hi), loc, ip), _v(Float32(lo), loc, ip)], - "cvt.rn.bf16x2.f32 $0, $1, $2;", "=r,f,f", False)) + return Int32( + _asm( + T.i32(), + [_v(Float32(hi), loc, ip), _v(Float32(lo), loc, ip)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + False, + ) + ) + @dsl_user_op def bits(x, *, loc=None, ip=None): return Int32(_asm(T.i32(), [_v(Float32(x), loc, ip)], "mov.b32 $0, $1;", "=r,f", False)) + @dsl_user_op def as_f32(x, *, loc=None, ip=None): return Float32(_asm(T.f32(), [_v(Int32(x), loc, ip)], "mov.b32 $0, $1;", "=f,r", False)) + @dsl_user_op def selp(a, b, p, *, loc=None, ip=None): """``p != 0 ? a : b`` (fp32).""" - return Float32(_asm(T.f32(), [_v(Float32(a), loc, ip), _v(Float32(b), loc, ip), _v(Int32(p), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\tselp.f32 $0, $1, $2, q;\n\t}", - "=f,f,f,r", False)) + return Float32( + _asm( + T.f32(), + [_v(Float32(a), loc, ip), _v(Float32(b), loc, ip), _v(Int32(p), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\tselp.f32 $0, $1, $2, q;\n\t}", + "=f,f,f,r", + False, + ) + ) + @dsl_user_op def isel(a, b, p, *, loc=None, ip=None): - return Int32(_asm(T.i32(), [_v(Int32(a), loc, ip), _v(Int32(b), loc, ip), _v(Int32(p), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\tselp.b32 $0, $1, $2, q;\n\t}", - "=r,r,r,r", False)) + return Int32( + _asm( + T.i32(), + [_v(Int32(a), loc, ip), _v(Int32(b), loc, ip), _v(Int32(p), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\tselp.b32 $0, $1, $2, q;\n\t}", + "=r,r,r,r", + False, + ) + ) + @dsl_user_op def lt(a, b, *, loc=None, ip=None): """``a < b`` (signed) as 0 / 1.""" - return Int32(_asm(T.i32(), [_v(Int32(a), loc, ip), _v(Int32(b), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.lt.s32 q, $1, $2;\n\tselp.u32 $0, 1, 0, q;\n\t}", - "=r,r,r", False)) + return Int32( + _asm( + T.i32(), + [_v(Int32(a), loc, ip), _v(Int32(b), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.lt.s32 q, $1, $2;\n\tselp.u32 $0, 1, 0, q;\n\t}", + "=r,r,r", + False, + ) + ) + @dsl_user_op def gt0(x, *, loc=None, ip=None): """``x > 0`` (fp32, ordered) as 0 / 1.""" - return Int32(_asm(T.i32(), [_v(Float32(x), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.gt.f32 q, $1, 0f00000000;\n\tselp.u32 $0, 1, 0, q;\n\t}", - "=r,f", False)) + return Int32( + _asm( + T.i32(), + [_v(Float32(x), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.gt.f32 q, $1, 0f00000000;\n\tselp.u32 $0, 1, 0, q;\n\t}", + "=r,f", + False, + ) + ) + @dsl_user_op def keep_sel(s, lo_e, hi_e, e, *, loc=None, ip=None): """The split's window mask on one score: ``lo_e <= e <= hi_e ? s : -inf``.""" - return Float32(_asm(T.f32(), [_v(Float32(s), loc, ip), _v(Int32(lo_e), loc, ip), _v(Int32(hi_e), loc, ip)], - "{\n\t.reg .pred a, b;\n\tsetp.le.s32 a, $2, " + str(e) + ";\n\tsetp.ge.s32 b, $3, " - + str(e) + ";\n\tand.pred a, a, b;\n\tselp.f32 $0, $1, 0fFF800000, a;\n\t}", - "=f,f,r,r", False)) + return Float32( + _asm( + T.f32(), + [_v(Float32(s), loc, ip), _v(Int32(lo_e), loc, ip), _v(Int32(hi_e), loc, ip)], + "{\n\t.reg .pred a, b;\n\tsetp.le.s32 a, $2, " + + str(e) + + ";\n\tsetp.ge.s32 b, $3, " + + str(e) + + ";\n\tand.pred a, a, b;\n\tselp.f32 $0, $1, 0fFF800000, a;\n\t}", + "=f,f,r,r", + False, + ) + ) + @dsl_user_op def shfl_bfly(x, m, *, loc=None, ip=None): - return Float32(_asm(T.f32(), [_v(Float32(x), loc, ip)], f"shfl.sync.bfly.b32 $0, $1, {m}, 31, -1;", - "=f,f", True)) + return Float32( + _asm( + T.f32(), + [_v(Float32(x), loc, ip)], + f"shfl.sync.bfly.b32 $0, $1, {m}, 31, -1;", + "=f,f", + True, + ) + ) + @dsl_user_op def shfl_idx(x, src, *, loc=None, ip=None): - return Float32(_asm(T.f32(), [_v(Float32(x), loc, ip)], f"shfl.sync.idx.b32 $0, $1, {src}, 31, -1;", - "=f,f", True)) + return Float32( + _asm( + T.f32(), + [_v(Float32(x), loc, ip)], + f"shfl.sync.idx.b32 $0, $1, {src}, 31, -1;", + "=f,f", + True, + ) + ) + # ---- memory @dsl_user_op def ld_g_v4(addr, pred, *, loc=None, ip=None): """16 bytes from global memory as 4 Int32 (zeros where ``pred`` is 0).""" - res = _asm(llvm.StructType.get_literal([T.i32()] * 4), [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $5, 0;\n\tmov.b32 $0, 0;\n\tmov.b32 $1, 0;\n\t" - "mov.b32 $2, 0;\n\tmov.b32 $3, 0;\n\t@q ld.global.v4.b32 {$0, $1, $2, $3}, [$4];\n\t}", - "=r,=r,=r,=r,l,r", True) + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $5, 0;\n\tmov.b32 $0, 0;\n\tmov.b32 $1, 0;\n\t" + "mov.b32 $2, 0;\n\tmov.b32 $3, 0;\n\t@q ld.global.v4.b32 {$0, $1, $2, $3}, [$4];\n\t}", + "=r,=r,=r,=r,l,r", + True, + ) return [Int32(llvm.extractvalue(T.i32(), res, [i])) for i in range(4)] + @dsl_user_op def ld_g_v2(addr, pred, *, loc=None, ip=None): - res = _asm(llvm.StructType.get_literal([T.i32()] * 2), [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\tmov.b32 $0, 0;\n\tmov.b32 $1, 0;\n\t" - "@q ld.global.v2.b32 {$0, $1}, [$2];\n\t}", "=r,=r,l,r", True) + res = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\tmov.b32 $0, 0;\n\tmov.b32 $1, 0;\n\t" + "@q ld.global.v2.b32 {$0, $1}, [$2];\n\t}", + "=r,=r,l,r", + True, + ) return [Int32(llvm.extractvalue(T.i32(), res, [i])) for i in range(2)] + @dsl_user_op def ld_g_b32(addr, pred, *, loc=None, ip=None): - return Int32(_asm(T.i32(), [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $2, 0;\n\tmov.b32 $0, 0;\n\t" - "@q ld.global.b32 $0, [$1];\n\t}", "=r,l,r", True)) + return Int32( + _asm( + T.i32(), + [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $2, 0;\n\tmov.b32 $0, 0;\n\t" + "@q ld.global.b32 $0, [$1];\n\t}", + "=r,l,r", + True, + ) + ) + @dsl_user_op def st_g_v2(addr, a, b, pred, *, loc=None, ip=None): - _asm(None, [_v(Int64(addr), loc, ip), _v(Int32(a), loc, ip), _v(Int32(b), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\t@q st.global.v2.b32 [$0], {$1, $2};\n\t}", - "l,r,r,r", True) + _asm( + None, + [ + _v(Int64(addr), loc, ip), + _v(Int32(a), loc, ip), + _v(Int32(b), loc, ip), + _v(Int32(pred), loc, ip), + ], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $3, 0;\n\t@q st.global.v2.b32 [$0], {$1, $2};\n\t}", + "l,r,r,r", + True, + ) + @dsl_user_op def st_g_b32(addr, a, pred, *, loc=None, ip=None): - _asm(None, [_v(Int64(addr), loc, ip), _v(Int32(a), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $2, 0;\n\t@q st.global.b32 [$0], $1;\n\t}", "l,r,r", True) + _asm( + None, + [_v(Int64(addr), loc, ip), _v(Int32(a), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $2, 0;\n\t@q st.global.b32 [$0], $1;\n\t}", + "l,r,r", + True, + ) + @dsl_user_op def st_s_v4(addr, w, pred, *, loc=None, ip=None): - _asm(None, [_v(Int32(addr), loc, ip)] + [_v(Int32(x), loc, ip) for x in w] + [_v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $5, 0;\n\t@q st.shared.v4.b32 [$0], {$1, $2, $3, $4};\n\t}", - "r,r,r,r,r,r", True) + _asm( + None, + [_v(Int32(addr), loc, ip)] + + [_v(Int32(x), loc, ip) for x in w] + + [_v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $5, 0;\n\t@q st.shared.v4.b32 [$0], {$1, $2, $3, $4};\n\t}", + "r,r,r,r,r,r", + True, + ) + def isel4(a, b, p): """Element-wise ``p ? a : b`` of two lists of Int32.""" @@ -328,8 +501,13 @@ def isel4(a, b, p): @dsl_user_op def st_s_b32_pred(addr, a, pred, *, loc=None, ip=None): - _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(a), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $2, 0;\n\t@q st.shared.b32 [$0], $1;\n\t}", "r,r,r", True) + _asm( + None, + [_v(Int32(addr), loc, ip), _v(Int32(a), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $2, 0;\n\t@q st.shared.b32 [$0], $1;\n\t}", + "r,r,r", + True, + ) @dsl_user_op @@ -341,127 +519,239 @@ def bar_named(bar_id, count, *, loc=None, ip=None): @dsl_user_op def stamp(addr, pred, *, loc=None, ip=None): """``%clock64`` to global memory (the debug build's phase stamps).""" - _asm(None, [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\t.reg .b64 c;\n\tsetp.ne.s32 q, $1, 0;\n\tmov.u64 c, %clock64;\n\t" - "@q st.global.b64 [$0], c;\n\t}", "l,r", True) + _asm( + None, + [_v(Int64(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\t.reg .b64 c;\n\tsetp.ne.s32 q, $1, 0;\n\tmov.u64 c, %clock64;\n\t" + "@q st.global.b64 [$0], c;\n\t}", + "l,r", + True, + ) @dsl_user_op def st_s_b32(addr, a, *, loc=None, ip=None): - _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(a), loc, ip)], "st.shared.b32 [$0], $1;", "r,r", True) + _asm( + None, + [_v(Int32(addr), loc, ip), _v(Int32(a), loc, ip)], + "st.shared.b32 [$0], $1;", + "r,r", + True, + ) + @dsl_user_op def ld_s_b32(addr, *, loc=None, ip=None): return Int32(_asm(T.i32(), [_v(Int32(addr), loc, ip)], "ld.shared.b32 $0, [$1];", "=r,r", True)) + @dsl_user_op def ld_s_v2f(addr, *, loc=None, ip=None): - res = _asm(llvm.StructType.get_literal([T.f32()] * 2), [_v(Int32(addr), loc, ip)], - "ld.shared.v2.f32 {$0, $1}, [$2];", "=f,=f,r", True) + res = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [_v(Int32(addr), loc, ip)], + "ld.shared.v2.f32 {$0, $1}, [$2];", + "=f,=f,r", + True, + ) return [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(2)] + @dsl_user_op def ld_s_v4f(addr, *, loc=None, ip=None): - res = _asm(llvm.StructType.get_literal([T.f32()] * 4), [_v(Int32(addr), loc, ip)], - "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", True) + res = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [_v(Int32(addr), loc, ip)], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,r", + True, + ) return [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(4)] + @dsl_user_op def mapa(addr, rank, *, loc=None, ip=None): - return Int32(_asm(T.i32(), [_v(Int32(addr), loc, ip), _v(Int32(rank), loc, ip)], - "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", False)) + return Int32( + _asm( + T.i32(), + [_v(Int32(addr), loc, ip), _v(Int32(rank), loc, ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + False, + ) + ) + @dsl_user_op def st_async(raddr, vals, rbar, pred, *, loc=None, ip=None): """``st.async`` of 2 or 4 fp32 into a cluster peer's shared memory, crediting its mbarrier.""" n = len(vals) regs = ", ".join(f"${i + 1}" for i in range(n)) - _asm(None, [_v(Int32(raddr), loc, ip)] + [_v(bits(x), loc, ip) for x in vals] - + [_v(Int32(rbar), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $" + str(n + 2) + ", 0;\n\t" - f"@q st.async.shared::cluster.mbarrier::complete_tx::bytes.v{n}.b32 [$0], {{{regs}}}, [${n + 1}];\n\t}}", - ",".join(["r"] * (n + 3)), True) + _asm( + None, + [_v(Int32(raddr), loc, ip)] + + [_v(bits(x), loc, ip) for x in vals] + + [_v(Int32(rbar), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $" + str(n + 2) + ", 0;\n\t" + f"@q st.async.shared::cluster.mbarrier::complete_tx::bytes.v{n}.b32 [$0], {{{regs}}}, [${n + 1}];\n\t}}", + ",".join(["r"] * (n + 3)), + True, + ) + # ---- barriers @dsl_user_op def bar_init(addr, pred, *, loc=None, ip=None): - _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t@q mbarrier.init.shared::cta.b64 [$0], 1;\n\t}", - "r,r", True) + _asm( + None, + [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t@q mbarrier.init.shared::cta.b64 [$0], 1;\n\t}", + "r,r", + True, + ) + @dsl_user_op def bar_expect(addr, tx, pred, *, loc=None, ip=None): - _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t" - f"@q mbarrier.arrive.expect_tx.shared::cta.b64 _, [$0], {tx};\n\t}}", "r,r", True) + _asm( + None, + [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t" + f"@q mbarrier.arrive.expect_tx.shared::cta.b64 _, [$0], {tx};\n\t}}", + "r,r", + True, + ) + @dsl_user_op def bar_arrive(addr, pred, *, loc=None, ip=None): """``mbarrier.arrive`` (release, CTA scope) when ``pred``.""" - _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t@q mbarrier.arrive.shared::cta.b64 _, [$0];\n\t}", - "r,r", True) + _asm( + None, + [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t@q mbarrier.arrive.shared::cta.b64 _, [$0];\n\t}", + "r,r", + True, + ) @dsl_user_op def bar_wait(addr, pred, cluster_scope, *, loc=None, ip=None): """Wait for phase 0 of the mbarrier at ``addr`` when ``pred``; acquire at cluster scope for peer data.""" scope = ".acquire.cluster" if cluster_scope else "" - _asm(None, [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q, done;\n\t.reg .b32 ph;\n\tmov.b32 ph, 0;\n\tsetp.ne.s32 q, $1, 0;\n\t" - "@!q bra.uni WDONE;\n\tWLOOP:\n\t" - "mbarrier.try_wait.parity" + scope + ".shared::cta.b64 done, [$0], ph;\n\t" - "@!done bra.uni WLOOP;\n\tWDONE:\n\t}", "r,r", True) + _asm( + None, + [_v(Int32(addr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q, done;\n\t.reg .b32 ph;\n\tmov.b32 ph, 0;\n\tsetp.ne.s32 q, $1, 0;\n\t" + "@!q bra.uni WDONE;\n\tWLOOP:\n\t" + "mbarrier.try_wait.parity" + scope + ".shared::cta.b64 done, [$0], ph;\n\t" + "@!done bra.uni WLOOP;\n\tWDONE:\n\t}", + "r,r", + True, + ) + @dsl_user_op def raw(text, *, loc=None, ip=None): _asm(None, [], text, "", True) + # ---- tensor memory and the MMA @dsl_user_op def tmem_alloc(slot, pred, *, loc=None, ip=None): - _asm(None, [_v(Int32(slot), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t" - f"@q tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [$0], {TMEM_COLS};\n\t" - "@q tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n\t}", "r,r", True) + _asm( + None, + [_v(Int32(slot), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t" + f"@q tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [$0], {TMEM_COLS};\n\t" + "@q tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n\t}", + "r,r", + True, + ) + @dsl_user_op def tmem_free(taddr, pred, *, loc=None, ip=None): - _asm(None, [_v(Int32(taddr), loc, ip), _v(Int32(pred), loc, ip)], - "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t" - f"@q tcgen05.dealloc.cta_group::1.sync.aligned.b32 $0, {TMEM_COLS};\n\t}}", "r,r", True) + _asm( + None, + [_v(Int32(taddr), loc, ip), _v(Int32(pred), loc, ip)], + "{\n\t.reg .pred q;\n\tsetp.ne.s32 q, $1, 0;\n\t" + f"@q tcgen05.dealloc.cta_group::1.sync.aligned.b32 $0, {TMEM_COLS};\n\t}}", + "r,r", + True, + ) + @dsl_user_op def tmem_ld(taddr, n, *, loc=None, ip=None): """``tcgen05.ld.16x32bx2.x{n}`` (half split 32): thread l gets lane l % 16, columns 32 (l / 16) + [0, n).""" regs = ", ".join(f"${i}" for i in range(n)) - res = _asm(llvm.StructType.get_literal([T.f32()] * n), [_v(Int32(taddr), loc, ip)], - f"tcgen05.ld.sync.aligned.16x32bx2.x{n}.b32 {{{regs}}}, [${n}], 32;\n\t" - "tcgen05.wait::ld.sync.aligned;", ",".join(["=f"] * n) + ",r", True) + res = _asm( + llvm.StructType.get_literal([T.f32()] * n), + [_v(Int32(taddr), loc, ip)], + f"tcgen05.ld.sync.aligned.16x32bx2.x{n}.b32 {{{regs}}}, [${n}], 32;\n\t" + "tcgen05.wait::ld.sync.aligned;", + ",".join(["=f"] * n) + ",r", + True, + ) return [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(n)] + @dsl_user_op def tmem_st_zero32(taddr, *, loc=None, ip=None): zeros = ", ".join(["$1"] * 32) - _asm(None, [_v(Int32(taddr), loc, ip), _v(Int32(0), loc, ip)], - "tcgen05.st.sync.aligned.16x32bx2.x32.b32 [$0], 32, {" + zeros + "};", "r,r", True) + _asm( + None, + [_v(Int32(taddr), loc, ip), _v(Int32(0), loc, ip)], + "tcgen05.st.sync.aligned.16x32bx2.x32.b32 [$0], 32, {" + zeros + "};", + "r,r", + True, + ) + @dsl_user_op def tmem_st16(taddr, words, *, loc=None, ip=None): regs = ", ".join(f"${i + 1}" for i in range(16)) - _asm(None, [_v(Int32(taddr), loc, ip)] + [_v(Int32(w), loc, ip) for w in words], - "tcgen05.st.sync.aligned.16x32bx2.x16.b32 [$0], 16, {" + regs + "};", "r" + ",r" * 16, True) + _asm( + None, + [_v(Int32(taddr), loc, ip)] + [_v(Int32(w), loc, ip) for w in words], + "tcgen05.st.sync.aligned.16x32bx2.x16.b32 [$0], 16, {" + regs + "};", + "r" + ",r" * 16, + True, + ) + @dsl_user_op -def mma4(d_taddr, a_lo, b_lo, bar, pred, a_tmem, accumulate_first, idesc, a_step, b_step, *, loc=None, - ip=None): +def mma4( + d_taddr, + a_lo, + b_lo, + bar, + pred, + a_tmem, + accumulate_first, + idesc, + a_step, + b_step, + *, + loc=None, + ip=None, +): """Four tcgen05 kind::f16 K16 steps as the split issues them, then a commit to ``bar``: D at ``d_taddr``; A in shared memory (descriptor low word ``a_lo`` + ``a_step`` per step) or in tensor memory (column ``a_lo`` + 8 per step); B descriptor low word ``b_lo`` + ``b_step`` per step; the first step accumulates iff ``accumulate_first``, the others always.""" - lines = ["{", ".reg .pred q, acc0, acc1;", ".reg .b64 a, b;", ".reg .b32 al, bl, hi, id;", - "setp.ne.s32 q, $4, 0;", f"mov.pred acc0, {1 if accumulate_first else 0};", "mov.pred acc1, 1;", - f"mov.b32 hi, {DESC_HI};", f"mov.b32 id, {idesc};"] + lines = [ + "{", + ".reg .pred q, acc0, acc1;", + ".reg .b64 a, b;", + ".reg .b32 al, bl, hi, id;", + "setp.ne.s32 q, $4, 0;", + f"mov.pred acc0, {1 if accumulate_first else 0};", + "mov.pred acc1, 1;", + f"mov.b32 hi, {DESC_HI};", + f"mov.b32 id, {idesc};", + ] for kk in range(4): acc = "acc0" if kk == 0 else "acc1" lines.append(f"add.s32 bl, $2, {b_step * kk};") @@ -475,16 +765,38 @@ def mma4(d_taddr, a_lo, b_lo, bar, pred, a_tmem, accumulate_first, idesc, a_step lines.append(f"@q tcgen05.mma.cta_group::1.kind::f16 [$0], a, b, id, {acc};") lines.append("@q tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [$3];") lines.append("}") - _asm(None, [_v(Int32(d_taddr), loc, ip), _v(Int32(a_lo), loc, ip), _v(Int32(b_lo), loc, ip), - _v(Int32(bar), loc, ip), _v(Int32(pred), loc, ip)], "\n\t".join(lines), "r,r,r,r,r", True) + _asm( + None, + [ + _v(Int32(d_taddr), loc, ip), + _v(Int32(a_lo), loc, ip), + _v(Int32(b_lo), loc, ip), + _v(Int32(bar), loc, ip), + _v(Int32(pred), loc, ip), + ], + "\n\t".join(lines), + "r,r,r,r,r", + True, + ) + @dsl_user_op def pair2(text, a0, a1, b0, b1, *, loc=None, ip=None): """A packed ``.f32x2`` op on (a0, a1), (b0, b1) (``text`` uses x, y -> z); two fp32.""" - res = _asm(llvm.StructType.get_literal([T.f32()] * 2), - [_v(Float32(a0), loc, ip), _v(Float32(a1), loc, ip), _v(Float32(b0), loc, ip), _v(Float32(b1), loc, ip)], - "{\n\t.reg .b64 x, y, z;\n\tmov.b64 x, {$2, $3};\n\tmov.b64 y, {$4, $5};\n\t" + text - + "\n\tmov.b64 {$0, $1}, z;\n\t}", "=f,=f,f,f,f,f", False) + res = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [ + _v(Float32(a0), loc, ip), + _v(Float32(a1), loc, ip), + _v(Float32(b0), loc, ip), + _v(Float32(b1), loc, ip), + ], + "{\n\t.reg .b64 x, y, z;\n\tmov.b64 x, {$2, $3};\n\tmov.b64 y, {$4, $5};\n\t" + + text + + "\n\tmov.b64 {$0, $1}, z;\n\t}", + "=f,=f,f,f,f,f", + False, + ) return [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(2)] @@ -501,14 +813,22 @@ def max3(a, b, c): def scale_mask2(a0, a1, qk_scale, nlo, span, e, *, loc=None, ip=None): """Scores ``e``, ``e + 1`` of a row half: ``mul.rn.f32x2`` by the QK scale, then the split's window mask as one unsigned range test (``e - lo <= span``) per score.""" - res = _asm(llvm.StructType.get_literal([T.f32()] * 2), - [_v(Float32(a0), loc, ip), _v(Float32(a1), loc, ip), _v(Float32(qk_scale), loc, ip), - _v(Int32(nlo), loc, ip), _v(Int32(span), loc, ip)], - "{\n\t.reg .b64 x, y, z;\n\t.reg .f32 s0, s1;\n\t.reg .b32 t;\n\t.reg .pred p;\n\t" - "mov.b64 x, {$2, $3};\n\tmov.b64 y, {$4, $4};\n\tmul.rn.f32x2 z, x, y;\n\tmov.b64 {s0, s1}, z;\n\t" - f"add.s32 t, $5, {e};\n\tsetp.le.u32 p, t, $6;\n\tselp.f32 $0, s0, 0fFF800000, p;\n\t" - f"add.s32 t, $5, {e + 1};\n\tsetp.le.u32 p, t, $6;\n\tselp.f32 $1, s1, 0fFF800000, p;\n\t}}", - "=f,=f,f,f,f,r,r", False) + res = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [ + _v(Float32(a0), loc, ip), + _v(Float32(a1), loc, ip), + _v(Float32(qk_scale), loc, ip), + _v(Int32(nlo), loc, ip), + _v(Int32(span), loc, ip), + ], + "{\n\t.reg .b64 x, y, z;\n\t.reg .f32 s0, s1;\n\t.reg .b32 t;\n\t.reg .pred p;\n\t" + "mov.b64 x, {$2, $3};\n\tmov.b64 y, {$4, $4};\n\tmul.rn.f32x2 z, x, y;\n\tmov.b64 {s0, s1}, z;\n\t" + f"add.s32 t, $5, {e};\n\tsetp.le.u32 p, t, $6;\n\tselp.f32 $0, s0, 0fFF800000, p;\n\t" + f"add.s32 t, $5, {e + 1};\n\tsetp.le.u32 p, t, $6;\n\tselp.f32 $1, s1, 0fFF800000, p;\n\t}}", + "=f,=f,f,f,f,r,r", + False, + ) return [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(2)] @@ -521,16 +841,22 @@ def tmem_ld8x(taddrs, *, loc=None, ip=None): regs = ", ".join(f"${8 * b + i}" for i in range(8)) lines.append(f"tcgen05.ld.sync.aligned.16x32bx2.x8.b32 {{{regs}}}, [${8 * n + b}], 32;") lines.append("tcgen05.wait::ld.sync.aligned;") - res = _asm(llvm.StructType.get_literal([T.f32()] * (8 * n)), [_v(Int32(a), loc, ip) for a in taddrs], - "\n\t".join(lines), ",".join(["=f"] * (8 * n)) + "," + ",".join(["r"] * n), True) + res = _asm( + llvm.StructType.get_literal([T.f32()] * (8 * n)), + [_v(Int32(a), loc, ip) for a in taddrs], + "\n\t".join(lines), + ",".join(["=f"] * (8 * n)) + "," + ",".join(["r"] * n), + True, + ) vals = [Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(8 * n)] - return [vals[8 * b:8 * b + 8] for b in range(n)] + return [vals[8 * b : 8 * b + 8] for b in range(n)] def desc_lo(saddr): """Low word of the split's shared-memory descriptor: the start address >> 4 (14 bits).""" return (saddr >> 4) & 0x3FFF + # ------------------------------------------------------------------------------------------ math blocks def unpack8(words): out = [] @@ -539,6 +865,7 @@ def unpack8(words): out.append(hi_bf16(w)) return out + def prep8(words, cosv, sinv, eps, scale, rope): """``_prep_rows_math`` of this lane's 8 elements (``words``: BF16x2 of elements 8 l .. 8 l + 7 of a head; 8 lanes per head): the squares (``mul.f32``) summed in order with ``add.rn.ftz``, the lanes paired over @@ -562,15 +889,18 @@ def prep8(words, cosv, sinv, eps, scale, rope): y = out return y + def to_cache_fp8(vals, scale_b): """``_to_cache_fp8`` of 8 BF16-exact values: ``mul.bf16`` by the BF16 scale, then E4M3 (RN, satfinite). 4 Int32, each two E4M3 in its low 16 bits (element 2 i in the low byte).""" y = [bf16_mul(v, scale_b) for v in vals] return [e4m3x2(y[2 * i + 1], y[2 * i]) for i in range(4)] + def fp8_to_bf16(pairs): return [e4m3x2_to_bf16x2(p) for p in pairs] + def fp8_words_to_bf16(words): """4 Int32 of 4 E4M3 each (16 elements) -> 8 BF16x2 words.""" out = [] @@ -579,16 +909,21 @@ def fp8_words_to_bf16(words): out.append(e4m3x2_to_bf16x2((w >> 16) & 0xFFFF)) return out + def store_tile_half(tile_row, key, kh, words8, valid): """32 BF16 (16 words) of a 64-wide tile row into chunks 4 kh .. 4 kh + 3 of the 128-byte swizzle.""" for c in range(4): chunk = kh * 4 + c - st_s_v4(tile_row + ((chunk ^ (key % 8)) * 16), words8[4 * c:4 * c + 4], valid) + st_s_v4(tile_row + ((chunk ^ (key % 8)) * 16), words8[4 * c : 4 * c + 4], valid) + def rope_row(g_cos, pos, l8): """``cos[pos, 4 l8 .. 4 l8 + 3]`` and ``sin[...]`` of the fp32 ``[positions, 2, 32]`` table.""" row = g_cos + Int64(pos) * 256 + Int64(l8 * 16) - return [as_f32(x) for x in ld_g_v4(row, Int32(1))], [as_f32(x) for x in ld_g_v4(row + 128, Int32(1))] + return [as_f32(x) for x in ld_g_v4(row, Int32(1))], [ + as_f32(x) for x in ld_g_v4(row + 128, Int32(1)) + ] + def softmax(acc, qk_scale, nlo, span): """The split's online-softmax step on one block (one row half per thread, 32 scores): the scale, the window @@ -631,6 +966,7 @@ def row_max(sml, row, half): m = max_tree([Float32(NEG_BIG)] + v) return fmax(m, shfl_bfly(m, 16)) + def class_sums(xs): """The combine's in-thread tree over the splits ``r + 8 k``, k = 0..7 (absent ones +0): ((x0 + x1) + (x2 + x3)) + ((x4 + x5) + (x6 + x7)), on column pairs with ``add.rn.f32x2``.""" @@ -638,9 +974,12 @@ def class_sums(xs): zero = (Float32(0.0), Float32(0.0)) for i in range(0, len(xs[0]), 2): x = [(xs[k][i], xs[k][i + 1]) for k in range(len(xs))] + [zero] * (8 - len(xs)) - out += add2(add2(add2(x[0], x[1]), add2(x[2], x[3])), add2(add2(x[4], x[5]), add2(x[6], x[7]))) + out += add2( + add2(add2(x[0], x[1]), add2(x[2], x[3])), add2(add2(x[4], x[5]), add2(x[6], x[7])) + ) return out + def denominators(sml, rows, lane, den_lane): """The combine's denominator of each row in ``rows`` over its 64 split lanes (two 32-lane halves): ``w l`` (``mul.f32``) shuffled 16, contracted into ``fma(w, l, partner)``, then lanes 8, 4, 2, 1 (``add.f32``); per @@ -649,8 +988,14 @@ def denominators(sml, rows, lane, den_lane): mm = [] for row in rows: lo_ = ld_s_v2f(sml + (lane * ROWS + row) * 8)[0] - hi_ = selp(ld_s_v2f(sml + ((isel(lane + 32, Int32(0), lt(lane, Int32(NUM_SPLITS - 32)))) * ROWS + row) * 8)[0], - Float32(NEG_BIG), lt(lane, Int32(NUM_SPLITS - 32))) + hi_ = selp( + ld_s_v2f( + sml + + ((isel(lane + 32, Int32(0), lt(lane, Int32(NUM_SPLITS - 32)))) * ROWS + row) * 8 + )[0], + Float32(NEG_BIG), + lt(lane, Int32(NUM_SPLITS - 32)), + ) mm.append(max3(Float32(NEG_BIG), lo_, hi_)) for m in (16, 8, 4, 2, 1): mm = [fmax(x, shfl_bfly(x, m)) for x in mm] @@ -666,7 +1011,9 @@ def denominators(sml, rows, lane, den_lane): for i, sok, ml in chains: m_s = selp(ml[0], Float32(NEG_BIG), sok) l_s = selp(ml[1], Float32(0.0), sok) - ws.append((selp(selp(ex2_ftz(sub(m_s, mm[i])), Float32(0.0), gt0(l_s)), Float32(0.0), sok), l_s)) + ws.append( + (selp(selp(ex2_ftz(sub(m_s, mm[i])), Float32(0.0), gt0(l_s)), Float32(0.0), sok), l_s) + ) parts = [shfl_bfly(mul(w, lv), 16) for w, lv in ws] a = [fma(w, lv, q) for (w, lv), q in zip(ws, parts)] for m in (8, 4, 2, 1): @@ -689,6 +1036,7 @@ def gate_rows(num, den, sig, gate_eps): rs = rsqrt_ftz(add(gate_eps, div_full(tot, Float32(64.0)))) return [bf16_mul(bf16r(mul(rs, att[i])), sig[i]) for i in range(4)] + class SwaClusterBlockKernel: """The kernel of one build: PDL, the rank's KV heads, the cache page, the QKV layout, RoPE, the debug output (0 none, 1 the partials dump, 2 per-phase clock stamps) and the denominator lane are compile-time constants. @@ -698,8 +1046,16 @@ class SwaClusterBlockKernel: class sums (8 of each half's 32 O columns each); warpgroup 4 runs the step keys' prep and the cache write, then the owner's denominators; warpgroup 0 finishes the owner's rows.""" - def __init__(self, enable_pdl: bool, kv_heads: int, page: int, permuted: bool, rope: bool, debug: int, - den_lane: int): + def __init__( + self, + enable_pdl: bool, + kv_heads: int, + page: int, + permuted: bool, + rope: bool, + debug: int, + den_lane: int, + ): self.enable_pdl = enable_pdl self.hkv = kv_heads self.page = page @@ -712,24 +1068,85 @@ def __init__(self, enable_pdl: bool, kv_heads: int, page: int, permuted: bool, r self.dbg = dbg_layout(kv_heads) @cute.jit - def __call__(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Pointer, - seq_ptr: cute.Pointer, scale_ptr: cute.Pointer, cos_ptr: cute.Pointer, pos_ptr: cute.Pointer, - gate_ptr: cute.Pointer, out_ptr: cute.Pointer, dbg_ptr: cute.Pointer, reqs: Int32, - qkv_stride: Int32, bt_stride_req: Int32, bt_stride_kv: Int32, gate_stride: Int32, - gate_hstride: Int32, out_stride: Int32, prep_eps: Float32, q_scale: Float32, - gate_eps: Float32, sm_scale_log2: Float32, stream): - self.kernel(qkv_ptr, pool_ptr, bt_ptr, seq_ptr, scale_ptr, cos_ptr, pos_ptr, gate_ptr, out_ptr, dbg_ptr, - qkv_stride, bt_stride_req, bt_stride_kv, gate_stride, gate_hstride, out_stride, prep_eps, - q_scale, gate_eps, sm_scale_log2).launch( - grid=(CLUSTER, self.hkv, reqs), block=[THREADS, 1, 1], cluster=(CLUSTER, 1, 1), smem=SMEM_BYTES + 1024, - stream=stream, use_pdl=self.enable_pdl) + def __call__( + self, + qkv_ptr: cute.Pointer, + pool_ptr: cute.Pointer, + bt_ptr: cute.Pointer, + seq_ptr: cute.Pointer, + scale_ptr: cute.Pointer, + cos_ptr: cute.Pointer, + pos_ptr: cute.Pointer, + gate_ptr: cute.Pointer, + out_ptr: cute.Pointer, + dbg_ptr: cute.Pointer, + reqs: Int32, + qkv_stride: Int32, + bt_stride_req: Int32, + bt_stride_kv: Int32, + gate_stride: Int32, + gate_hstride: Int32, + out_stride: Int32, + prep_eps: Float32, + q_scale: Float32, + gate_eps: Float32, + sm_scale_log2: Float32, + stream, + ): + self.kernel( + qkv_ptr, + pool_ptr, + bt_ptr, + seq_ptr, + scale_ptr, + cos_ptr, + pos_ptr, + gate_ptr, + out_ptr, + dbg_ptr, + qkv_stride, + bt_stride_req, + bt_stride_kv, + gate_stride, + gate_hstride, + out_stride, + prep_eps, + q_scale, + gate_eps, + sm_scale_log2, + ).launch( + grid=(CLUSTER, self.hkv, reqs), + block=[THREADS, 1, 1], + cluster=(CLUSTER, 1, 1), + smem=SMEM_BYTES + 1024, + stream=stream, + use_pdl=self.enable_pdl, + ) @cute.kernel - def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Pointer, seq_ptr: cute.Pointer, - scale_ptr: cute.Pointer, cos_ptr: cute.Pointer, pos_ptr: cute.Pointer, gate_ptr: cute.Pointer, - out_ptr: cute.Pointer, dbg_ptr: cute.Pointer, qkv_stride: Int32, bt_stride_req: Int32, - bt_stride_kv: Int32, gate_stride: Int32, gate_hstride: Int32, out_stride: Int32, - prep_eps: Float32, q_scale: Float32, gate_eps: Float32, sm_scale_log2: Float32): + def kernel( + self, + qkv_ptr: cute.Pointer, + pool_ptr: cute.Pointer, + bt_ptr: cute.Pointer, + seq_ptr: cute.Pointer, + scale_ptr: cute.Pointer, + cos_ptr: cute.Pointer, + pos_ptr: cute.Pointer, + gate_ptr: cute.Pointer, + out_ptr: cute.Pointer, + dbg_ptr: cute.Pointer, + qkv_stride: Int32, + bt_stride_req: Int32, + bt_stride_kv: Int32, + gate_stride: Int32, + gate_hstride: Int32, + out_stride: Int32, + prep_eps: Float32, + q_scale: Float32, + gate_eps: Float32, + sm_scale_log2: Float32, + ): tid, _, _ = cute.arch.thread_idx() rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) _, h, b = cute.arch.block_idx() @@ -756,7 +1173,11 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi g_pos = pos_ptr.toint() g_out = out_ptr.toint() g_dbg = dbg_ptr.toint() - g_ts = g_dbg + Int64(self.dbg["ts"]) + Int64(((b * self.hkv + h) * CLUSTER + rank) * STAMPS * 8) + g_ts = ( + g_dbg + + Int64(self.dbg["ts"]) + + Int64(((b * self.hkv + h) * CLUSTER + rank) * STAMPS * 8) + ) sml = sbase + ML_OFF s_bar = sbase + BAR_OFF o_bar = s_bar + 40 @@ -831,8 +1252,14 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi # ------------------------------------------------ after the wait: the projection, one round trip qr = tid // 8 # query row of warpgroups 0-3 (token qr / 8, head qr % 8) - qa = g_qkv + (Int64(b * TOKENS + qr // 8) * Int64(qkv_stride) - + Int64((h * GROUP + qr % 8) * self.q_stride + l8 * 8)) * 2 + qa = ( + g_qkv + + ( + Int64(b * TOKENS + qr // 8) * Int64(qkv_stride) + + Int64((h * GROUP + qr % 8) * self.q_stride + l8 * 8) + ) + * 2 + ) row_k = Int64(b * TOKENS + tw) * Int64(qkv_stride) step = 1 - prep_q qw = ld_g_v4(qa, prep_q) @@ -840,8 +1267,14 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi vraw = ld_g_v4(g_qkv + (row_k + Int64(self.v_col0 + h * D + l8 * 8)) * 2, step) og = wt // 16 od = wt % 16 - ga = gate_ptr.toint() + (Int64(b * TOKENS + rank) * Int64(gate_stride) - + Int64((h * GROUP + og) * gate_hstride + od * 4)) * 2 + ga = ( + gate_ptr.toint() + + ( + Int64(b * TOKENS + rank) * Int64(gate_stride) + + Int64((h * GROUP + og) * gate_hstride + od * 4) + ) + * 2 + ) gw = ld_g_v2(ga, lt(tid, Int32(128))) # Tensor memory only now, with the projection loads in flight: after the wait every predecessor has @@ -856,23 +1289,44 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi tmem = ld_s_b32(sbase + SLOT_OFF) # warpgroups 0-3: the query rows; warpgroup 4: the step keys (one prep each: 8 lanes per head) - y = prep8(isel4(qw, kraw, prep_q), cosv, sinv, prep_eps, selp(q_scale, Float32(1.0), prep_q), self.rope) + y = prep8( + isel4(qw, kraw, prep_q), + cosv, + sinv, + prep_eps, + selp(q_scale, Float32(1.0), prep_q), + self.rope, + ) qb = fp8_to_bf16(to_cache_fp8(y, scale_b)) st_s_v4(sbase + Q_OFF + qr * 128 + ((l8 ^ (qr % 8)) * 16), qb, prep_q) if cutlass.const_expr(self.debug == 1): - dq = g_dbg + Int64(self.dbg["q"]) + Int64(((b * self.hkv + h) * ROWS + qr) * D + l8 * 8) * 2 + dq = ( + g_dbg + + Int64(self.dbg["q"]) + + Int64(((b * self.hkv + h) * ROWS + qr) * D + l8 * 8) * 2 + ) st_g_v2(dq, qb[0], qb[1], prep_q) st_g_v2(dq + 8, qb[2], qb[3], prep_q) k8 = to_cache_fp8(y, scale_b) v8 = to_cache_fp8(unpack8(vraw), scale_b) inw = Int64((jw % self.page) * D + l8 * 8) - st_g_v2(g_pool + (Int64(wpk) * self.hkv + Int64(h)) * (self.page * D) + inw, k8[0] | (k8[1] << 16), - k8[2] | (k8[3] << 16), writer) - st_g_v2(g_pool + (Int64(wpv) * self.hkv + Int64(h)) * (self.page * D) + inw, v8[0] | (v8[1] << 16), - v8[2] | (v8[3] << 16), writer) + st_g_v2( + g_pool + (Int64(wpk) * self.hkv + Int64(h)) * (self.page * D) + inw, + k8[0] | (k8[1] << 16), + k8[2] | (k8[3] << 16), + writer, + ) + st_g_v2( + g_pool + (Int64(wpv) * self.hkv + Int64(h)) * (self.page * D) + inw, + v8[0] | (v8[1] << 16), + v8[2] | (v8[3] << 16), + writer, + ) sj = (jw - base) // BLOCK kj = (jw - base) % BLOCK - mine = step * lt(wt, Int32(64)) * (1 - lt(sj % CLUSTER, rank)) * (1 - lt(rank, sj % CLUSTER)) + mine = ( + step * lt(wt, Int32(64)) * (1 - lt(sj % CLUSTER, rank)) * (1 - lt(rank, sj % CLUSTER)) + ) prow = (sj // CLUSTER) * TILE_BYTES + kj * 128 + ((l8 ^ (kj % 8)) * 16) st_s_v4(sbase + K_OFF + prow, qb, mine) st_s_v4(sbase + V_OFF + prow, fp8_to_bf16(v8), mine) @@ -882,8 +1336,18 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi stamp(g_ts + 24, is_t0) # ------------------------------------------------ QK: each warpgroup's lead thread issues its block's - mma4(tmem + (16 << 16) + COLS_PER_BLOCK * wg, desc_lo(sbase + Q_OFF), desc_lo(sbase + K_OFF + wg * TILE_BYTES), - s_bar + 8 * wg, lead * live, False, False, IDESC_QK, 2, 2) + mma4( + tmem + (16 << 16) + COLS_PER_BLOCK * wg, + desc_lo(sbase + Q_OFF), + desc_lo(sbase + K_OFF + wg * TILE_BYTES), + s_bar + 8 * wg, + lead * live, + False, + False, + IDESC_QK, + 2, + 2, + ) # ------------------------------------------------ each warpgroup: its block's softmax, PV, (max, sum) out row = 16 * warp + l16 @@ -897,8 +1361,12 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi lo_e = cutlass.max(lo, q_pos - (WINDOW - 1)) - jb hi_e = q_pos - jb empty = lt(hi_e, lo_e) - m_new, l_new, pw = softmax(tmem_ld(col + (16 << 16), 32), qk_scale, isel(Int32(-64), 0 - lo_e, empty), - isel(Int32(0), hi_e - lo_e, empty)) + m_new, l_new, pw = softmax( + tmem_ld(col + (16 << 16), 32), + qk_scale, + isel(Int32(-64), 0 - lo_e, empty), + isel(Int32(0), hi_e - lo_e, empty), + ) m_k = selp(m_new, Float32(NEG_BIG), live) l_k = selp(l_new, Float32(0.0), live) tmem_st16(col + 64, pw) @@ -907,15 +1375,29 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi raw("tcgen05.fence::before_thread_sync;") bar_named(1 + wg, 128) raw("tcgen05.fence::after_thread_sync;") - mma4(tmem + COLS_PER_BLOCK * wg, tmem + COLS_PER_BLOCK * wg + 64, - desc_lo(sbase + V_OFF + wg * TILE_BYTES), o_bar + 8 * wg, lead * live, True, True, IDESC_PV, 0, 128) + mma4( + tmem + COLS_PER_BLOCK * wg, + tmem + COLS_PER_BLOCK * wg + 64, + desc_lo(sbase + V_OFF + wg * TILE_BYTES), + o_bar + 8 * wg, + lead * live, + True, + True, + IDESC_PV, + 0, + 128, + ) # a split with no keys: O stays the zeros just stored (the split's acc * dequant of an empty loop) bar_arrive(o_bar + 8 * wg, lead * (1 - live)) slot = sml + (s * ROWS + row) * 8 for dd in cutlass.range_constexpr(4): st_async(mapa(slot, half * 4 + dd), [m_k, l_k], mapa(ml_bar, half * 4 + dd), valid) if cutlass.const_expr(self.debug == 1): - dml = g_dbg + Int64(self.dbg["ml"]) + Int64((((b * NUM_SPLITS + s) * self.hkv + h) * ROWS + row) * 8) + dml = ( + g_dbg + + Int64(self.dbg["ml"]) + + Int64((((b * NUM_SPLITS + s) * self.hkv + h) * ROWS + row) * 8) + ) st_g_v2(dml, bits(m_k), bits(l_k), valid * (1 - half)) if cutlass.const_expr(self.debug == 2): stamp(g_ts + 32, is_t0) @@ -941,7 +1423,9 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi mls = [] for k in cutlass.range_constexpr(NB_MAX): sk = rank + CLUSTER * k - mls.append(ld_s_v2f(sml + (isel(sk, Int32(0), lt(sk, Int32(NUM_SPLITS))) * ROWS + row) * 8)) + mls.append( + ld_s_v2f(sml + (isel(sk, Int32(0), lt(sk, Int32(NUM_SPLITS))) * ROWS + row) * 8) + ) m_max = row_max(sml, row, half) wks = [] for k in cutlass.range_constexpr(NB_MAX): @@ -951,7 +1435,9 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi raw("tcgen05.fence::after_thread_sync;") if cutlass.const_expr(self.debug == 2): stamp(g_ts + 72, is_t0) - os_ = tmem_ld8x([tmem + lane_tm + COLS_PER_BLOCK * k + 8 * chunk for k in range(NB_MAX)]) + os_ = tmem_ld8x( + [tmem + lane_tm + COLS_PER_BLOCK * k + 8 * chunk for k in range(NB_MAX)] + ) for k in cutlass.range_constexpr(NB_MAX): sk = rank + CLUSTER * k lk = lt(sk, nblk) @@ -964,9 +1450,24 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi for i in cutlass.range_constexpr(8): ok = mul(o[i], dequant) if cutlass.const_expr(self.debug == 1): - dob = g_dbg + Int64(self.dbg["o"]) + Int64( - ((((b * NUM_SPLITS + sk) * self.hkv + h) * ROWS + row) * D + 32 * half + 8 * chunk + i) * 4) - st_g_b32(dob, bits(selp(ok, Float32(0.0), lk)), lt(sk, Int32(NUM_SPLITS)) * cls_ok) + dob = ( + g_dbg + + Int64(self.dbg["o"]) + + Int64( + ( + (((b * NUM_SPLITS + sk) * self.hkv + h) * ROWS + row) * D + + 32 * half + + 8 * chunk + + i + ) + * 4 + ) + ) + st_g_b32( + dob, + bits(selp(ok, Float32(0.0), lk)), + lt(sk, Int32(NUM_SPLITS)) * cls_ok, + ) xs.append(xk) tsum = class_sums(xs) dst = sbase + T_OFF + ((rank * GROUP + row % 8) * D + 32 * half + 8 * chunk) * 4 @@ -977,7 +1478,9 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi stamp(g_ts + 88, is_t0) # warpgroup 4: the owner's denominators (warp i: heads 2 i, 2 i + 1), the combine's butterfly if wg == 4: - dens = denominators(sml, [rank * GROUP + 2 * warp, rank * GROUP + 2 * warp + 1], lane, self.den_lane) + dens = denominators( + sml, [rank * GROUP + 2 * warp, rank * GROUP + 2 * warp + 1], lane, self.den_lane + ) st_s_b32_pred(sbase + DEN_OFF + (warp * 2) * 4, bits(dens[0]), step) st_s_b32_pred(sbase + DEN_OFF + (warp * 2 + 1) * 4, bits(dens[1]), step) if cutlass.const_expr(self.debug == 2): @@ -1005,7 +1508,11 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi pw4.append(add(tv[2 * w2][i], tv[2 * w2 + 1][i])) num.append(add(add(pw4[0], pw4[2]), add(pw4[1], pw4[3]))) res = gate_rows(num, den, sig, gate_eps) - oa = g_out + (Int64(b * TOKENS + rank) * Int64(out_stride) + Int64((h * GROUP + og) * D + od * 4)) * 2 + oa = ( + g_out + + (Int64(b * TOKENS + rank) * Int64(out_stride) + Int64((h * GROUP + og) * D + od * 4)) + * 2 + ) st_g_v2(oa, pack_bf16x2(res[1], res[0]), pack_bf16x2(res[3], res[2]), fin) if cutlass.const_expr(self.debug == 2): stamp(g_ts + 64, is_t0) @@ -1014,10 +1521,26 @@ def kernel(self, qkv_ptr: cute.Pointer, pool_ptr: cute.Pointer, bt_ptr: cute.Poi _compiled: dict = {} -def compiled_key(device: int, enable_pdl: bool, kv_heads: int, page: int, permuted: bool, rope: bool, - debug: int = 0, den_lane: int = 0) -> tuple: - return (device, bool(enable_pdl), int(kv_heads), int(page), bool(permuted), bool(rope), int(debug), - int(den_lane)) +def compiled_key( + device: int, + enable_pdl: bool, + kv_heads: int, + page: int, + permuted: bool, + rope: bool, + debug: int = 0, + den_lane: int = 0, +) -> tuple: + return ( + device, + bool(enable_pdl), + int(kv_heads), + int(page), + bool(permuted), + bool(rope), + int(debug), + int(den_lane), + ) def _compile(key: tuple): @@ -1029,43 +1552,114 @@ def _compile(key: tuple): device, pdl, hkv, page, permuted, rope, debug, den_lane = key kernel = SwaClusterBlockKernel(pdl, hkv, page, permuted, rope, debug, den_lane) - ptrs = [make_ptr(dt, 16, cute.AddressSpace.gmem, assumed_align=16) for dt in ( - cutlass.BFloat16, cutlass.Uint8, cutlass.Int32, cutlass.Int32, cutlass.Float32, cutlass.Float32, - cutlass.Int32, cutlass.BFloat16, cutlass.BFloat16, cutlass.Uint8)] + ptrs = [ + make_ptr(dt, 16, cute.AddressSpace.gmem, assumed_align=16) + for dt in ( + cutlass.BFloat16, + cutlass.Uint8, + cutlass.Int32, + cutlass.Int32, + cutlass.Float32, + cutlass.Float32, + cutlass.Int32, + cutlass.BFloat16, + cutlass.BFloat16, + cutlass.Uint8, + ) + ] with torch.cuda.device(device): stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - fn = cute.compile(kernel, *ptrs, *[cutlass.Int32(1)] * 7, *[cutlass.Float32(0.0)] * 4, stream, - options="--enable-tvm-ffi") + fn = cute.compile( + kernel, + *ptrs, + *[cutlass.Int32(1)] * 7, + *[cutlass.Float32(0.0)] * 4, + stream, + options="--enable-tvm-ffi", + ) _compiled[key] = fn return fn -def swa_cluster_block(qkv: torch.Tensor, kv_pool: torch.Tensor, block_tables: torch.Tensor, seq_lens: torch.Tensor, - kv_scale: torch.Tensor, num_kv_heads: int, cos_sin: Optional[torch.Tensor], - positions: Optional[torch.Tensor], prep_eps: float, q_scale: float, qkv_permuted: bool, - gate: torch.Tensor, gate_head_stride: int, gate_eps: float, enable_pdl: bool, - debug: Optional[torch.Tensor] = None, den_lane: int = 0, stamps: bool = False) -> torch.Tensor: +def swa_cluster_block( + qkv: torch.Tensor, + kv_pool: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + kv_scale: torch.Tensor, + num_kv_heads: int, + cos_sin: Optional[torch.Tensor], + positions: Optional[torch.Tensor], + prep_eps: float, + q_scale: float, + qkv_permuted: bool, + gate: torch.Tensor, + gate_head_stride: int, + gate_eps: float, + enable_pdl: bool, + debug: Optional[torch.Tensor] = None, + den_lane: int = 0, + stamps: bool = False, +) -> torch.Tensor: """The block of ``seq_lens.numel()`` sequences of 8 rows each; the caller checked ``fits`` and the multi-request op's preconditions. ``debug`` (uint8, ``dbg_layout(G)["bytes"]``) receives the Q tiles and the split partials in the split kernel's layout; ``den_lane`` exists for the build's tests.""" reqs = seq_lens.numel() - out = torch.empty((qkv.shape[0], GROUP * num_kv_heads * D), dtype=torch.bfloat16, device=qkv.device) + out = torch.empty( + (qkv.shape[0], GROUP * num_kv_heads * D), dtype=torch.bfloat16, device=qkv.device + ) rope = cos_sin is not None - key = compiled_key(qkv.device.index, enable_pdl, num_kv_heads, kv_pool.shape[2], qkv_permuted, rope, - 0 if debug is None else (2 if stamps else 1), den_lane) + key = compiled_key( + qkv.device.index, + enable_pdl, + num_kv_heads, + kv_pool.shape[2], + qkv_permuted, + rope, + 0 if debug is None else (2 if stamps else 1), + den_lane, + ) if torch.cuda.is_current_stream_capturing() and key not in _compiled: - raise RuntimeError("swa_cluster_block compiles on its first call per device and build; make that call " - "before capturing a CUDA graph") + raise RuntimeError( + "swa_cluster_block compiles on its first call per device and build; make that call " + "before capturing a CUDA graph" + ) fn = _compile(key) import cuda.bindings.driver as cuda - fn(qkv.data_ptr(), kv_pool.data_ptr(), block_tables.data_ptr(), seq_lens.data_ptr(), kv_scale.data_ptr(), - (cos_sin if rope else kv_scale).data_ptr(), (positions if rope else seq_lens).data_ptr(), gate.data_ptr(), - out.data_ptr(), (debug if debug is not None else out).data_ptr(), reqs, qkv.stride(0), block_tables.stride(0), - block_tables.stride(1), gate.stride(0), gate_head_stride, out.stride(0), float(prep_eps), float(q_scale), - float(gate_eps), math.log2(math.e) / math.sqrt(D), - cuda.CUstream(torch.cuda.current_stream(qkv.device).cuda_stream)) + fn( + qkv.data_ptr(), + kv_pool.data_ptr(), + block_tables.data_ptr(), + seq_lens.data_ptr(), + kv_scale.data_ptr(), + (cos_sin if rope else kv_scale).data_ptr(), + (positions if rope else seq_lens).data_ptr(), + gate.data_ptr(), + out.data_ptr(), + (debug if debug is not None else out).data_ptr(), + reqs, + qkv.stride(0), + block_tables.stride(0), + block_tables.stride(1), + gate.stride(0), + gate_head_stride, + out.stride(0), + float(prep_eps), + float(q_scale), + float(gate_eps), + math.log2(math.e) / math.sqrt(D), + cuda.CUstream(torch.cuda.current_stream(qkv.device).cuda_stream), + ) return out -__all__ = ["CLUSTER", "NUM_SPLITS", "dbg_layout", "fits", "owned_splits", "swa_cluster_block", "compiled_key"] +__all__ = [ + "CLUSTER", + "NUM_SPLITS", + "dbg_layout", + "fits", + "owned_splits", + "swa_cluster_block", + "compiled_key", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_step_prep.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_step_prep.py index dffdbac5eeac..c2187c07062f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_step_prep.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_step_prep.py @@ -144,9 +144,7 @@ def _columns_kernel( @triton.jit -def _layer_offset( - layer, o0, o1, o2, o3, o4, o5, o6, o7 -): +def _layer_offset(layer, o0, o1, o2, o3, o4, o5, o6, o7): offset = o0 offset = tl.where(layer == 1, o1, offset) offset = tl.where(layer == 2, o2, offset) @@ -330,7 +328,9 @@ def refresh_block_tables(self, worker, attn_metadata, num_seqs: int) -> bool: tables = worker._ctx_block_tables offsets_name = ( - "kv_cache_block_offsets" if worker._ctx_unified_cache else "draft_kv_cache_block_offsets" + "kv_cache_block_offsets" + if worker._ctx_unified_cache + else "draft_kv_cache_block_offsets" ) src = getattr(attn_metadata, offsets_name, None) if src is None or num_seqs <= 0 or src.dim() != 4: @@ -360,7 +360,9 @@ def refresh_block_tables(self, worker, attn_metadata, num_seqs: int) -> bool: ) return True - def store_step_context(self, worker, draft_model, gen_hs, slots, rows, per_request: int) -> None: + def store_step_context( + self, worker, draft_model, gen_hs, slots, rows, per_request: int + ) -> None: """``DFlashWorker._store_step_context`` without the mask (the fork's side chain) over the manager's pages: the positions and columns in one launch (from one read of the lengths, as the built-in's ``ctx_len``), the built-in projection and ``precompute_context_kv``, From 5a73305afb9b9d48a24cd9046e9742ed789ff6ec Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Mon, 9 Nov 2026 21:40:46 -0800 Subject: [PATCH 4/9] [None][feat] MuseSpark GB300: drafter tail folds as two default-off levers Folds two drafter-tail launch pairs of the sm103 TEP8 decode step, each behind a dispatch lever that is default off: * drafter_conv_epilogue (on drafter_block_conv and drafter_fixed_gemms): on blocks of 9..32 rows the coefficient GEMM computes the input-side convolution in its epilogue and prefetches its operands under the mainloop, over a row-permuted copy of the projection weight (32 MiB per convolution). The weight values are unchanged. * drafter_finish_norm (on drafter_block_decode): on blocks the drafter hops do not take, each hop's finish convolution, residual add and RMSNorm are one launch that releases its dependents at entry. Both are lossless-bitwise and knob-only, so an environment that sets nothing runs exactly the paths it ran before. The two environment switches the work carried are retired by name in levers.RETIRED, so a deployment still exporting one is told which knob replaces it instead of being silently ignored. Served at TEP8 only: these are base fast paths, and the sm103 TEP4 composition lists its own set, so the resolver refuses the knob there rather than serving an unmeasured shape. Proved rather than asserted: a verifier rebuilds FAST_PATHS, DEPENDS_ON and both compositions by AST and resolves them with an EMPTY environment. The resolved path sets are identical to the base at all twelve deployments checked, and each new lever then turns on by its own knob. Run on GB300 (2 nodes, TEP8 holder): all 16 touched modules import, including this commit's two new CuTe DSL files; the real lever table resolves to 44 TEP8 paths with an EMPTY environment, matching the AST verifier; conv_coef_gemm compiles and runs with finite outputs at 16 rows. The model itself was not booted. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 34d9a4d0b161a427d0e893b73906fbcb695c550d) --- .../dispatch.py | 13 + .../drafter.py | 8 +- .../drafter_block_decode/block_decode.py | 294 +++++++- .../levers.py | 5 + .../ops/gemm/conv_coef_gemm.py | 647 ++++++++++++++++++ .../ops/spec/drafter_finish_norm.py | 440 ++++++++++++ 6 files changed, 1399 insertions(+), 8 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/conv_coef_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_finish_norm.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py index 716fd90febef..237e6c39c4a9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py @@ -411,6 +411,19 @@ class FastPath: ), 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("context_kv_fork", "drafter context-K/V store on a side stream", "draft"), FastPath( "drafter_fused_prep", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py index 2d7201adb531..2ab5cfee705d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py @@ -22,7 +22,13 @@ 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"), + ) 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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py index f90a92ed6087..cbb16497aa0f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py @@ -41,6 +41,18 @@ attention and append on a multi-request split geometry) instead of ``dflash_window_attention``. +Two levers refine them, each resolved by the plan and passed to ``install``: + +* ``drafter_conv_epilogue`` (on ``drafter_block_conv`` and ``drafter_fixed_gemms``): on blocks of + 9..32 rows whose coefficient projection runs a ``drafter_fixed_gemms`` tactic, ``prepare`` is one + launch, the coefficient GEMM with the input-side convolution in its epilogue + (``ops/gemm/conv_coef_gemm.py``, ``ConvEpilogue``) over a copy of the projection weight with + side 0's two taps per group on adjacent rows (32 MiB per convolution), bitwise the fixed-tactic + GEMM and ``dflash2_block_conv``. +* ``drafter_finish_norm`` (on ``drafter_block_decode``): on blocks the drafter + hops do not take, each hop's finish convolution and the next norm's fused residual add + RMSNorm + are one launch (``ops/spec/drafter_finish_norm.py``, ``FinishNorm``), bitwise the two. + The fast paths compose. ``windowed_block_decode``, like the built-in forward, reaches the block convolutions through the modules' ``prepare`` / ``finish``, which call ``_convolve``; with ``drafter_block_conv`` installed that is the fused convolution @@ -50,12 +62,13 @@ from __future__ import annotations import types -from typing import Callable +from typing import Callable, Optional import torch from torch import nn from tensorrt_llm._torch.models.modeling_dflash import DFlash2BlockConv, DFlashForCausalLM +from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm.logger import logger from ...dispatch import LAUNCH_PDL @@ -65,6 +78,8 @@ geometry_unsupported_reason, unsupported_reason, ) +from ...ops.spec.drafter_finish_norm import drafter_finish_norm +from ...ops.spec.drafter_finish_norm import takes as _finish_norm_takes #: Operand dtypes the fused convolution takes; anything else runs the built-in. SUPPORTED_DTYPES = (torch.bfloat16, torch.float16, torch.float32) @@ -77,6 +92,165 @@ _ENTRY_ROWS = (8,) +#: Block rows ``ConvEpilogue`` serves: two to four requests' blocks (the one-request block keeps +#: the PDL chain's convolution trigger). +CONV_EPILOGUE_ROWS = (9, 32) + + +class ConvEpilogue: + """One convolution's ``prepare`` as the coefficient GEMM with the convolution in its epilogue. + + Holds the ``interleave``d copy of the module's ``kernel_projection`` weight. ``prepare`` + returns None for a call it does not serve: rows outside ``CONV_EPILOGUE_ROWS``, operands the + kernel does not take, or a projection not on a ``drafter_fixed_gemms`` tactic at that row + count (whose bits the kernel reproduces); the module then runs the built-in ``prepare``. + """ + + logged = False + + def __init__(self, weight: torch.Tensor) -> None: + from ...ops.gemm import conv_coef_gemm as _gemm + + self.weight = _gemm.interleave(weight.detach()) + self.source = weight + + def prepare(self, conv: "FusedDFlash2BlockConv", x: torch.Tensor, block_size: int): + rows = x.shape[0] if x.dim() == 2 else 0 + if not CONV_EPILOGUE_ROWS[0] <= rows <= CONV_EPILOGUE_ROWS[1]: + return None + proj = conv.kernel_projection + tactics = getattr(proj, "_fixed_gemm_tactics", None) + if ( + tactics is None + or proj.weight is not self.source + or x.dtype != torch.bfloat16 + or not x.is_contiguous() + or x.data_ptr() % 32 + or torch.is_grad_enabled() + ): + return None + from ...ops.gemm import batch_splitk_gemm as _bsk + from ...ops.gemm import conv_coef_gemm as _gemm + from ...ops.gemm.fixed_low_m_gemm import tactic_for + + tactic = tactic_for(rows, x.shape[1], self.weight.shape[0], tactics) + if tactic is None: + return None + conv_out, coefficients = _gemm.run( + x, + self.weight, + conv.base_kernel[0], + block_size, + _bsk.BatchSplitKTactic(*tactic), + bool(getattr(proj, "_fixed_gemm_pdl", True)), + ) + if not ConvEpilogue.logged: + ConvEpilogue.logged = True + logger.info( + f"[musespark_open] drafter_conv_epilogue: the first fused prepare ran ({rows} rows, " + f"split-K tactic {tuple(tactic)}): coefficient GEMM + input-side convolution in one launch" + ) + return conv_out, coefficients.view(rows, 2, conv.taps, conv.num_groups)[:, 1] + + +def _conv_epilogue_reason(conv: nn.Module) -> Optional[str]: + weight = conv.kernel_projection.weight + if ( + conv.taps != 2 + or conv.group_size != 16 + or tuple(weight.shape) != (2048, 8192) + or weight.dtype != torch.bfloat16 + or conv.base_kernel.dtype != torch.bfloat16 + or conv.kernel_projection.bias is not None + ): + return "not a two-tap, group-16, BF16 [2048, 8192] coefficient projection" + return None + + +class FinishNorm: + """A hop's finish convolution and the next norm's fused add + RMSNorm in one launch. + + ``__call__`` returns the normed rows and updates ``residual`` in place, as the built-in + ``finish`` then ``norm(hidden, residual)`` do; operands the fused op does not take run those + two. + """ + + logged = False + + def __init__(self, enable_pdl: bool, entry_trigger: bool = True) -> None: + self.enable_pdl = enable_pdl + #: Release the successor at entry (its next coefficient GEMM streams its weight under this + #: kernel), as the fused drafter hop releases the same successors early; False: after the + #: stores. Every successor reads this kernel's output and residual after its own wait. + self.entry_trigger = entry_trigger + + def __call__( + self, + hidden: torch.Tensor, + coefficients: torch.Tensor, + conv: "FusedDFlash2BlockConv", + norm: nn.Module, + residual: torch.Tensor, + block_size: int, + ) -> torch.Tensor: + base = conv.base_kernel[1] + if not _finish_norm_takes(hidden, coefficients, base, residual, norm.weight): + hidden = conv.finish(hidden, coefficients, block_size) + normed, _ = norm(hidden, residual) + return normed + normed = drafter_finish_norm( + hidden, + coefficients, + base, + residual, + norm.weight, + norm.variance_epsilon, + block_size, + self.enable_pdl, + self.entry_trigger, + ) + if not FinishNorm.logged: + FinishNorm.logged = True + logger.info( + f"[musespark_open] drafter_finish_norm: the first fused hop tail ran ({hidden.shape[0]} " + "rows): finish convolution + residual add + RMSNorm in one launch" + ) + return normed + + +def _finish_norm_reason(draft_model: nn.Module) -> Optional[str]: + """Why ``FinishNorm`` cannot stand in for this drafter's finish + norm pairs, or None.""" + if not getattr(draft_model, "has_block_conv", False): + return "the drafter has no DFlash 2 block convolutions" + layers = draft_model.model.layers + norms = [("model.norm", draft_model.model.norm)] + for index, layer in enumerate(layers): + norms.append((f"layers.{index}.input_layernorm", layer.input_layernorm)) + norms.append((f"layers.{index}.post_attention_layernorm", layer.post_attention_layernorm)) + for where, norm in norms: + if type(norm) is not RMSNorm: + return f"{where} is {type(norm).__name__}, not the built-in RMSNorm" + if norm.use_gemma or norm.use_cuda_tile or norm.return_hp_output: + return f"{where} is not the plain flashinfer RMSNorm" + if norm.is_nvfp4 and norm.nvfp4_scale is not None: + return f"{where} quantizes its output" + weight = getattr(norm, "weight", None) + if weight is None or weight.dtype != torch.bfloat16 or tuple(weight.shape) != (8192,): + return f"{where} has no BF16 weight of 8192" + for kind in ("attention_convs", "mlp_convs"): + for index, conv in enumerate(getattr(draft_model, kind, None) or ()): + if ( + conv.taps != 2 + or conv.group_size != 16 + or tuple(conv.base_kernel.shape) != (2, 2, 8192) + or conv.base_kernel.dtype != torch.bfloat16 + ): + return f"{kind}.{index} is not a two-tap, group-16, BF16 convolution over 8192" + if type(conv).finish is not DFlash2BlockConv.finish: + return f"{kind}.{index} overrides finish" + return None + + class FusedDFlash2BlockConv(DFlash2BlockConv): """The built-in DFlash 2 block convolution with ``_convolve`` as one kernel. @@ -88,6 +262,17 @@ class FusedDFlash2BlockConv(DFlash2BlockConv): built-in convolution. """ + def prepare( + self, hidden_states: torch.Tensor, block_size: int + ) -> tuple[torch.Tensor, torch.Tensor]: + """The built-in ``prepare``, or ``ConvEpilogue``'s one launch where it serves the call.""" + fused = self.__dict__.get("_conv_epilogue") + if fused is not None: + prepared = fused.prepare(self, hidden_states, block_size) + if prepared is not None: + return prepared + return super().prepare(hidden_states, block_size) + def _convolve( self, hidden_states: torch.Tensor, delta: torch.Tensor, side: int, block_size: int ) -> torch.Tensor: @@ -333,6 +518,13 @@ def windowed_block_decode( merge_entry = bool(self.__dict__.get("_pdl_merge")) and B * block_size in self.__dict__.get( "_pdl_merge_rows", _ENTRY_ROWS ) + # ``drafter_finish_norm``: each hop's finish convolution + next norm in one launch on + # blocks the drafter hops do not take. + finish_norm = self.__dict__.get("_finish_norm") + if finish_norm is not None and ( + hops is not None or not has_block_conv or noise_embedding.dtype != torch.bfloat16 + ): + finish_norm = None # wide_drafter_window_attention: a block decode of several requests on the wide op. wide = self.__dict__.get("_wide_window_attention") if wide is not None and not wide.takes(B * block_size): @@ -414,6 +606,17 @@ def windowed_block_decode( layer.post_attention_layernorm, block_size, ) + elif finish_norm is not None: + # The finish convolution and post_attention_layernorm in one launch; the residual is + # updated in place. + hidden_out = finish_norm( + attn_mod.o_proj(attn_output), + attn_conv_coefficients, + self.attention_convs[layer_idx], + layer.post_attention_layernorm, + res_flat, + block_size, + ) else: hidden_out = attn_mod.o_proj(attn_output) if has_block_conv: @@ -445,6 +648,24 @@ def windowed_block_decode( ) continue hidden_out = layer.mlp(hidden_out) + if finish_norm is not None: + # The MLP finish convolution and the next norm (the next layer's input_layernorm, + # or model.norm) in one launch, as the hop does; the residual is updated in place. + next_norm = ( + layers[layer_idx + 1].input_layernorm + if layer_idx + 1 < len(layers) + else self.model.norm + ) + prenormed = finish_norm( + hidden_out, + mlp_conv_coefficients, + self.mlp_convs[layer_idx], + next_norm, + res_flat, + block_size, + ) + residual = res_flat.reshape(B, block_size, -1) + continue if has_block_conv: hidden_out = self.mlp_convs[layer_idx].finish( hidden_out, mlp_conv_coefficients, block_size @@ -453,7 +674,7 @@ def windowed_block_decode( hidden_states = hidden_out.reshape(B, block_size, -1) residual = res_flat.reshape(B, block_size, -1) - if hops is not None: + if hops is not None or finish_norm is not None: return prenormed hidden_states_out, _ = self.model.norm( hidden_states.reshape(-1, hidden_states.shape[-1]), @@ -480,7 +701,14 @@ def install_block_decode(draft_model: nn.Module, enable_pdl: bool) -> str | None return None -def install(draft_model: nn.Module | None, block_conv: bool, block_decode: bool) -> None: +def install( + draft_model: nn.Module | None, + block_conv: bool, + block_decode: bool, + *, + conv_epilogue: bool = False, + finish_norm: bool = False, +) -> None: """Install the fused block convolution and the windowed block decode on ``draft_model``. Each installs independently of the other; the order does not matter @@ -496,9 +724,9 @@ def install(draft_model: nn.Module | None, block_conv: bool, block_decode: bool) TypeError: from ``install_fused_block_conv``. """ if block_conv: - _install_fused_block_conv(draft_model) + _install_fused_block_conv(draft_model, conv_epilogue) if block_decode: - _install_windowed_attention(draft_model, True) + _install_windowed_attention(draft_model, True, finish_norm) def _builtin_block_convs(draft_model: nn.Module) -> int: @@ -513,7 +741,7 @@ def _builtin_block_convs(draft_model: nn.Module) -> int: ) -def _install_fused_block_conv(draft_model: nn.Module | None) -> None: +def _install_fused_block_conv(draft_model: nn.Module | None, conv_epilogue: bool) -> None: fresh = 0 if draft_model is None else _builtin_block_convs(draft_model) count = 0 if draft_model is None else install_fused_block_conv(draft_model) if not count: @@ -524,9 +752,50 @@ def _install_fused_block_conv(draft_model: nn.Module | None) -> None: logger.info( f"[musespark_open] drafter_block_conv: fused DFlash 2 block convolution on {count} modules" ) + if conv_epilogue: + install_conv_epilogue(draft_model) + + +def install_conv_epilogue(draft_model: nn.Module) -> int: + """Give every fused convolution of ``draft_model`` its ``ConvEpilogue``; return how many are new. + + Idempotent (a convolution whose epilogue holds the current projection weight is left alone); + logged by the call that added any. + + Raises: + ValueError: on a convolution the epilogue cannot take (``_conv_epilogue_reason``), on every + rank alike. + """ + convs = [ + (f"{kind}.{index}", conv) + for kind in ("attention_convs", "mlp_convs") + for index, conv in enumerate(getattr(draft_model, kind, None) or ()) + ] + for where, conv in convs: + reason = _conv_epilogue_reason(conv) + if reason is not None: + raise ValueError( + f"[musespark_open] drafter_conv_epilogue cannot run here: {where} is {reason}" + ) + added = 0 + for _, conv in convs: + fused = conv.__dict__.get("_conv_epilogue") + if fused is None or fused.source is not conv.kernel_projection.weight: + conv._conv_epilogue = ConvEpilogue(conv.kernel_projection.weight) + added += 1 + if added: + logger.info( + f"[musespark_open] drafter_conv_epilogue: drafter blocks of {CONV_EPILOGUE_ROWS[0]}.." + f"{CONV_EPILOGUE_ROWS[1]} rows run each convolution's prepare as one launch (coefficient " + f"GEMM on its fixed split-K tactic + input-side convolution) on {len(convs)} modules " + f"({added * 2048 * 8192 * 2 / 2**20:.0f} MiB of tap-interleaved projection weights)" + ) + return added -def _install_windowed_attention(draft_model: nn.Module | None, block_decode: bool) -> None: +def _install_windowed_attention( + draft_model: nn.Module | None, block_decode: bool, finish_norm: bool +) -> None: enable_pdl = LAUNCH_PDL wrapped = isinstance( getattr(draft_model, "_dflash_flash_attention", None), WindowedBlockAttention @@ -549,3 +818,14 @@ def _install_windowed_attention(draft_model: nn.Module | None, block_decode: boo logger.info( "[musespark_open] drafter_block_decode: DFlash block decode with the K/V append fused" ) + if finish_norm: + reason = _finish_norm_reason(draft_model) + if reason is not None: + raise ValueError(f"[musespark_open] drafter_finish_norm cannot run here: {reason}") + if not isinstance(draft_model.__dict__.get("_finish_norm"), FinishNorm): + draft_model._finish_norm = FinishNorm(enable_pdl) + logger.info( + "[musespark_open] drafter_finish_norm: drafter blocks the drafter hops do not " + "take run each hop's finish convolution + residual add + RMSNorm as one launch " + f"(pdl={enable_pdl}, dependents released at entry)" + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py index 80d0f9b9069e..702ce3e2456f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py @@ -136,6 +136,8 @@ "TRTLLM_MUSESPARK_TEP8_EXPERTS_INTERLEAVED_FETCH": ( PREFIX + "TEP8_EXPERTS_INTERLEAVED_FETCH=on" ), + "TRTLLM_MUSESPARK_DRAFTER_CONV_EPILOGUE": PREFIX + "DRAFTER_CONV_EPILOGUE=on", + "TRTLLM_MUSESPARK_DRAFTER_FINISH_NORM": PREFIX + "DRAFTER_FINISH_NORM=on", } ) #: Name prefixes only retired switches carry; any other name under them raises too. @@ -704,6 +706,9 @@ def _vocab_parallel_head(d: Deployment, t: Resolving, p: Mapping[str, object]) - "tep8_experts_interleaved_fetch": _ported_tep8( Kind.FLAGGED, "the expert kernel's interleaved weight fetch" ), + # New this round: default off, not covered by the TP8 default-on gate. + "drafter_conv_epilogue": _ported(Kind.BITWISE), + "drafter_finish_norm": _ported(Kind.BITWISE), **{name: _tep4_path_rule() for name in TEP4_PATHS}, } ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/conv_coef_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/conv_coef_gemm.py new file mode 100644 index 000000000000..5d3242701cc0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/conv_coef_gemm.py @@ -0,0 +1,647 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A DFlash 2 convolution's ``prepare`` in one launch: the coefficient GEMM with the input-side +block convolution in its epilogue. + +``DFlash2BlockConv.prepare`` projects the sublayer input onto ``2 x taps x groups`` coefficients +(``kernel_projection``, ``[2048, 8192]``: side, tap, group) and convolves the input with side 0's +(``dflash2_block_conv``, its own launch between this GEMM and the sublayer's first GEMM); side 1's +go to the sublayer's ``finish``. On the drafter's 9..32-row blocks the GEMM is flashinfer's +cluster split-K kernel at the fixed tactic of ``fixed_low_m_gemm`` (``drafter_fixed_gemms``). + +This entry is that GEMM on ``batch_splitk_gemm``'s copy of the same kernel (same tiles, same +``split_k``, so the same bits: each output element is the FP32 sum of the ``split_k`` K-slice +partials in cluster-rank order, rounded once to BF16) over a row-permuted copy of the weight in +which each side-0 group's two taps are adjacent rows (``interleave``), so a CTA tile holds both +coefficients of each of its groups. After the reduction the CTA that owns a side-0 tile's values +stages them (rounded to BF16, as the GEMM output is) in shared memory and convolves the input +with them, statement for statement as ``dflash2_block_conv``: each coefficient ``base + delta``, +tap product, mask product and running sum rounded to BF16, IEEE FP32 in between, the masked tap-1 +term multiplied by 0.0, row ``t - 1`` read across a block boundary and zero before row 0. +Side-0 coefficients are not stored; side-1 tiles (the original row order) store theirs as the +GEMM does. So the convolution output and the side-1 coefficients are bitwise the GEMM + +``dflash2_block_conv`` pair's. + +Ownership: the reduce-scatter epilogue gives cluster rank ``r`` the values of epilogue thread +group ``r`` (``128 / split_k`` threads, whole warps); a warp's TMEM sub-partition holds a +contiguous quarter of the tile's kernel-M rows (``tmem_warp_shape`` (4, 1)), so rank ``r`` owns +kernel-M rows ``[r M / s, (r + 1) M / s)`` of every public row, an even number of whole groups. + +Programmatic dependent launch: as ``batch_splitk_gemm`` (weight before the wait, activation +after it, dependents released once the weight warp has issued its last load). The epilogue +warps issue every load of the convolution (input rows ``t`` and ``t - 1``, both base taps) +after their own ``griddepcontrol.wait`` and before they wait for the mainloop, so the loads +complete under it and only the arithmetic and the stores follow the reduction. +""" + +from __future__ import annotations + +import functools +from typing import Tuple + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass import Int32, Int64, Uint32 +from cutlass.cute import experimental as cute_ext +from cutlass.cutlass_dsl import dsl_user_op + +from ..allreduce.drafter_hop import ( + _add_rn, + _asm, + _i64, + _keep, + _load_x4_if, + _mul_rn, + _pack_bf16x2, + _round_bf16, + _u32, + _unpack_bf16x2, +) +from . import batch_splitk_gemm as _bsk + +#: The drafter's convolution geometry: hidden 8192, two taps, groups of 16 channels. +HIDDEN, TAPS, GROUP_SIZE = 8192, 2, 16 +GROUPS = HIDDEN // GROUP_SIZE +#: Rows of the coefficient projection per side (taps x groups). +SIDE_ROWS = TAPS * GROUPS +#: Channels one epilogue item covers (one 16-byte vector). +_VECTOR = 8 +_STAGE_ALIGN_BYTES = 128 + + +@dsl_user_op +def _store_x4_if(address, enabled, w0, w1, w2, w3, *, loc=None, ip=None): + """Plain 16-byte global store when ``enabled``; else nothing.""" + _asm( + None, + [_i64(address, loc, ip), _u32(enabled, loc, ip)] + + [_u32(w, loc, ip) for w in (w0, w1, w2, w3)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $1, 0;\n@p st.global.v4.b32 [$0], {$2, $3, $4, $5};\n}", + "l,r,r,r,r,r", + loc, + ip, + ) + + +def interleave(weight: torch.Tensor) -> torch.Tensor: + """The kernel's weight: side 0's rows reordered ``[tap][group] -> [group][tap]``, side 1 as is. + + ``weight`` is a ``kernel_projection`` weight ``[2 * taps * groups, hidden]``. Returns a new + contiguous tensor (the module's weight is not modified). + """ + side0 = weight[:SIDE_ROWS].view(TAPS, GROUPS, -1).transpose(0, 1).reshape(SIDE_ROWS, -1) + return torch.cat([side0, weight[SIDE_ROWS:]], dim=0).contiguous() + + +def stage_bytes(tactic: _bsk.BatchSplitKTactic) -> int: + return tactic.mma_m * tactic.mma_n * 4 + + +def smem_bytes(tactic: _bsk.BatchSplitKTactic, ab_stages: int | None = None) -> int: + """``batch_splitk_gemm``'s layout plus the FP32 staging tile.""" + cursor = _bsk.smem_bytes(tactic, ab_stages) + return _bsk._align_up(cursor, _STAGE_ALIGN_BYTES) + stage_bytes(tactic) + + +def fitted(tactic: _bsk.BatchSplitKTactic) -> _bsk.BatchSplitKTactic: + """``tactic`` with as many of its stages as fit beside the staging tile (pipeline depth only: + the stages do not change any output bit).""" + stages = tactic.ab_stages + while stages > 1 and smem_bytes(tactic, stages) > _bsk.SMEM_CAPACITY_BYTES: + stages -= 1 + return _bsk.BatchSplitKTactic(tactic.mma_m, tactic.mma_n, tactic.split_k, stages) + + +def validate(tactic: _bsk.BatchSplitKTactic, m: int, n: int, k: int) -> None: + """``batch_splitk_gemm.validate`` plus the epilogue's own bounds. + + Raises: + ValueError: naming the first violated bound. + """ + _bsk.validate(tactic, m, n, k) + if n != 2 * SIDE_ROWS or k != HIDDEN: + raise ValueError(f"the convolution epilogue takes a [{2 * SIDE_ROWS}, {HIDDEN}] weight") + rows_owned = tactic.mma_m // tactic.split_k + if rows_owned % 2 or (rows_owned // 2 * 2 * tactic.mma_n) % (128 // tactic.split_k): + raise ValueError(f"tactic {tactic}: a rank's tile share is not whole groups per thread") + if SIDE_ROWS % tactic.mma_m: + raise ValueError(f"mma_m={tactic.mma_m} does not tile one side's {SIDE_ROWS} rows") + if smem_bytes(tactic) > _bsk.SMEM_CAPACITY_BYTES: + raise ValueError(f"tactic {tactic} with its staging tile exceeds shared memory") + + +class ConvCoefGemmKernel(_bsk.BatchSplitKGemmKernel): + """``batch_splitk_gemm``'s kernel with the side-0 convolution in its epilogue.""" + + def __init__(self, *, tactic: _bsk.BatchSplitKTactic, use_pdl: bool) -> None: + super().__init__(tactic=tactic, use_pdl=use_pdl) + self.rows_owned = tactic.mma_m // tactic.split_k + # Items (8 channels of one group of one public row) per owner thread. + self.items = self.rows_owned * tactic.mma_n // self.rs_group + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + x: Int64, + base: Int64, + out: Int64, + rows: Int32, + block: Int32, + stream: _cuda.CUstream, + ): + self.kernel(a, b, c, x, base, out, rows, block).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(_bsk.SMEM_CAPACITY_BYTES), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mC: cute.Tensor, + x: Int64, + base: Int64, + out: Int64, + rows: Int32, + block: Int32, + ): + # batch_splitk_gemm.BatchSplitKGemmKernel.kernel, the epilogue call aside. + stages = self.num_ab_stage + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_bsk._AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_bsk._AB_BUFFER_ALIGN_BYTES, + ) + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + bar_spare = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_bsk._TMEM_POINTER_BYTES, + ).iterator + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.rs_mailbox_elements), + alignment=_bsk._MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + else: + mailbox = sA + bar_reduce = bar_spare + # The owner's BF16-rounded coefficients of the tile, as FP32, kernel-M fastest. + stage = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n), stride=(1, self.cta_m)), + alignment=_STAGE_ALIGN_BYTES, + ) + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + if cutlass.const_expr(self.split_k > 1): + cute.arch.mbarrier_init(bar_reduce, 1) + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_conv_warp( + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + stage, + bidx, + n_idx, + x, + base, + out, + rows, + block, + ) + + @cute.experimental.jit + def epilog_conv_warp( + self, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + stage: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + x: Int64, + base: Int64, + out: Int64, + rows: Int32, + block: Int32, + ): + # batch_splitk_gemm.BatchSplitKGemmKernel.epilog_warp up to the owner's rounding. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + _bsk._tcgen05_fence(before=False) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), acc_layout + )[((None, None), 0, 0, 0)] + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate(self.acc_dtype, cute.AddressSpace.rmem, rmem_layout, alignment=32) + rD = cute_ext.allocate(c_dtype, cute.AddressSpace.rmem, rmem_layout, alignment=32) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + # The convolution's operands, loaded while the mainloop runs: this thread's items of the + # tile's side-0 groups when its group owns them (predicated off, no access, otherwise). + side0 = bidx * self.cta_m < SIDE_ROWS + dest = Int32(0) + lt = epi_tid + if cutlass.const_expr(self.split_k > 1): + dest = epi_tid // self.rs_group + lt = epi_tid % self.rs_group + conv_on = Uint32(side0) & Uint32(dest == split_rank) + if cutlass.const_expr(self.use_pdl): + # The input rows are the predecessor's output. + cute.arch.griddepcontrol_wait() + items = self._prefetch(conv_on, split_rank, lt, bidx, n_idx, x, base, rows) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + _bsk._tcgen05_fence(before=False) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + cute.arch.fence_view_async_tmem_load() + _bsk._tcgen05_fence(before=True) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + if cutlass.const_expr(self.split_k > 1): + vpt = cutlass.const_expr(cute.size(rmem_layout)) + group_values = cutlass.const_expr(self.rs_group * vpt) + if dest != split_rank: + j = split_rank - cutlass.select_(split_rank > dest, Int32(1), Int32(0)) + for value_idx in cutlass.range_constexpr(0, vpt, 4): + slot = (value_idx // 4) * (self.rs_group * 4) + lt * 4 + _bsk._store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + j * group_values + slot, + bar_reduce, + dest, + ) + else: + if lt == 0: + cute.arch.mbarrier_arrive_and_expect_tx(bar_reduce, self.rs_expected_bytes) + _bsk._mbarrier_wait_acquire_cluster(bar_reduce, Int32(0)) + for value_idx in cutlass.range_constexpr(vpt): + slot = (value_idx // 4) * (self.rs_group * 4) + lt * 4 + value_idx % 4 + own = rAcc[value_idx] + acc = own + for r in cutlass.range_constexpr(self.split_k): + jr = cutlass.min( + Int32(r) - cutlass.select_(Int32(r) > split_rank, Int32(1), Int32(0)), + Int32(self.split_k - 2), + ) + term = cutlass.select_( + Int32(r) == split_rank, own, mailbox[jr * group_values + slot] + ) + if cutlass.const_expr(r == 0): + acc = term + else: + acc = acc + term + rAcc[value_idx] = acc + rD.store(rAcc.load().to(c_dtype)) + if side0: + # The coefficients as the GEMM output holds them (BF16), widened back. + rAcc.store(rD.load().to(self.acc_dtype)) + cute_ext.partition_and_copy(thr_t2r, rAcc, stage) + else: + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + else: + rD.store(rAcc.load().to(c_dtype)) + if side0: + rAcc.store(rD.load().to(self.acc_dtype)) + cute_ext.partition_and_copy(thr_t2r, rAcc, stage) + else: + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + + if side0: + if dest == split_rank: + # The owner group's staged rows are complete for every one of its threads. + cute.arch.barrier(barrier_id=1, number_of_threads=self.rs_group) + self._convolve(stage, items, out, block) + + def _prefetch( + self, + conv_on: Uint32, + split_rank: cutlass.Int32, + lt: cutlass.Int32, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + x: Int64, + base: Int64, + rows: Int32, + ): + """This thread's items (8 channels of one owned group of one public row): coordinates and + the input rows ``t``, ``t - 1`` and both base taps, all loads issued before any use.""" + groups_owned = self.rows_owned // 2 + m_first = split_rank * self.rows_owned + items = [] + for i in range(self.items): + item = lt + i * self.rs_group + half = item % 2 + gl = (item // 2) % groups_owned + n = item // (2 * groups_owned) + t = n_idx * self.cta_n + n + m = m_first + 2 * gl + group = (bidx * self.cta_m + m) // 2 + column_bytes = Int64(group * GROUP_SIZE + half * _VECTOR) * 2 + row_bytes = Int64(t) * (HIDDEN * 2) + on = conv_on & Uint32(t < rows) + cur = _load_x4_if(x + row_bytes + column_bytes, on) + prev = _load_x4_if(x + row_bytes - HIDDEN * 2 + column_bytes, on & Uint32(t >= 1)) + b0 = _load_x4_if(base + column_bytes, on) + b1 = _load_x4_if(base + Int64(HIDDEN * 2) + column_bytes, on) + items.append((t, m, n, row_bytes + column_bytes, Uint32(t < rows), cur, prev, b0, b1)) + return items + + def _convolve(self, stage: cute.Tensor, items, out: Int64, block: Int32): + """``dflash2_block_conv``'s statements on the prefetched items, coefficients from the stage + (rows past the call's are computed and not stored: no branch).""" + for t, m, n, offset, valid, cur, prev, b0, b1 in items: + d0 = stage[m, n] + d1 = stage[m + 1, n] + keep = _keep(t % block) + words = [] + for k in range(_VECTOR // 2): + pair = [] + for e in range(2): + coefficient = _round_bf16(_add_rn(_unpack_bf16x2(b0[k])[e], d0)) + acc = _round_bf16(_mul_rn(coefficient, _unpack_bf16x2(cur[k])[e])) + coefficient = _round_bf16(_add_rn(_unpack_bf16x2(b1[k])[e], d1)) + term = _round_bf16(_mul_rn(coefficient, _unpack_bf16x2(prev[k])[e])) + term = _round_bf16(_mul_rn(term, keep)) + pair.append(_round_bf16(_add_rn(acc, term))) + words.append(_pack_bf16x2(pair[0], pair[1])) + _store_x4_if(out + offset, valid, words[0], words[1], words[2], words[3]) + + +@cute.experimental.jit +def _bmm_conv( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + x: Int64, + base: Int64, + out: Int64, + rows: Int32, + block: Int32, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + x, + base, + out, + rows, + block, + stream, + ) + + +@functools.lru_cache(maxsize=None) +def _compiled( + tactic: _bsk.BatchSplitKTactic, use_pdl: bool, leading: Tuple[int, int, int], device: int +): + m, n, k = 64, 8, _bsk.CTA_K + dtype = torch.bfloat16 + with torch.cuda.device(device): + reprs = tuple( + _bsk._from_dlpack_dynamic(_bsk._make_layout_tensor(shape, dtype, d), d) + for shape, d in ( + ((1, n, k), leading[0]), + ((1, k, m), leading[1]), + ((1, n, m), leading[2]), + ) + ) + return cute_ext.compile( + _bmm_conv, + ConvCoefGemmKernel(tactic=tactic, use_pdl=use_pdl), + *reprs, + Int64(0), + Int64(0), + Int64(0), + Int32(8), + Int32(8), + _cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +def run( + x: torch.Tensor, + weight: torch.Tensor, + base0: torch.Tensor, + block_size: int, + tactic: _bsk.BatchSplitKTactic, + pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """``prepare``'s two outputs from ``x [M, 8192]`` and an ``interleave``d weight ``[2048, 8192]``. + + Returns ``(conv, coefficients)``: the input-side convolution ``[M, 8192]`` (bitwise + ``dflash2_block_conv(x, coef[:, 0], base0, block_size, 16)`` of the GEMM's coefficients) and + the ``[M, 2048]`` coefficient buffer whose side-1 half holds the GEMM's side-1 coefficients in + the original order (side 0's half is not written). ``base0`` is the input-side base kernel + ``[2, 8192]``, contiguous. ``tactic`` is taken with its stages cut to fit (``fitted``). + + Raises: + ValueError: for operands or a tactic the kernel does not take. + """ + tactic = fitted(tactic) + if not (x.dtype == weight.dtype == base0.dtype == torch.bfloat16): + raise ValueError("the convolution epilogue takes BF16 operands only") + if x.dim() != 2 or weight.dim() != 2 or tuple(base0.shape) != (TAPS, HIDDEN): + raise ValueError( + f"shapes x {tuple(x.shape)}, weight {tuple(weight.shape)}, base {tuple(base0.shape)}" + ) + m, k = (int(d) for d in x.shape) + n = int(weight.shape[0]) + if int(weight.shape[1]) != k: + raise ValueError(f"shapes x {tuple(x.shape)}, weight {tuple(weight.shape)}") + if not (x.is_contiguous() and weight.is_contiguous() and base0.is_contiguous()): + raise ValueError("the convolution epilogue takes contiguous operands") + if any(t.data_ptr() % 32 for t in (x, weight)) or base0.data_ptr() % 16: + raise ValueError("the convolution epilogue takes 32-byte aligned operands (base 16)") + if not (x.is_cuda and weight.device == x.device == base0.device): + raise ValueError("the convolution epilogue takes CUDA operands on one device") + validate(tactic, m, n, k) + coefficients = torch.empty((m, n), dtype=x.dtype, device=x.device) + conv = torch.empty((m, k), dtype=x.dtype, device=x.device) + with torch.cuda.device(x.device): + tensors, leading = _bsk._swapped(x, weight, coefficients) + compiled = _compiled(tactic, bool(pdl), leading, x.get_device()) + compiled( + *tensors, + Int64(x.data_ptr()), + Int64(base0.data_ptr()), + Int64(conv.data_ptr()), + Int32(m), + Int32(block_size), + _cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream), + ) + return conv, coefficients diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_finish_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_finish_norm.py new file mode 100644 index 000000000000..38837eb13c94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/drafter_finish_norm.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A DFlash 2 drafter hop's finish convolution, residual add and RMSNorm in one launch. + +After each all-reduced drafter projection (``o_proj``, ``down_proj``) the built-in block decode +runs two kernels on drafter blocks the fused drafter hop does not take (``drafter_hops``: 8 rows; +``wide_drafter_hops``: up to its row bound): the finish-side block convolution +(``ops/spec/dflash2_block_conv.py``, ``rows x 8`` programs) and flashinfer's CuTe-DSL fused +residual add + RMSNorm of the next norm (``FusedAddRMSNormKernel``, which the built-in +``RMSNorm`` reaches on this GPU; in place on both operands). This entry is the two in one +kernel, bit for bit: the fused drafter hop's kernel (``ops/allreduce/drafter_hop.py``, its +thread map, arithmetic and helpers) with the all-reduce transport left out -- rows ``t`` and +``t - 1`` of the already reduced projection are read from global memory. + +* convolution: ``dflash2_block_conv``'s statements -- each coefficient ``base + delta``, tap + product, mask product and running sum rounded to BF16 where it rounds, IEEE FP32 add / + multiply in between, the masked tap-1 term multiplied by 0.0 (NaN and signed zeros propagate + as the torch chain's do), row ``t - 1`` read across a block boundary and zero before row 0; +* add + RMSNorm: ``h = conv + r`` in FP32 (not rounded); per flashinfer thread ``f`` of the + row's 128, one FMA chain ``s = fma(h, h, s)`` from ``+0`` over its 64 columns + ``8 f + 1024 b + e`` (``b`` then ``e`` ascending); the butterfly over each warp (lane offsets + 1, 2, 4, 8, 16) and over the 4 warp sums padded with zeros; + ``rstd = rsqrt.approx.ftz(fma(s, 2**-13, eps))``; ``y = (w + 0) * (h * rstd)``; the residual + becomes ``bf16(h)``. Every floating-point operation is an inline PTX instruction. + +Thread map (the hop's): a row is one cluster of 8 CTAs of 128 threads; CTA ``c`` owns +flashinfer threads ``16 c .. 16 c + 15``, its thread ``j`` block ``j mod 8`` of flashinfer +thread ``16 c + j / 8`` (8 columns at ``8 f + 1024 b``); the 16 chain sums of each CTA go to +every CTA of the cluster through distributed shared memory, and every CTA runs both butterflies +on the 128 sums. Every global load of a thread is issued before its first use. + +Programmatic dependent launch: the norm weight and the base kernel (model parameters) are read +before ``griddepcontrol.wait``; the projection rows, the coefficients and the residual after +it. The dependents are released at entry (``entry_trigger``) or after the stores. +""" + +from __future__ import annotations + +import functools + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper + +from ..allreduce import drafter_hop as _hop + +HIDDEN = _hop.HIDDEN +TAPS = 2 +GROUP_SIZE = 16 + +if _hop.HAS_CUTE_DSL: + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + import cutlass.utils as utils + from cutlass import Float32, Int32, Int64, Uint32 + + from ..allreduce.drafter_hop import ( + CLUSTER, + OWNER_STRIDE, + OWNERS, + ROW_BYTES, + SMEM_BYTES, + SUMS_BYTES, + THREADS, + VECTOR, + WARPS, + _add_rn, + _bf16_bits_to_f32, + _float_bits, + _fma_rn, + _keep, + _load_u16, + _load_x4, + _load_x4_if, + _mul_rn, + _pack_bf16x2, + _round_bf16, + _rsqrt_ftz, + _shfl_bfly, + _store_x4, + _unpack_bf16x2, + ) + + class _FinishNormKernel: + """Finish convolution + flashinfer add + RMSNorm, one row per cluster of 8 CTAs.""" + + def __init__(self, eps: float, pdl: bool, entry_trigger: bool): + self.eps = float(eps) + self.pdl = pdl + self.entry_trigger = entry_trigger + + @cute.jit + def __call__( + self, + x: Int64, + residual: Int64, + delta: Int64, + base: Int64, + weight: Int64, + output: Int64, + rows: Int32, + block: Int32, + x_row: Int32, + residual_row: Int32, + delta_row: Int32, + delta_tap: Int32, + stream: cuda.CUstream, + ): + self.kernel( + x, + residual, + delta, + base, + weight, + output, + rows, + block, + x_row, + residual_row, + delta_row, + delta_tap, + ).launch( + grid=[rows * CLUSTER, 1, 1], + block=[THREADS, 1, 1], + cluster=[CLUSTER, 1, 1], + smem=SMEM_BYTES, + stream=stream, + use_pdl=self.pdl, + ) + + @cute.kernel + def kernel( + self, + x: Int64, + residual: Int64, + delta: Int64, + base: Int64, + weight: Int64, + output: Int64, + rows: Int32, + block: Int32, + x_row: Int32, + residual_row: Int32, + delta_row: Int32, + delta_tap: Int32, + ): + tid = cute.arch.thread_idx()[0] + lane = tid % 32 + warp = tid // 32 + crank = cute.arch.block_idx()[0] % CLUSTER + row = cute.arch.block_idx()[0] // CLUSTER + owner = tid // CLUSTER + block_b = tid % CLUSTER + column = (crank * OWNERS + owner) * VECTOR + block_b * (THREADS * VECTOR) + smem = utils.SmemAllocator() + h_local = smem.allocate_tensor( + Float32, cute.make_layout((OWNERS * OWNER_STRIDE,)), byte_alignment=16 + ) + chains = smem.allocate_tensor(Float32, cute.make_layout((OWNERS,)), byte_alignment=16) + all_chains = smem.allocate_tensor( + Float32, cute.make_layout((THREADS,)), byte_alignment=16 + ) + sums = smem.allocate_tensor(Float32, cute.make_layout((WARPS,)), byte_alignment=16) + bars = smem.allocate_tensor(cutlass.Uint64, cute.make_layout((1,)), byte_alignment=8) + if tid == 0: + cute.arch.mbarrier_init(bars.iterator, 1) + cute.arch.mbarrier_arrive_and_expect_tx(bars.iterator, SUMS_BYTES) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + cute.arch.cluster_arrive_relaxed() + + # Model parameters (written at load time only) before the wait. + column_bytes = Int64(column) * 2 + w = _load_x4(weight + column_bytes) + b0 = _load_x4(base + column_bytes) + b1 = _load_x4(base + Int64(ROW_BYTES) + column_bytes) + if cutlass.const_expr(self.pdl): + if cutlass.const_expr(self.entry_trigger): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + has_prev = Uint32(row != 0) + x_offset = Int64(row) * Int64(x_row) * 2 + column_bytes + current_row = _load_x4(x + x_offset) + previous_row = _load_x4_if(x + x_offset - Int64(x_row) * 2, has_prev) + residual_offset = Int64(row) * Int64(residual_row) * 2 + column_bytes + r = _load_x4(residual + residual_offset) + group = Int64(column // GROUP_SIZE) + delta_address = delta + (Int64(row) * Int64(delta_row) + group) * 2 + d0 = _bf16_bits_to_f32(_load_u16(delta_address)) + d1 = _bf16_bits_to_f32(_load_u16(delta_address + Int64(delta_tap) * 2)) + # Peers' barriers are initialized before any distributed store below. + cute.arch.cluster_wait() + + keep = _keep(row % block) + h = cute.make_rmem_tensor(cute.make_layout((VECTOR,)), Float32) + for e in cutlass.range_constexpr(VECTOR): + k = e // 2 + current = _unpack_bf16x2(current_row[k])[e % 2] + shifted = _unpack_bf16x2(previous_row[k])[e % 2] + # The finish-side block convolution, rounded where dflash2_block_conv rounds. + coefficient = _round_bf16(_add_rn(_unpack_bf16x2(b0[k])[e % 2], d0)) + acc = _round_bf16(_mul_rn(coefficient, current)) + coefficient = _round_bf16(_add_rn(_unpack_bf16x2(b1[k])[e % 2], d1)) + term = _round_bf16(_mul_rn(coefficient, shifted)) + term = _round_bf16(_mul_rn(term, keep)) + acc = _round_bf16(_add_rn(acc, term)) + # flashinfer's residual add. + h[e] = _add_rn(acc, _unpack_bf16x2(r[k])[e % 2]) + + # Block ``block_b`` of local flashinfer thread ``owner``: at [owner][block_b]. + for e in cutlass.range_constexpr(VECTOR): + h_local[owner * OWNER_STRIDE + block_b * VECTOR + e] = h[e] + cute.arch.barrier() + # Local flashinfer thread ``tid``'s FMA chain over blocks 0..7, 8 columns each. + if tid < OWNERS: + chain = Float32(0.0) + for k in cutlass.range_constexpr(CLUSTER * VECTOR): + value = h_local[tid * OWNER_STRIDE + k] + chain = _fma_rn(value, value, chain) + chains[tid] = chain + cute.arch.barrier() + # The 16 sums of flashinfer threads 16 crank .. 16 crank + 15, to every CTA of the + # row's cluster (this one included), 4 per store. + if tid < OWNERS // 4: + target = all_chains.iterator + crank * OWNERS + tid * 4 + for peer in cutlass.range_constexpr(CLUSTER): + cute.arch.store_async_dsmem( + target.align(16), + ( + _float_bits(chains[4 * tid]), + _float_bits(chains[4 * tid + 1]), + _float_bits(chains[4 * tid + 2]), + _float_bits(chains[4 * tid + 3]), + ), + bars.iterator, + peer, + ) + cute.arch.mbarrier_wait(bars.iterator, 0) + # Lane ``l`` of warp ``w`` holds flashinfer thread ``32 w + l``'s sum. + sum_sq = all_chains[tid] + for i in cutlass.range_constexpr(5): + sum_sq = _add_rn(sum_sq, _shfl_bfly(sum_sq, Uint32(1 << i))) + if lane == 0: + sums[warp] = sum_sq + cute.arch.barrier() + row_sum = Float32(0.0) + if lane < WARPS: + row_sum = sums[lane] + for i in cutlass.range_constexpr(5): + row_sum = _add_rn(row_sum, _shfl_bfly(row_sum, Uint32(1 << i))) + rms_rcp = _rsqrt_ftz(_fma_rn(row_sum, Float32(1.0 / HIDDEN), Float32(self.eps))) + + updated = [] + normed = [] + for k in cutlass.range_constexpr(4): + lo_w, hi_w = _unpack_bf16x2(w[k]) + lo = _mul_rn(_add_rn(lo_w, Float32(0.0)), _mul_rn(h[2 * k], rms_rcp)) + hi = _mul_rn(_add_rn(hi_w, Float32(0.0)), _mul_rn(h[2 * k + 1], rms_rcp)) + normed.append(_pack_bf16x2(lo, hi)) + updated.append(_pack_bf16x2(h[2 * k], h[2 * k + 1])) + _store_x4(residual + residual_offset, *updated) + _store_x4(output + Int64(row) * (HIDDEN * 2) + column_bytes, *normed) + if cutlass.const_expr(self.pdl and not self.entry_trigger): + cute.arch.griddepcontrol_launch_dependents() + + @functools.lru_cache(maxsize=64) + def _compiled_kernel(device_index: int, eps: float, pdl: bool, entry_trigger: bool): + kernel = _FinishNormKernel(eps, pdl, entry_trigger) + with torch.cuda.device(device_index): + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile( + kernel, + *([Int64(0)] * 6), + Int32(1), + Int32(8), + Int32(HIDDEN), + Int32(HIDDEN), + Int32(2 * TAPS * HIDDEN // GROUP_SIZE), + Int32(HIDDEN // GROUP_SIZE), + stream, + options="--enable-tvm-ffi", + ) + + +def takes( + x: torch.Tensor, + delta: torch.Tensor, + base: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, +) -> bool: + """Whether this entry computes these operands (else the caller runs the built-in two).""" + rows = x.shape[0] if x.dim() == 2 else -1 + return ( + _hop.HAS_CUTE_DSL + and rows >= 1 + and all(t.dtype == torch.bfloat16 and t.is_cuda for t in (x, delta, base, residual, weight)) + and all(t.device == x.device for t in (delta, base, residual, weight)) + and tuple(x.shape) == (rows, HIDDEN) + and tuple(residual.shape) == (rows, HIDDEN) + and tuple(delta.shape) == (rows, TAPS, HIDDEN // GROUP_SIZE) + and tuple(base.shape) == (TAPS, HIDDEN) + and tuple(weight.shape) == (HIDDEN,) + and x.stride(1) == 1 + and residual.stride(1) == 1 + and delta.stride(2) == 1 + and base.is_contiguous() + and weight.is_contiguous() + and all(t.data_ptr() % 16 == 0 for t in (base, weight)) + and all(t.stride(0) % 8 == 0 and t.data_ptr() % 16 == 0 for t in (x, residual)) + and not torch.is_grad_enabled() + ) + + +class DrafterFinishNorm(OpWrapper): + """``y, residual <- rmsnorm(conv(x) + residual)``: the drafter hop without its all-reduce. + + Arguments: ``x`` the all-reduced projection output ``[rows, 8192]``; ``delta`` the finish + coefficients ``[rows, 2, 512]`` (``coefficients[:, 1]`` of the convolution's ``prepare``, + unit group stride); ``base`` the finish base kernel ``[2, 8192]`` (a parameter: read before + the dependency wait); ``residual`` ``[rows, 8192]``, updated in place to + ``bf16(conv(x) + residual)``; ``weight`` / ``eps`` the next RMSNorm's (a parameter); + ``block_size`` rows per request; ``pdl`` launch with programmatic dependent launch; + ``entry_trigger`` release the dependents at entry instead of after the stores. Returns the + normed rows, a fresh tensor. + + Fusion boundary: the convolution and the add + norm; the all-reduce and the next GEMM stay + outside. + """ + + ARCHS = frozenset({Arch.SM_103}) + + note = """ + Metadata consumed: none. One CuTe-DSL kernel, compiled per (device, eps, pdl, trigger): + grid `rows x 8` CTAs of 128 threads in clusters of 8 (one row per cluster), the fused + drafter hop's thread map without its transport. + + Call site: `drafter_block_decode/block_decode.py` `windowed_block_decode`, both hops of + every drafter layer, on blocks the fused drafter hop does not take, when + `TRTLLM_MUSESPARK_DRAFTER_FINISH_NORM=1` (`FinishNorm`). + + Numerics: bitwise the built-in `dflash2_block_conv` then flashinfer's CuTe-DSL fused add + + RMSNorm for a 8192-wide BF16 row (the fused drafter hop's arithmetic). `reference` is + float64 between the rounding points. + + Hard limits + hidden 8192, two taps over groups of 16, BF16, 16-byte aligned rows (`takes`). + """ + + def __call__( + self, + x: torch.Tensor, + delta: torch.Tensor, + base: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + block_size: int, + pdl: bool, + entry_trigger: bool = False, + ) -> torch.Tensor: + rows = x.shape[0] + out = torch.empty((rows, HIDDEN), dtype=x.dtype, device=x.device) + if rows == 0: + return out + compiled = _compiled_kernel(x.get_device(), float(eps), bool(pdl), bool(entry_trigger)) + compiled( + x.data_ptr(), + residual.data_ptr(), + delta.data_ptr(), + base.data_ptr(), + weight.data_ptr(), + out.data_ptr(), + rows, + block_size, + x.stride(0), + residual.stride(0), + delta.stride(0), + delta.stride(1), + cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream), + ) + return out + + def reference( + self, + x: torch.Tensor, + delta: torch.Tensor, + base: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + block_size: int, + pdl: bool, + entry_trigger: bool = False, + ) -> torch.Tensor: + """The rounding points, float64 in between; updates ``residual`` as the op does. + + ``pdl`` and ``entry_trigger`` change how the op launches, not what it computes; ignored. + """ + rows = x.shape[0] + xf = x.float() + coef = (base.float().unsqueeze(0) + delta.float().repeat_interleave(GROUP_SIZE, -1)).to( + torch.bfloat16 + ) + position = torch.arange(rows, device=x.device) % block_size + acc = (coef[:, 0].float() * xf).to(torch.bfloat16).float() + shifted = torch.zeros_like(xf) + shifted[1:] = xf[:-1] + term = (coef[:, 1].float() * shifted).to(torch.bfloat16).float() + term = term * (position >= 1).float().unsqueeze(-1) + acc = (acc + term).to(torch.bfloat16).double() + s = acc + residual.double() + normed = s * torch.rsqrt(s.square().mean(-1, keepdim=True) + eps) * weight.double() + residual.copy_(s.to(torch.bfloat16)) + return normed.to(torch.bfloat16) + + def is_valid( + self, + x: torch.Tensor, + delta: torch.Tensor, + base: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + block_size: int, + pdl: bool, + entry_trigger: bool = False, + ) -> None: + assert takes(x, delta, base, residual, weight), ( + "drafter_finish_norm takes BF16 x / residual [rows, 8192] (16-byte aligned rows), " + "delta [rows, 2, 512] with unit group stride, contiguous base [2, 8192] and weight " + f"[8192]; got x {tuple(x.shape)} {x.dtype}, delta {tuple(delta.shape)} " + f"{delta.stride()}, base {tuple(base.shape)}, residual {tuple(residual.shape)}, " + f"weight {tuple(weight.shape)}" + ) + assert block_size >= 1, f"block_size must be >= 1, got {block_size}" + + +drafter_finish_norm = DrafterFinishNorm() From 4f3322aef81b0a02aa212c25e97aa618da082985 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Mon, 9 Nov 2026 21:41:00 -0800 Subject: [PATCH 5/9] [None][feat] MuseSpark GB300: drafter GEMM cells, SiLU epilogue and attention-block split-K Three more default-off dispatch levers on the sm103 TEP8 decode step: * drafter_multi_request_cells (on drafter_fixed_gemms): the drafter GEMMs of 16, 24 and 32-row blocks run measured TP8 split-K cells. The qkv, gate_up and coefficient cells keep their table row's split_k and are bitwise that row; the o_proj and down_proj cells are single K-slice GEMMs that reproduced the library call's bits on every checked shape, which cannot be proven against a closed library, so the lever is flagged rather than called bitwise. * drafter_silu_epilogue (on drafter_block_decode): gate_up_proj and silu(gate) * up of 9..32-row blocks are one launch over an interleaved copy of the gate_up weights, which costs about 560 MiB per generation rank. * tep8_attn_block_splitk (on dense_gemms): the attention block's QKV and o_proj of 9..32-row decode calls run batch split-K cells, so each stages its weight tile before its dependency wait. o_proj at split_k 1 keeps the library call's summation; QKV at split_k 4 keeps the served cell's. A tactic the drafter table marks OnBatchSplitK is run on this target's own split-K copy, whose tiles and summation match the served kernel at equal split_k. Every lever is knob-only, so an environment that sets nothing runs exactly the paths it ran before; each environment switch this work carried is retired by name in levers.RETIRED. Proved by the AST verifier: rebuilding FAST_PATHS, DEPENDS_ON and both compositions and resolving with an EMPTY environment gives path sets identical to the base at all twelve deployments, and each lever then turns on by its own knob. Run on GB300 (2 nodes, TEP8 holder): all 13 drafter cells validate against the split-K kernel's own bounds; the whole attention-block table passes its checker and is refused at TP4; gate_up_silu_gemm is bitwise the GEMM plus silu_and_mul it replaces at 16, 24 and 32 rows (0 differing elements of 57344, 86016 and 114688), once the reference does its arithmetic in FP32 and rounds once; every attention-block cell compiles and runs, and the o_proj cells at split_k 1 are bitwise cuBLAS at 9, 17 and 25 rows. The model itself was not booted. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 8bb1c39fa67854227cef8030a29737ad7103903d) --- .../dispatch.py | 19 + .../drafter.py | 1 + .../kernels/dense_gemms/dispatcher.py | 216 ++++++- .../drafter_block_decode/block_decode.py | 170 +++++- .../lever_settings.py | 12 +- .../levers.py | 9 + .../ops/gemm/fixed_low_m_gemm.py | 88 +++ .../ops/gemm/gate_up_silu_gemm.py | 569 ++++++++++++++++++ 8 files changed, 1071 insertions(+), 13 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/gate_up_silu_gemm.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py index 237e6c39c4a9..23ca2672e63b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py @@ -424,6 +424,25 @@ class FastPath: "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( + "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", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py index 2ab5cfee705d..6f3079f5e7d9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py @@ -28,6 +28,7 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None: 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. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py index 2cc4c7ad452d..8399cebef02e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py @@ -32,6 +32,16 @@ names on ``ops/gemm/batch_splitk_gemm`` instead of the built-in cuBLAS call; the pooler and unpooler slices of those rows run it through ``projection_publish`` (``batch_splitk_apply``). +With ``tep8_attn_block_splitk`` (sm103 TEP8 only) the +attention block's two projections of a 9..32-row decode call run on ``ops/gemm/batch_splitk_gemm`` +(this target's copy of flashinfer's split-K kernel, launched with programmatic dependent launch) +at the cells ``SM103_TP8_ATTN_BLOCK_KERNELS`` names: the attention output projection at split_k 1 +(one K slice, cuBLAS's summation order) instead of the built-in cuBLAS call, so its weight k-tiles +are loaded into shared memory before its ``griddepcontrol.wait``, under the attention that +launches it early (the sliding-window cluster block, the full layers' combine); and QKV at split_k +4 (flashinfer's K slices and order) with a deeper ring where the copy's reduce-scatter mailbox +leaves room. Every row keeps its bits; the forks and joins are the served ones. + With ``dense_gemm_prefetch`` a projection that ran also forks a side stream that asks the L2 to fetch the next projection's weight (QKV -> the same layer's attention output -> the same layer's pooler slice; the unpooler -> the next layer's QKV), so the @@ -79,6 +89,8 @@ CUBLAS, DIRECT, MAX_ROWS, + O_PROJ_SHAPE, + QKV_SHAPE, SPLITK, WIDE_MAX_ROWS, dense_gemm, @@ -163,17 +175,19 @@ _BATCH_SPLITK = False -def configure(*, batch_splitk: bool = False) -> None: - """Set whether ``install`` arms the batched split-K table. +def configure(*, batch_splitk: bool = False, attn_block_splitk: bool = False) -> None: + """Set whether ``install`` arms the batched split-K tables. - ``lever_settings.apply`` sets it from the model's plan before any install; the default - keeps the built-in cuBLAS calls at every row count above the wide band. + ``lever_settings.apply`` sets them from the model's plan before any install; the defaults + keep the built-in cuBLAS calls above the wide band and the served calls inside it. Args: batch_splitk: the ``tep8_batch_splitk`` lever (``_BATCH_SPLITK``). + attn_block_splitk: the ``tep8_attn_block_splitk`` lever (``_ATTN_BLOCK_SPLITK``). """ - global _BATCH_SPLITK + global _BATCH_SPLITK, _ATTN_BLOCK_SPLITK _BATCH_SPLITK = bool(batch_splitk) + _ATTN_BLOCK_SPLITK = bool(attn_block_splitk) #: ``batch_splitk`` on sm103 at TP8 (per-rank QKV ``[2304, 8192]``, pooler slice ``[512, 8192]``, @@ -231,6 +245,105 @@ def configure(*, batch_splitk: bool = False) -> None: ] = {((10, 3), 8): SM103_TP8_BATCH_KERNELS} +#: ``tep8_attn_block_splitk``, pushed here by ``lever_settings.apply`` (default off). On sm103 at +#: TP8 with the sharded pooler and unpooler it runs the 9..32-row calls of the sites and rows +#: ``SM103_TP8_ATTN_BLOCK_KERNELS`` names on ``ops/gemm/batch_splitk_gemm``; anywhere else it +#: changes nothing but a warning. +_ATTN_BLOCK_SPLITK = False + +#: ``attn_block_splitk`` on sm103 at TP8, as ``(last row count, kernel)`` steps from 9 rows (``None``: +#: the served call at those rows). ``bsplitk:,,,`` with one row tile +#: per call (``mma_n`` covers the rows), so each weight byte is streamed once. +#: +#: * o_proj ``[8192, 1024]``: split_k 1, the bits of the built-in cuBLAS call (which does not split +#: K = 1024 at these rows); ``mma_m`` 64, 128 CTAs. Launched with PDL right behind the attention's +#: early trigger, it loads its weight k-tiles before its wait. +#: * QKV ``[2304, 8192]``: split_k 4, the served cell's K slices and order (the reduce-scatter +#: copy is bitwise with flashinfer's kernel at equal split_k), one more stage than flashinfer's +#: cell and ``mma_n`` 24 at 17..24 rows; 9..16 rows keep flashinfer's cell (the copy was slower there). +#: +#: Measured on the E-MK chain of kernel stack 3 (62 layers, 8 ranks, two tray pairs; mb-emk-h217-1..4). +SM103_TP8_ATTN_BLOCK_KERNELS: Mapping[str, Tuple[Tuple[int, Optional[str]], ...]] = { + "qkv": ((16, None), (24, "bsplitk:64,24,4,10"), (WIDE_MAX_ROWS, "bsplitk:64,32,4,9")), + "o_proj": ( + (16, "bsplitk:64,16,1,8"), + (24, "bsplitk:64,24,1,8"), + (WIDE_MAX_ROWS, "bsplitk:64,32,1,8"), + ), +} + +#: The ``attn_block_splitk`` table of each ``(compute capability, TP size)`` it is measured for. +ATTN_BLOCK_KERNELS_BY_SM_TP: Mapping[ + Tuple[Tuple[int, int], int], Mapping[str, Tuple[Tuple[int, Optional[str]], ...]] +] = {((10, 3), 8): SM103_TP8_ATTN_BLOCK_KERNELS} + +#: The ``split_k`` an ``attn_block_splitk`` cell of each site must have: QKV's 8-row cell's (flashinfer +#: split-K 4), and 1 for o_proj, whose 8-row cell is cuBLAS, which sums K = 1024 in one slice. +_ATTN_BLOCK_SPLIT_K = {"qkv": 4, "o_proj": 1} + + +def attn_block_kernels( + table: Mapping[str, Tuple[Tuple[int, Optional[str]], ...]] = SM103_TP8_ATTN_BLOCK_KERNELS, +) -> Dict[str, Tuple[Optional[str], ...]]: + """The ``attn_block_splitk`` kernel of each site per row count 9..32 (index ``rows - WIDE_MIN_ROWS``), + ``None`` where the site keeps its served call.""" + out = {} + for site, steps in table.items(): + kernels: List[Optional[str]] = [] + first = WIDE_MIN_ROWS + for last, kernel in steps: + assert first <= last <= WIDE_MAX_ROWS, steps + kernels += [kernel] * (last - first + 1) + first = last + 1 + assert first == WIDE_MAX_ROWS + 1, steps + out[site] = tuple(kernels) + return out + + +def attn_block_kernel_table( + capability: Tuple[int, int], tp_size: int, sharded: bool +) -> Dict[str, Tuple[Optional[str], ...]]: + """``attn_block_splitk``'s kernel per row count 9..32 of each site its table names on ``capability`` + at ``tp_size``, each cell checked: its ``split_k`` (``_ATTN_BLOCK_SPLIT_K``), one row tile at every + row count it serves, and the site's shape (``batch_splitk_gemm.validate``). Reads only the tables. + + Raises: + ValueError: if ``ATTN_BLOCK_KERNELS_BY_SM_TP`` has no table for ``(capability, tp_size)``, the + projections are not ``sharded``, or a cell fails a check. + """ + from ...ops.gemm import batch_splitk_gemm as _bsk + + key = (tuple(capability), int(tp_size)) + steps = ATTN_BLOCK_KERNELS_BY_SM_TP.get(key) + if steps is None or not sharded: + raise ValueError( + f"[musespark_open] tep8_attn_block_splitk is measured for SM {(10, 3)} at TP8 with the " + f"sharded " + f"pooler and unpooler; got SM {key[0]} at TP{key[1]}, " + f"{'sharded' if sharded else 'replicated'} projections" + ) + shapes = {"qkv": QKV_SHAPE, "o_proj": O_PROJ_SHAPE} + table = attn_block_kernels(steps) + for role, cells in table.items(): + k, n = shapes[role] + for rows, kernel in enumerate(cells, start=WIDE_MIN_ROWS): + if kernel is None: + continue + tactic = _bsk.parse(kernel) + if tactic.split_k != _ATTN_BLOCK_SPLIT_K[role]: + raise ValueError( + f"[musespark_open] tep8_attn_block_splitk: {role} kernel {kernel} has split_k " + f"{tactic.split_k}; the site's bits need {_ATTN_BLOCK_SPLIT_K[role]}" + ) + if rows > tactic.mma_n: + raise ValueError( + f"[musespark_open] tep8_attn_block_splitk: {role} kernel {kernel} takes {rows} rows in " + f"more than one row tile" + ) + _bsk.validate(tactic, rows, n, k) + return table + + def batch_kernels( table: Mapping[str, Tuple[Tuple[int, str], ...]] = SM103_TP8_BATCH_KERNELS, ) -> Dict[str, Tuple[Optional[str], ...]]: @@ -507,8 +620,12 @@ def __init__( prefetch: bool, enable_pdl: bool, capability: Optional[Tuple[int, int]] = None, + attn_block: Optional[Mapping[str, Tuple[Optional[str], ...]]] = None, ) -> None: self.kernels = dict(kernels) + #: ``attn_block_splitk``: the batch split-K cell of a site's 9..32-row calls + #: (``attn_block[site][rows - WIDE_MIN_ROWS]``, ``None`` where the served call runs); empty when off. + self.attn_block = dict(attn_block or {}) #: The 9..32-row table (``set_wide_kernels``) and whether the current forward arms it. self.wide_kernels: Dict[str, Tuple[str, ...]] = {} self.wide_armed = False @@ -556,6 +673,13 @@ def kernel(self, role: str, rows: int) -> Optional[str]: """``kernel_for``: the kernel of a ``role`` call of ``rows`` rows, or ``None``.""" return self.kernel_for(role, rows) + def attn_block_kernel_for(self, role: str, rows: int) -> Optional[str]: + """``attn_block_splitk``'s cell of a 9..32-row ``role`` call, or ``None`` (the served call).""" + cells = self.attn_block.get(role) + if cells is None or not WIDE_MIN_ROWS <= rows <= WIDE_MAX_ROWS: + return None + return cells[rows - WIDE_MIN_ROWS] + def batch_kernel_for(self, role: str, rows: int) -> Optional[str]: """``batch_splitk``'s kernel of a 33..160-row ``role`` call, or ``None`` (the built-in call).""" cells = self.batch_splitk.get(role) @@ -815,6 +939,10 @@ def apply( return self._apply_substitute(input_tensor, weight) if self.publish is not None: return self._apply_publish(module, input_tensor, weight, bias, force_active) + if plan.attn_block: + out = self._attn_block_apply(input_tensor, weight, bias) + if out is not None: + return out if not self.fits(input_tensor, weight, bias): if input_tensor.dim() >= 1 and self.is_decode_call(input_tensor): # The built-in dispatch picks its kernel per rank: say so once. @@ -855,6 +983,54 @@ def _apply_substitute(self, input_tensor: torch.Tensor, weight: torch.Tensor) -> self.site.plan.fork_prefetch(self.site.next_weight()) return product.view(*input_tensor.shape[:-1], int(weight.shape[0])) + def _attn_block_apply( + self, input_tensor: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] + ) -> Optional[torch.Tensor]: + """``attn_block_splitk``: a 9..32-row QKV or o_proj call on the plan's batch split-K cell for its + rows, or ``None`` (the served call): BF16, no bias, the site's shape, contiguous, 32-byte aligned, + not compiled or differentiated. Joins happened in ``apply``. Forks what the served call forks: + QKV (``fits``' path) the next weight's prefetch; o_proj (the built-in call's path) nothing.""" + plan = self.site.plan + if bias is not None or torch.is_grad_enabled() or torch.compiler.is_compiling(): + return None + if not ( + input_tensor.is_cuda + and input_tensor.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and weight.dim() == 2 + and input_tensor.dim() >= 1 + and (int(weight.shape[1]), int(weight.shape[0])) == self.site.shape + and int(input_tensor.shape[-1]) == int(weight.shape[1]) + and input_tensor.is_contiguous() + and weight.is_contiguous() + and input_tensor.data_ptr() % 32 == 0 + and weight.data_ptr() % 32 == 0 + ): + return None + k = int(weight.shape[1]) + rows = input_tensor.numel() // k + kernel = plan.attn_block_kernel_for(self.site.role, rows) + if kernel is None: + return None + from ...ops.gemm import batch_splitk_gemm as _bsk + + x = input_tensor.view(-1, k) + out = torch.empty((rows, int(weight.shape[0])), dtype=x.dtype, device=x.device) + _bsk.run(x, weight, out, _bsk.parse(kernel), plan.enable_pdl) + if ( + self.site.role != "o_proj" + and plan.prefetch + and self.site.next_weight is not None + and rows not in self.site.skip_prefetch_rows + ): + plan.fork_prefetch(self.site.next_weight()) + logger.info_once( + f"[musespark_open] tep8_attn_block_splitk: a {rows}-row {self.site.role} call (layer " + f"{self.site.layer_idx}) ran {kernel}", + key=f"musespark_open.attn_block_splitk.{self.site.role}.{rows}", + ) + return out.view(*input_tensor.shape[:-1], int(weight.shape[0])) + def batch_splitk_apply( self, input_tensor: torch.Tensor, @@ -1111,7 +1287,23 @@ def install( f"TP8 with the sharded pooler and unpooler; SM {tuple(capability)} at TP{tp_size} " f"({'sharded' if sharded else 'replicated'}) keeps the built-in calls" ) - plan = DenseGemmPlan(kernels, prefetch=prefetch, enable_pdl=LAUNCH_PDL, capability=capability) + attn_table: Mapping[str, Tuple[Optional[str], ...]] = {} + if _ATTN_BLOCK_SPLITK: + if (tuple(capability), tp_size) in ATTN_BLOCK_KERNELS_BY_SM_TP and sharded: + attn_table = attn_block_kernel_table(capability, tp_size, sharded) + else: + logger.warning( + "[musespark_open] tep8_attn_block_splitk is measured for SM (10, 3) at " + f"TP8 with the sharded pooler and unpooler; SM {tuple(capability)} at TP{tp_size} " + f"({'sharded' if sharded else 'replicated'}) keeps the served calls" + ) + plan = DenseGemmPlan( + kernels, + prefetch=prefetch, + enable_pdl=LAUNCH_PDL, + capability=capability, + attn_block=attn_table, + ) plan.tp_size = tp_size if batch_table: plan.set_batch_splitk_kernels(batch_table) @@ -1190,4 +1382,16 @@ def install( f"{BATCH_MAX_ROWS} rows by {steps} (split_k 4, each row the bits of its {MAX_ROWS}-row " "kernel); o_proj and the rows past each site's last step keep the built-in cuBLAS call" ) + if attn_table: + table = ATTN_BLOCK_KERNELS_BY_SM_TP[(tuple(capability), tp_size)] + steps = "; ".join( + f"{site}=" + ", ".join(f"{k or 'served'}@..{last}" for last, k in table[site]) + for site in table + ) + logger.info( + f"[musespark_open] tep8_attn_block_splitk: TP8 calls of {WIDE_MIN_ROWS}.." + f"{WIDE_MAX_ROWS} rows by {steps} (o_proj split_k 1: the built-in cuBLAS call's bits, " + "its weight loaded before its dependency wait under the attention; QKV split_k 4: the " + "served cell's bits)" + ) return plan diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py index cbb16497aa0f..6ce374d41204 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py @@ -41,7 +41,7 @@ attention and append on a multi-request split geometry) instead of ``dflash_window_attention``. -Two levers refine them, each resolved by the plan and passed to ``install``: +Three levers refine them, each resolved by the plan and passed to ``install``: * ``drafter_conv_epilogue`` (on ``drafter_block_conv`` and ``drafter_fixed_gemms``): on blocks of 9..32 rows whose coefficient projection runs a ``drafter_fixed_gemms`` tactic, ``prepare`` is one @@ -52,6 +52,11 @@ * ``drafter_finish_norm`` (on ``drafter_block_decode``): on blocks the drafter hops do not take, each hop's finish convolution and the next norm's fused residual add + RMSNorm are one launch (``ops/spec/drafter_finish_norm.py``, ``FinishNorm``), bitwise the two. +* ``drafter_silu_epilogue`` (on ``drafter_block_decode``): on blocks of 9..32 + rows whose ``gate_up_proj`` runs a ``drafter_fixed_gemms`` tactic, the MLP's ``gate_up_proj`` and + ``silu(gate) * up`` are one launch (``ops/gemm/gate_up_silu_gemm.py``, ``SiluEpilogue``) over a + copy of the gate_up weight with gate and up rows interleaved (112 MiB per layer), bitwise the + fixed-tactic GEMM and the built-in ``silu_and_mul``. The fast paths compose. ``windowed_block_decode``, like the built-in forward, reaches the block convolutions through the modules' ``prepare`` / ``finish``, @@ -92,6 +97,9 @@ _ENTRY_ROWS = (8,) +#: Block rows ``SiluEpilogue`` serves: two to four requests' blocks (the one-request block keeps +#: the PDL chain's silu trigger). +SILU_EPILOGUE_ROWS = (9, 32) #: Block rows ``ConvEpilogue`` serves: two to four requests' blocks (the one-request block keeps #: the PDL chain's convolution trigger). CONV_EPILOGUE_ROWS = (9, 32) @@ -251,6 +259,153 @@ def _finish_norm_reason(draft_model: nn.Module) -> Optional[str]: return None +class SiluEpilogue: + """One drafter MLP with ``gate_up_proj`` and ``silu(gate) * up`` as one launch. + + Holds the ``interleave``d copy of the MLP's ``gate_up_proj`` weight. Called as the MLP is + (``x``, ``final_all_reduce_params``); a call it does not serve -- rows outside + ``SILU_EPILOGUE_ROWS``, operands the kernel does not take, a ``gate_up_proj`` not on a + ``drafter_fixed_gemms`` tactic at that row count (whose bits the kernel reproduces) or no longer + holding the copied weight -- runs the MLP itself. + """ + + logged = False + + def __init__(self, mlp: nn.Module) -> None: + from ...ops.gemm import gate_up_silu_gemm as _gemm + + weight = mlp.gate_up_proj.weight + self.mlp = mlp + self.weight = _gemm.interleave(weight.detach()) + self.source = weight + self._errors: dict = {} + + def _tactic(self, rows: int, k: int): + """The ``drafter_fixed_gemms`` tactic ``gate_up_proj`` runs at ``rows``, if the fused kernel + takes it, else None (cached per row count and tactic).""" + from ...spec_worker import FixedTacticLinearMethod + + proj = self.mlp.gate_up_proj + method = proj.quant_method + if type(method) is not FixedTacticLinearMethod or proj.weight is not self.source: + return None + if getattr(proj, "use_custom_cublas_mm", False): + return None + tactic = method.tactic_for(rows, k, self.weight.shape[0]) + if tactic is None: + return None + key = (rows, k, tuple(tactic)) + error = self._errors.get(key, False) + if error is False: + from ...ops.gemm import batch_splitk_gemm as _bsk + from ...ops.gemm import gate_up_silu_gemm as _gemm + + error = _gemm.tactic_error( + _bsk.BatchSplitKTactic(*tactic), rows, self.weight.shape[0], k + ) + self._errors[key] = error + return None if error is not None else (tactic, method.enable_pdl) + + def __call__(self, x: torch.Tensor, final_all_reduce_params=None) -> torch.Tensor: + rows = x.shape[0] if x.dim() == 2 else 0 + picked = None + if ( + SILU_EPILOGUE_ROWS[0] <= rows <= SILU_EPILOGUE_ROWS[1] + and x.dtype == torch.bfloat16 + and x.is_cuda + and x.is_contiguous() + and x.data_ptr() % 32 == 0 + and not torch.is_grad_enabled() + ): + picked = self._tactic(rows, x.shape[1]) + if picked is None: + return self.mlp(x, final_all_reduce_params=final_all_reduce_params) + from ...ops.gemm import batch_splitk_gemm as _bsk + from ...ops.gemm import gate_up_silu_gemm as _gemm + + tactic, enable_pdl = picked + activated = _gemm.run(x, self.weight, _bsk.BatchSplitKTactic(*tactic), enable_pdl) + if not SiluEpilogue.logged: + SiluEpilogue.logged = True + logger.info( + f"[musespark_open] drafter_silu_epilogue: the first fused gate_up ran ({rows} rows, " + f"split-K tactic {tuple(tactic)}): gate_up GEMM + silu(gate) * up in one launch" + ) + mlp = self.mlp + return mlp.down_proj( + activated, all_reduce_params=final_all_reduce_params, layer_idx=mlp.layer_idx + ) + + +def _silu_epilogue_reason(draft_model: nn.Module) -> Optional[str]: + """Why ``SiluEpilogue`` cannot stand in for this drafter's MLPs, or None. + + The MLP must run exactly ``down_proj(silu_and_mul(gate_up_proj(x)))``: the built-in + ``GatedMLP`` (or the PDL chain's re-class of it), a plain SwiGLU on a fused, bias-free BF16 + ``gate_up_proj`` and a BF16 ``down_proj`` input. + """ + from tensorrt_llm._torch.modules.gated_mlp import GatedMLP + + for index, layer in enumerate(draft_model.model.layers): + mlp = layer.mlp + where = f"layers.{index}.mlp" + if not isinstance(mlp, GatedMLP) or type(mlp).forward is not GatedMLP.forward: + return f"{where} is {type(mlp).__name__}, not the built-in GatedMLP" + if ( + mlp.activation is not torch.nn.functional.silu + or not mlp._is_plain_swiglu() + or mlp.swiglu_limit + ): + return f"{where} is not a plain SwiGLU" + if mlp.split_gate_up: + return f"{where} splits its gate/up projection" + if mlp._can_fuse_gate_up_swiglu() or mlp._can_fuse_swiglu_fp8_quant(): + return f"{where} fuses its activation into a GEMM" + down = mlp.down_proj + if getattr(down, "has_fp8_qdq", False) or getattr(down, "has_w4a8_nvfp4_fp8", False): + return f"{where} quantizes the activation for down_proj" + weight = getattr(mlp.gate_up_proj, "weight", None) + if ( + weight is None + or weight.dtype != torch.bfloat16 + or weight.dim() != 2 + or weight.shape[0] % 16 + or mlp.gate_up_proj.bias is not None + ): + return f"{where}.gate_up_proj is not a bias-free BF16 [2d, K] projection" + return None + + +def install_silu_epilogue(draft_model: nn.Module) -> int: + """Give every drafter MLP its ``SiluEpilogue``; return how many are new. + + Idempotent (an MLP whose epilogue holds the current gate_up weight is left alone). + + Raises: + ValueError: on an MLP the epilogue cannot stand in for (``_silu_epilogue_reason``), on + every rank alike. + """ + reason = _silu_epilogue_reason(draft_model) + if reason is not None: + raise ValueError(f"[musespark_open] drafter_silu_epilogue cannot run here: {reason}") + added, total_bytes = 0, 0 + for layer in draft_model.model.layers: + mlp = layer.mlp + fused = mlp.__dict__.get("_silu_epilogue") + if fused is None or fused.source is not mlp.gate_up_proj.weight: + mlp._silu_epilogue = SiluEpilogue(mlp) + added += 1 + total_bytes += mlp._silu_epilogue.weight.numel() * 2 + if added: + logger.info( + f"[musespark_open] drafter_silu_epilogue: drafter blocks of {SILU_EPILOGUE_ROWS[0]}.." + f"{SILU_EPILOGUE_ROWS[1]} rows run each MLP's gate_up_proj + silu(gate) * up as one launch " + f"on {len(draft_model.model.layers)} layers ({total_bytes / 2**20:.0f} MiB of interleaved " + "gate_up weights)" + ) + return added + + class FusedDFlash2BlockConv(DFlash2BlockConv): """The built-in DFlash 2 block convolution with ``_convolve`` as one kernel. @@ -628,6 +783,8 @@ def windowed_block_decode( hidden_out, mlp_conv_coefficients = self.mlp_convs[layer_idx].prepare( hidden_out, block_size ) + # ``drafter_silu_epilogue``: gate_up_proj + silu(gate) * up in one launch. + mlp = layer.mlp.__dict__.get("_silu_epilogue") or layer.mlp if hops is not None: # The MLP, its all-reduce, the MLP finish convolution and the next # norm (the next layer's input_layernorm, or model.norm) in one hop. @@ -638,7 +795,7 @@ def windowed_block_decode( ) prenormed = hops.mlp( layer_idx, - layer.mlp, + mlp, hidden_out, res_flat, mlp_conv_coefficients, @@ -647,7 +804,7 @@ def windowed_block_decode( block_size, ) continue - hidden_out = layer.mlp(hidden_out) + hidden_out = mlp(hidden_out) if finish_norm is not None: # The MLP finish convolution and the next norm (the next layer's input_layernorm, # or model.norm) in one launch, as the hop does; the residual is updated in place. @@ -708,6 +865,7 @@ def install( *, conv_epilogue: bool = False, finish_norm: bool = False, + silu_epilogue: bool = False, ) -> None: """Install the fused block convolution and the windowed block decode on ``draft_model``. @@ -726,7 +884,7 @@ def install( if block_conv: _install_fused_block_conv(draft_model, conv_epilogue) if block_decode: - _install_windowed_attention(draft_model, True, finish_norm) + _install_windowed_attention(draft_model, True, finish_norm, silu_epilogue) def _builtin_block_convs(draft_model: nn.Module) -> int: @@ -794,7 +952,7 @@ def install_conv_epilogue(draft_model: nn.Module) -> int: def _install_windowed_attention( - draft_model: nn.Module | None, block_decode: bool, finish_norm: bool + draft_model: nn.Module | None, block_decode: bool, finish_norm: bool, silu_epilogue: bool ) -> None: enable_pdl = LAUNCH_PDL wrapped = isinstance( @@ -829,3 +987,5 @@ def _install_windowed_attention( "take run each hop's finish convolution + residual add + RMSNorm as one launch " f"(pdl={enable_pdl}, dependents released at entry)" ) + if silu_epilogue: + install_silu_epilogue(draft_model) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py index c3f810824ee5..aff50ab089f8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py @@ -88,8 +88,11 @@ def apply(plan: "_dispatch.Plan") -> None: # Dense GEMMs. ( "kernels.dense_gemms.dispatcher", - dict(batch_splitk=on("tep8_batch_splitk")), - dict(batch_splitk=False), + dict( + batch_splitk=on("tep8_batch_splitk"), + attn_block_splitk=on("tep8_attn_block_splitk"), + ), + dict(batch_splitk=False, attn_block_splitk=False), ), # Attention. ( @@ -152,6 +155,11 @@ def apply(plan: "_dispatch.Plan") -> None: dict(tep8_max_rows=None), ), # Drafter. + ( + "ops.gemm.fixed_low_m_gemm", + dict(multi_request_cells=on("drafter_multi_request_cells")), + dict(multi_request_cells=False), + ), ( "ops.allreduce.drafter_hop", dict( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py index 702ce3e2456f..c266f76160da 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py @@ -138,6 +138,9 @@ ), "TRTLLM_MUSESPARK_DRAFTER_CONV_EPILOGUE": PREFIX + "DRAFTER_CONV_EPILOGUE=on", "TRTLLM_MUSESPARK_DRAFTER_FINISH_NORM": PREFIX + "DRAFTER_FINISH_NORM=on", + "TRTLLM_MUSESPARK_DRAFTER_SILU_EPILOGUE": PREFIX + "DRAFTER_SILU_EPILOGUE=on", + "TRTLLM_MUSESPARK_TEP8_ATTN_BLOCK_SPLITK": PREFIX + "TEP8_ATTN_BLOCK_SPLITK=on", + "TRTLLM_MUSESPARK_DRAFTER_MULTI_REQUEST_CELLS": (PREFIX + "DRAFTER_MULTI_REQUEST_CELLS=on"), } ) #: Name prefixes only retired switches carry; any other name under them raises too. @@ -709,6 +712,12 @@ def _vocab_parallel_head(d: Deployment, t: Resolving, p: Mapping[str, object]) - # New this round: default off, not covered by the TP8 default-on gate. "drafter_conv_epilogue": _ported(Kind.BITWISE), "drafter_finish_norm": _ported(Kind.BITWISE), + "drafter_silu_epilogue": _ported(Kind.BITWISE), + "tep8_attn_block_splitk": _ported(Kind.BITWISE), + # The qkv / gate_up / coefficient cells keep their row's split_k and are bitwise the + # table's row; the o_proj and down_proj cells reproduced cuBLAS's bits on every checked + # shape, which cannot be proven against a closed library. + "drafter_multi_request_cells": _ported(Kind.FLAGGED), **{name: _tep4_path_rule() for name in TEP4_PATHS}, } ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py index 106f88be6f9a..a80231056722 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py @@ -19,6 +19,8 @@ from __future__ import annotations +import functools + import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( @@ -82,7 +84,85 @@ (32, 8192, 14336): (128, 32, 1, 5), } + #: The measured table of each ``(compute capability, drafter TP size)``. +class OnBatchSplitK(tuple): + """A tactic ``(mma_m, mma_n, split_k, ab_stages)`` run on ``batch_splitk_gemm``'s copy of the + split-K kernel instead of flashinfer's. + + The copy has the same tiles and the same summation (each output element the FP32 sum of the + ``split_k`` K-slice partials in cluster-rank order, rounded once), so its bits are flashinfer's + at the same ``split_k``; its epilogue reduce-scatters the partials over the cluster instead of + sending every peer's whole tile to one owner, which leaves room for a deeper ring. A plain + 4-tuple everywhere else (``len``, unpacking, equality). + """ + + __slots__ = () + + def __new__(cls, *fields: int) -> "OnBatchSplitK": + return super().__new__(cls, fields) + + def __repr__(self) -> str: + return f"bsplitk{tuple(self)!r}" + + +#: ``drafter_multi_request_cells`` (sm103 TP8 only): the blocks of 2..4 requests (16, 24 and 32 +#: rows) run these cells in place of the TP8 table's rows chosen from the shapes. Each qkv / +#: gate_up / convolution coefficient cell keeps its row's ``split_k`` (bitwise the table's row). +#: The o_proj and down_proj cells, which the low-M dispatcher's per-rank race served (cuBLAS on +#: the measured ranks), are single K-slice GEMMs (``split_k`` 1). With every successor listed at a +#: width, ``drafter_pdl_chain``'s entry triggers act on these widths too. +MULTI_REQUEST_CELLS_TP8: dict[tuple[int, int, int], tuple[int, int, int, int]] = { + # qkv_proj / context K/V: 16 rows one stage deeper; 24 / 32 rows on the reduce-scatter copy, + # one more stage than flashinfer fits. + (16, 8192, 1280): (64, 16, 4, 10), + (24, 8192, 1280): OnBatchSplitK(64, 32, 4, 9), + (32, 8192, 1280): OnBatchSplitK(64, 32, 4, 9), + # Convolution coefficient projection: 16 rows one stage deeper (``ConvEpilogue`` fits its + # stages beside its staging tile). The 24 / 32-row rows stay: ``mma_n`` 32 there was a large + # loss with the convolution epilogue. + (16, 8192, 2048): (64, 16, 4, 10), + # gate_up_proj: the same tiles on the reduce-scatter copy. + (16, 8192, 7168): OnBatchSplitK(128, 16, 2, 6), + (24, 8192, 7168): OnBatchSplitK(128, 32, 2, 5), + (32, 8192, 7168): OnBatchSplitK(128, 32, 2, 5), + # o_proj per rank: one K slice, the fastest tile per row count in the chain. + (16, 1024, 8192): OnBatchSplitK(128, 16, 1, 6), + (24, 1024, 8192): OnBatchSplitK(128, 32, 1, 5), + (32, 1024, 8192): (64, 32, 1, 8), + # down_proj per rank: one K slice, the fastest tile per row count in the chain. + (16, 3584, 8192): (64, 16, 1, 11), + (24, 3584, 8192): (64, 32, 1, 9), + (32, 3584, 8192): OnBatchSplitK(64, 32, 1, 9), +} + +#: Whether ``splitk_tactics`` merges ``MULTI_REQUEST_CELLS_TP8`` into the sm103 TP8 table: +#: ``drafter_multi_request_cells``, pushed here by ``lever_settings.apply`` (default off). +_MULTI_REQUEST_CELLS = False + + +def configure(*, multi_request_cells: bool) -> None: + """Set this module's lever-backed settings (``lever_settings.apply``), once per process.""" + global _MULTI_REQUEST_CELLS + _MULTI_REQUEST_CELLS = bool(multi_request_cells) + + +@functools.lru_cache(maxsize=None) +def _with_multi_request_cells(key: tuple) -> dict: + """The table of ``key`` with ``MULTI_REQUEST_CELLS_TP8`` merged in; built (and logged) once per + process.""" + from tensorrt_llm.logger import logger + + logger.info( + "[musespark_open] drafter_multi_request_cells: drafter blocks of 16, 24 and 32 rows run " + f"{len(MULTI_REQUEST_CELLS_TP8)} measured TP8 cells (qkv / gate_up / coefficient rows at " + "their split_k, o_proj / down_proj on one K slice; " + f"{sum(isinstance(v, OnBatchSplitK) for v in MULTI_REQUEST_CELLS_TP8.values())} on the " + "reduce-scatter split-K copy)" + ) + return {**SPLITK_TACTICS_BY_SM_TP[key], **MULTI_REQUEST_CELLS_TP8} + + SPLITK_TACTICS_BY_SM_TP: dict[ tuple[tuple[int, int], int], dict[tuple[int, int, int], tuple[int, int, int, int]] ] = { @@ -128,6 +208,8 @@ def splitk_tactics( f"fixed_low_m_gemm: no split-K tactics for SM {key[0]} at TP {key[1]} " f"(supported: {sorted(SPLITK_TACTICS_BY_SM_TP)})" ) + if key == ((10, 3), 8) and _MULTI_REQUEST_CELLS: + return _with_multi_request_cells(key) return table @@ -197,6 +279,12 @@ def __call__( from .splitk_gemm import SplitKTactic, run_splitk_dense out = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device) + if isinstance(tactic, OnBatchSplitK): + from . import batch_splitk_gemm as _bsk + + return _bsk.run( + x.detach(), weight.detach(), out, _bsk.BatchSplitKTactic(*tactic), enable_pdl + ) # DLPack refuses tensors that require grad; a trainable nn.Linear # weight still does under inference mode. return run_splitk_dense( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/gate_up_silu_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/gate_up_silu_gemm.py new file mode 100644 index 000000000000..e7a4c20bdd82 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/gate_up_silu_gemm.py @@ -0,0 +1,569 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The drafter MLP's ``gate_up_proj`` and ``silu(gate) * up`` in one launch. + +``GatedMLP`` runs ``gate_up_proj`` (on the drafter's 9..32-row blocks flashinfer's cluster split-K +kernel at the fixed tactic of ``fixed_low_m_gemm``, ``drafter_fixed_gemms``) and then the built-in +Triton ``silu_and_mul_kernel`` (``modules/swiglu.py``) as its own launch before ``down_proj``. + +This entry is that GEMM on ``batch_splitk_gemm``'s copy of the same kernel (same tiles, same +``split_k``, so the same bits: each output element is the FP32 sum of the ``split_k`` K-slice +partials in cluster-rank order, rounded once to BF16) over a row-permuted copy of the weight in +which gate row ``c`` and up row ``c`` are adjacent (``interleave``: rows ``2c``, ``2c + 1``). A +row permutation changes no output element's K-ordered sum. After the reduction the CTA that owns +a share of the tile rounds it to BF16 (the GEMM output's values), stages it widened back to FP32 +(exact) in shared memory, and applies the Triton kernel's statements to each (gate, up) pair, as +the PTX that kernel compiles to (``SILU_PTX``: ``a * -log2(e)``, ``ex2.approx.f32``, ``+ 1``, +``div.full.f32``, ``* a``, ``b + 0.0``, the product), rounding the result to BF16 to nearest even +as its store does. So the output is bitwise ``silu_and_mul(fixed_low_m_gemm(x, weight))``. The +GEMM output itself is never stored. + +Ownership: the reduce-scatter epilogue gives cluster rank ``r`` the values of epilogue thread group +``r`` (``128 / split_k`` threads, whole warps); a warp's TMEM sub-partition holds a contiguous +quarter of the tile's kernel-M rows, so rank ``r`` owns kernel-M rows ``[r M / s, (r + 1) M / s)`` +of every public row: whole (gate, up) pairs, whatever the tactic. One interleaved copy serves +every tactic. + +Programmatic dependent launch: as ``batch_splitk_gemm`` (weight before the wait, activation after +it, dependents released once the weight warp has issued its last load). The epilogue reads no +global memory; it stores after the reduction, which follows the dependency wait. + +The SwiGLU statements are the ones upstream's sm107 drafter folds compile into their split-K +GEMM's epilogue (TensorRT-LLM ``feat/musespark-open-agentx`` #21754, +``muse_spark_open_816b_a42b_nvfp4__sm_107__tep/ops/gemm/dr3_fold_splitk_gemm.py``, ``SILU_PTX``; +flashinfer's split-K kernel at ``mma_n`` 8 / 16 with a whole-tile owner). This copy runs on the +reduce-scatter kernel at any ``mma_n`` the tactic tables use (up to 32 at 32 rows). +""" + +from __future__ import annotations + +import functools +from typing import Tuple + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass import Int32, Uint32 +from cutlass.cute import experimental as cute_ext +from cutlass.cutlass_dsl import dsl_user_op + +from ..allreduce.drafter_hop import _asm, _f32_asm, _i64, _pack_bf16x2, _u32 +from . import batch_splitk_gemm as _bsk + +#: Output columns per epilogue item (one 16-byte BF16 vector): 16 kernel-M rows (8 pairs). +_VECTOR = 8 +_STAGE_ALIGN_BYTES = 128 + +#: The Triton ``silu_and_mul_kernel`` statements for one element (``swiglu_alpha`` 1.0, +#: ``swiglu_beta`` 0.0), as that kernel compiles: ``1.0 * a`` and the negation fold into one +#: multiply by ``-log2(e)``, ``exp`` is ``ex2.approx.f32``, ``1 / (1 + e)`` is ``div.full.f32``, +#: then ``a * sigmoid``, ``b + 0.0`` (keeps the sign of zero) and their product. No two of them +#: can contract (no product feeds an add). From upstream's ``SILU_PTX`` (#21754). +SILU_PTX = ( + "{\n" + ".reg .f32 t, e, d, s, p, q;\n" + "mul.f32 t, $1, 0fBFB8AA3B;\n" + "ex2.approx.f32 e, t;\n" + "add.f32 d, e, 0f3F800000;\n" + "div.full.f32 s, 0f3F800000, d;\n" + "mul.f32 p, $1, s;\n" + "add.f32 q, $2, 0f00000000;\n" + "mul.f32 $0, p, q;\n" + "}" +) + + +@dsl_user_op +def _silu_and_mul(a, b, *, loc=None, ip=None): + """``a * sigmoid(a) * (b + 0.0)`` in FP32, the Triton kernel's instructions.""" + return _f32_asm(SILU_PTX, (a, b), loc, ip) + + +@dsl_user_op +def _store_x4_if(address, enabled, w0, w1, w2, w3, *, loc=None, ip=None): + """Plain 16-byte global store when ``enabled``; else nothing.""" + _asm( + None, + [_i64(address, loc, ip), _u32(enabled, loc, ip)] + + [_u32(w, loc, ip) for w in (w0, w1, w2, w3)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $1, 0;\n@p st.global.v4.b32 [$0], {$2, $3, $4, $5};\n}", + "l,r,r,r,r,r", + loc, + ip, + ) + + +def interleave(weight: torch.Tensor) -> torch.Tensor: + """The kernel's weight: ``[gate; up]`` rows ``[2d, K]`` -> ``[d, 2, K]`` (gate ``c``, up ``c`` + adjacent). Returns a new contiguous tensor (the module's weight is not modified).""" + d = weight.shape[0] // 2 + return torch.stack((weight[:d], weight[d:]), dim=1).reshape(weight.shape).contiguous() + + +def stage_bytes(tactic: _bsk.BatchSplitKTactic) -> int: + return tactic.mma_m * tactic.mma_n * 4 + + +def smem_bytes(tactic: _bsk.BatchSplitKTactic, ab_stages: int | None = None) -> int: + """``batch_splitk_gemm``'s layout plus the FP32 staging tile.""" + cursor = _bsk.smem_bytes(tactic, ab_stages) + return _bsk._align_up(cursor, _STAGE_ALIGN_BYTES) + stage_bytes(tactic) + + +def fitted(tactic: _bsk.BatchSplitKTactic) -> _bsk.BatchSplitKTactic: + """``tactic`` with as many of its stages as fit beside the staging tile (pipeline depth only: + the stages do not change any output bit).""" + stages = tactic.ab_stages + while stages > 1 and smem_bytes(tactic, stages) > _bsk.SMEM_CAPACITY_BYTES: + stages -= 1 + return _bsk.BatchSplitKTactic(tactic.mma_m, tactic.mma_n, tactic.split_k, stages) + + +def tactic_error(tactic: _bsk.BatchSplitKTactic, m: int, n: int, k: int) -> str | None: + """Why the fused kernel cannot run ``x[m, k] @ weight[n, k]^T`` at ``tactic`` (stages as + ``fitted``), or None. ``n`` is the GEMM's width (gate and up).""" + tactic = fitted(tactic) + try: + _bsk.validate(tactic, m, n, k) + except ValueError as error: + return str(error) + if tactic.split_k > 4: + return f"split_k={tactic.split_k}: an owner group is under one warp" + rows_owned = tactic.mma_m // tactic.split_k + if rows_owned % (2 * _VECTOR): + return ( + f"tactic {tactic}: a rank's {rows_owned} kernel-M rows are not whole 8-column vectors" + ) + if n % tactic.mma_m: + return f"mma_m={tactic.mma_m} does not tile the {n} rows" + if smem_bytes(tactic) > _bsk.SMEM_CAPACITY_BYTES: + return f"tactic {tactic} with its staging tile exceeds shared memory" + return None + + +def items(tactic: _bsk.BatchSplitKTactic) -> Tuple[int, int, int]: + """``(vectors per public row, items per tile, items per owner thread)`` of an owner group.""" + rows_owned = tactic.mma_m // tactic.split_k + per_row = rows_owned // (2 * _VECTOR) + total = per_row * tactic.mma_n + group = 128 // tactic.split_k + return per_row, total, (total + group - 1) // group + + +class GateUpSiluGemmKernel(_bsk.BatchSplitKGemmKernel): + """``batch_splitk_gemm``'s kernel whose owners store ``silu(gate) * up`` of their share.""" + + def __init__(self, *, tactic: _bsk.BatchSplitKTactic, use_pdl: bool) -> None: + super().__init__(tactic=tactic, use_pdl=use_pdl) + self.rows_owned = tactic.mma_m // tactic.split_k + self.per_row, self.total_items, self.items = items(tactic) + + @cute.experimental.jit + def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: _cuda.CUstream): + # The parent's grid over the GEMM's kernel-M (the interleaved weight's rows: twice the + # output's), its cluster, shared-memory request and PDL flag. + self.kernel(a, b, c).launch( + grid=( + cute.ceil_div(a.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(_bsk.SMEM_CAPACITY_BYTES), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major: the interleaved weight + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major: the activation + mC: cute.Tensor, # (Gemm_M / 2, Gemm_N, Gemm_L), M-major: the SwiGLU output + ): + # batch_splitk_gemm.BatchSplitKGemmKernel.kernel, the staging tile and the epilogue aside. + stages = self.num_ab_stage + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_bsk._AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_bsk._AB_BUFFER_ALIGN_BYTES, + ) + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + bar_spare = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_bsk._TMEM_POINTER_BYTES, + ).iterator + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.rs_mailbox_elements), + alignment=_bsk._MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_bsk._MBARRIER_BYTES, + ).iterator + else: + mailbox = sA + bar_reduce = bar_spare + # The owner's BF16-rounded GEMM values of the tile, as FP32, kernel-M fastest. + stage = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n), stride=(1, self.cta_m)), + alignment=_STAGE_ALIGN_BYTES, + ) + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + if cutlass.const_expr(self.split_k > 1): + cute.arch.mbarrier_init(bar_reduce, 1) + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_silu_warp( + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.arch.thread_idx()[0] - 128, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + stage, + bidx, + n_idx, + l_idx, + mC, + ) + + @cute.experimental.jit + def epilog_silu_warp( + self, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + epi_tid: cutlass.Int32, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + stage: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + l_idx: cutlass.Int32, + mC: cute.Tensor, + ): + # batch_splitk_gemm.BatchSplitKGemmKernel.epilog_warp up to the owner's rounding. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + _bsk._tcgen05_fence(before=False) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), acc_layout + )[((None, None), 0, 0, 0)] + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + cutlass.BFloat16, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + stage_epi = cute.flat_divide(stage, epi_tile) + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, stage_epi, epi_tid) + rAcc = cute_ext.allocate(self.acc_dtype, cute.AddressSpace.rmem, rmem_layout, alignment=32) + rD = cute_ext.allocate(cutlass.BFloat16, cute.AddressSpace.rmem, rmem_layout, alignment=32) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + dest = Int32(0) + lt = epi_tid + if cutlass.const_expr(self.split_k > 1): + dest = epi_tid // self.rs_group + lt = epi_tid % self.rs_group + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + _bsk._tcgen05_fence(before=False) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + cute.arch.fence_view_async_tmem_load() + _bsk._tcgen05_fence(before=True) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + if cutlass.const_expr(self.split_k > 1): + vpt = cutlass.const_expr(cute.size(rmem_layout)) + group_values = cutlass.const_expr(self.rs_group * vpt) + if dest != split_rank: + j = split_rank - cutlass.select_(split_rank > dest, Int32(1), Int32(0)) + for value_idx in cutlass.range_constexpr(0, vpt, 4): + slot = (value_idx // 4) * (self.rs_group * 4) + lt * 4 + _bsk._store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + j * group_values + slot, + bar_reduce, + dest, + ) + else: + if lt == 0: + cute.arch.mbarrier_arrive_and_expect_tx(bar_reduce, self.rs_expected_bytes) + _bsk._mbarrier_wait_acquire_cluster(bar_reduce, Int32(0)) + for value_idx in cutlass.range_constexpr(vpt): + slot = (value_idx // 4) * (self.rs_group * 4) + lt * 4 + value_idx % 4 + own = rAcc[value_idx] + acc = own + for r in cutlass.range_constexpr(self.split_k): + jr = cutlass.min( + Int32(r) - cutlass.select_(Int32(r) > split_rank, Int32(1), Int32(0)), + Int32(self.split_k - 2), + ) + term = cutlass.select_( + Int32(r) == split_rank, own, mailbox[jr * group_values + slot] + ) + if cutlass.const_expr(r == 0): + acc = term + else: + acc = acc + term + rAcc[value_idx] = acc + # The GEMM output's values (BF16), widened back exactly. + rD.store(rAcc.load().to(cutlass.BFloat16)) + rAcc.store(rD.load().to(self.acc_dtype)) + cute_ext.partition_and_copy(thr_t2r, rAcc, stage_epi[None, None, 0, 0]) + else: + rD.store(rAcc.load().to(cutlass.BFloat16)) + rAcc.store(rD.load().to(self.acc_dtype)) + cute_ext.partition_and_copy(thr_t2r, rAcc, stage_epi[None, None, 0, 0]) + + if dest == split_rank: + # The owner group's staged rows are complete for every one of its threads. + cute.arch.barrier(barrier_id=1, number_of_threads=self.rs_group) + self._silu(stage, lt, split_rank, bidx, n_idx, l_idx, mC) + + def _silu(self, stage, lt, split_rank, bidx, n_idx, l_idx, mC): + """``silu(gate) * up`` of this thread's items: 8 output columns of one public row each, + from 16 staged kernel-M rows (gate, up pairs); rows past the call's are not stored. Traced + from the jit epilogue: its Python loops unroll.""" + rows = cute.size(mC, mode=[1]) + m_first = split_rank * self.rows_owned + for i in range(self.items): + item = lt + i * self.rs_group + vec = item % self.per_row + n = item // self.per_row + t = n_idx * self.cta_n + n + m0 = m_first + vec * (2 * _VECTOR) + column = (bidx * self.cta_m + m0) // 2 + valid = Uint32(t < rows) & Uint32(item < self.total_items) + # Clamp the stage coordinates of an item past the tile (never stored). + n_safe = cutlass.min(n, Int32(self.cta_n - 1)) + words = [] + for k in range(_VECTOR // 2): + lo = _silu_and_mul(stage[m0 + 4 * k, n_safe], stage[m0 + 4 * k + 1, n_safe]) + hi = _silu_and_mul(stage[m0 + 4 * k + 2, n_safe], stage[m0 + 4 * k + 3, n_safe]) + words.append(_pack_bf16x2(lo, hi)) + address = ( + mC.iterator + cute.crd2idx((column, cutlass.min(t, rows - 1), l_idx), mC.layout) + ).toint() + _store_x4_if(address, valid, words[0], words[1], words[2], words[3]) + + +@cute.experimental.jit +def _bmm_silu( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + stream, + ) + + +@functools.lru_cache(maxsize=None) +def _compiled( + tactic: _bsk.BatchSplitKTactic, use_pdl: bool, leading: Tuple[int, int, int], device: int +): + m, n, k = 64, 8, _bsk.CTA_K + dtype = torch.bfloat16 + with torch.cuda.device(device): + reprs = tuple( + _bsk._from_dlpack_dynamic(_bsk._make_layout_tensor(shape, dtype, d), d) + for shape, d in ( + ((1, n, k), leading[0]), + ((1, k, m), leading[1]), + ((1, n, m), leading[2]), + ) + ) + return cute_ext.compile( + _bmm_silu, + GateUpSiluGemmKernel(tactic=tactic, use_pdl=use_pdl), + *reprs, + _cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +def run( + x: torch.Tensor, + weight: torch.Tensor, + tactic: _bsk.BatchSplitKTactic, + pdl: bool = True, +) -> torch.Tensor: + """``silu(x @ gate.T) * (x @ up.T)`` -> ``[M, d]`` BF16 for an ``interleave``d weight ``[2d, K]``. + + Bitwise ``silu_and_mul(fixed_low_m_gemm(x, [gate; up], tactic))``. ``tactic`` is taken with its + stages cut to fit the staging tile (``fitted``). The first call per tactic and layout compiles + (CuTe-DSL JIT) and must not be inside a CUDA-graph capture. + + Raises: + ValueError: for operands or a tactic the kernel does not take (``tactic_error``). + """ + tactic = fitted(tactic) + if not (x.dtype == weight.dtype == torch.bfloat16): + raise ValueError("the SwiGLU epilogue takes BF16 operands only") + if x.dim() != 2 or weight.dim() != 2 or int(weight.shape[1]) != int(x.shape[1]): + raise ValueError(f"shapes x {tuple(x.shape)}, weight {tuple(weight.shape)}") + if not (x.is_contiguous() and weight.is_contiguous()): + raise ValueError("the SwiGLU epilogue takes contiguous operands") + if any(t.data_ptr() % 32 for t in (x, weight)): + raise ValueError("the SwiGLU epilogue takes 32-byte aligned operands") + if not (x.is_cuda and weight.device == x.device): + raise ValueError("the SwiGLU epilogue takes CUDA operands on one device") + m, k = (int(s) for s in x.shape) + n = int(weight.shape[0]) + error = tactic_error(tactic, m, n, k) + if error is not None: + raise ValueError(f"the SwiGLU epilogue cannot run {tactic}: {error}") + out = torch.empty((m, n // 2), dtype=x.dtype, device=x.device) + with torch.cuda.device(x.device): + tensors, leading = _bsk._swapped(x, weight[: n // 2], out) + # ``_swapped`` gives A from the half-height view; the kernel reads the whole weight. + tensors = (_bsk._from_dlpack_dynamic(weight.unsqueeze(0), leading[0]),) + tensors[1:] + compiled = _compiled(tactic, bool(pdl), leading, x.get_device()) + compiled(*tensors, _cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream)) + return out From cd86beae2ef4e510b503b1ee484ca6737be650f8 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Mon, 9 Nov 2026 21:41:11 -0800 Subject: [PATCH 6/9] [None][feat] MuseSpark GB300: let the wide expert placement reach 512 decode rows Raises the ceiling of tep8_wide_balance's MAX_ROWS from 256 to 512 rows. Above 256 rows (32 requests at seven draft tokens) the balanced NVFP4 placement had no band left, so a step past that count fell back where the row-count census, the ring-buddy cap and the placement band all ended at once. One wide ring_buddy_assign program now covers a 4096-slot tile. Its work is per expert -- the census is a histogram of the selected ids, the move decision is one per expert and the gather is one per slot -- so the larger tile costs no [SLOTS, RANKS] tensor. The narrow program keeps its 2048-slot bound, and a call takes the wider one only where wide_serves already admits it: the balanced TP8 placement above the wide program's minimum, with no half window and no static move. The lever's default stays 256, so nothing reaches the new band until the parameter is set, and the lever itself is still default off. The AST verifier resolves the same path sets as the base at all twelve deployments with an EMPTY environment; MAX_ROWS now accepts 288 and 512 and still refuses 128 and 513. Run on GB300 (2 nodes, TEP8 holder): ring_buddy_assign's widened validation accepts 288 and 256 rows on the wide program and still takes 32 on the narrow one, exercised on the unconditional path. The model itself was not booted. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 90a880ae9fe98b908cfd95e8f69cda7ebec4bf52) --- .../levers.py | 15 ++++++++++-- .../ops/moe/ring_buddy_assign.py | 24 ++++++++++++++----- 2 files changed, 31 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py index c266f76160da..97fca5d58c5e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py @@ -570,9 +570,13 @@ def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, st RING_BUDDY_MAX_ROWS = 128 RING_BUDDY_ROW_CEILING = 256 #: ``tep8_wide_balance__MAX_ROWS``: the band the balanced NVFP4 placement adds above the -#: stacked cap, from ``ring_buddy_assign.WIDE_BALANCE_MIN_ROWS`` (129) to the same ceiling. +#: stacked cap, from ``ring_buddy_assign.WIDE_BALANCE_MIN_ROWS`` (129) to one wide +#: ``ring_buddy_assign`` program's rows (``WIDE_MAX_SLOTS // TOP_K``). The default stays at the +#: narrow program's ceiling; the band above it is reached only by setting the parameter, and +#: the wide program places every call of more than 64 rows. WIDE_BALANCE_ROWS = 256 WIDE_BALANCE_MIN_ROWS = 129 +WIDE_BALANCE_ROW_CEILING = 512 #: ``ring_buddy_stats__INTERVAL_S``: seconds between the placement-census reports TP rank 0 #: logs. Whole seconds only (the knob grammar has no fractional parameter). RING_BUDDY_STATS_INTERVAL_S = 60 @@ -694,7 +698,14 @@ def _vocab_parallel_head(d: Deployment, t: Resolving, p: Mapping[str, object]) - "tep8_wide_balance": _ported_tep8( Kind.FLAGGED, "the 256-row wide balance", - (Param("MAX_ROWS", WIDE_BALANCE_ROWS, WIDE_BALANCE_MIN_ROWS, RING_BUDDY_ROW_CEILING),), + ( + Param( + "MAX_ROWS", + WIDE_BALANCE_ROWS, + WIDE_BALANCE_MIN_ROWS, + WIDE_BALANCE_ROW_CEILING, + ), + ), ), # Diagnostics: one extra assign launch per placement, and the in-kernel placement # stands aside for it (``ring_buddy``'s ``_stats is None``), so it is never on by diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/ring_buddy_assign.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/ring_buddy_assign.py index 98bce3fb484e..6ea5e823b24d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/ring_buddy_assign.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/ring_buddy_assign.py @@ -95,9 +95,14 @@ def block_of(ranks: int) -> int: #: ``tep8_wide_balance``'s ``MAX_ROWS`` (default off): on tep8 the balanced placement of the #: NVFP4 layers also serves the decode calls above the stacked path's row cap, up to those rows -#: (129..256, ``8 * B`` for B requests at K = 7), each call of more than 64 rows placed by the +#: (129..512, ``8 * B`` for B requests at K = 7), each call of more than 64 rows placed by the #: wide program. WIDE_BALANCE_MIN_ROWS = 129 +#: Largest ``M * TOP_K`` the wide program handles: 512 decode rows (64 requests at K = 7), a +#: 4096-slot tile from 257 rows on. Its work is per expert (the census is a histogram of the +#: ids, the move decision one per expert, one gather per slot), so the tile above ``MAX_SLOTS`` +#: costs no ``[SLOTS, RANKS]`` tensor; the narrow program keeps ``MAX_SLOTS``. +WIDE_MAX_SLOTS = 4096 _WIDE_BALANCE_ROWS: Optional[int] = None @@ -109,7 +114,7 @@ def configure(*, wide_balance_rows: Optional[int] = None) -> None: Args: wide_balance_rows: ``tep8_wide_balance``'s ``MAX_ROWS`` (``_WIDE_BALANCE_ROWS``), - in ``WIDE_BALANCE_MIN_ROWS``..``MAX_SLOTS // TOP_K``. + in ``WIDE_BALANCE_MIN_ROWS``..``WIDE_MAX_SLOTS // TOP_K``. Raises: ValueError: for a row cap outside that range. @@ -119,11 +124,11 @@ def configure(*, wide_balance_rows: Optional[int] = None) -> None: _WIDE_BALANCE_ROWS = None return rows = int(wide_balance_rows) - if not WIDE_BALANCE_MIN_ROWS <= rows <= MAX_SLOTS // TOP_K: + if not WIDE_BALANCE_MIN_ROWS <= rows <= WIDE_MAX_SLOTS // TOP_K: raise ValueError( f"[musespark_open] tep8_wide_balance MAX_ROWS {rows}: expected the balanced " f"placement's row cap for the NVFP4 layers' decode calls, a whole number in " - f"{WIDE_BALANCE_MIN_ROWS}..{MAX_SLOTS // TOP_K}" + f"{WIDE_BALANCE_MIN_ROWS}..{WIDE_MAX_SLOTS // TOP_K}" ) _WIDE_BALANCE_ROWS = rows @@ -745,8 +750,15 @@ def is_valid( assert selected.ndim == 2 and selected.shape[1] == TOP_K, ( f"selected must be [M, {TOP_K}], got {tuple(selected.shape)}" ) - assert 0 < selected.numel() <= MAX_SLOTS, ( - f"one program covers at most {MAX_SLOTS} slots (M <= {MAX_SLOTS // TOP_K}), got {selected.numel()}" + slots = max(16, triton.next_power_of_2(max(1, selected.numel()))) + most = ( + WIDE_MAX_SLOTS + if wide and wide_serves(slots, ranks, balance, static_from, half) + else MAX_SLOTS + ) + assert 0 < selected.numel() <= most, ( + f"one program covers at most {most} slots (M <= {most // TOP_K}), " + f"got {selected.numel()}" ) assert weights.dtype == torch.float32 and weights.is_contiguous(), ( "weights must be contiguous float32" From 434826bfeb771c9c1a20167a861db6fb51868f28 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Mon, 9 Nov 2026 21:41:22 -0800 Subject: [PATCH 7/9] [None][feat] MuseSpark GB300: arm the one-request SWA block's early loads The sliding-window block's split and combine kernels already take an EARLY_LOADS build, but nothing asked for it, so the capability was inert. swa_block_single_early_loads (on swa_block, default off) arms it on the one-request block: * the split issues every projection load the program makes after its wait at once -- the query rows, the key block's step K/V rows and, in program (0, h, 0), the step's K/V rows it writes to the cache -- before any prep arithmetic, so the step costs one dependent round trip instead of one per prep. It needs the fused prep, whose operands those loads are, and is passed only when both hold; * the combine loads the partial O with the (max, sum) words, and the gate head before its wait where there is a gate. The same statements on the same values, so the same bits. The blocks-per-program split runs a different kernel that has no such build, so it keeps the served launch. Knob-only and default off; the environment switch this work carried is retired by name in levers.RETIRED. The AST verifier resolves path sets identical to the base at all twelve deployments with an EMPTY environment, and the lever turns on by its own knob at TEP8. Run on GB300 (2 nodes, TEP8 holder): the EARLY_LOADS build this commit arms -- which nothing had compiled before, since the halves were ported without a caller -- compiles, runs, and is bitwise the served build (0 differing elements). The model itself was not booted. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 2506fb61c71ed363e3ca4f14e8b1941c8eba2cbe) --- .../dispatch.py | 7 ++++ .../lever_settings.py | 5 +++ .../levers.py | 4 +++ .../ops/attention/swa_block.py | 36 +++++++++++++++++++ 4 files changed, 52 insertions(+) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py index 23ca2672e63b..705e931fcb49 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py @@ -430,6 +430,13 @@ class FastPath: "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 " diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py index aff50ab089f8..96b10dbc8508 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/lever_settings.py @@ -95,6 +95,11 @@ def apply(plan: "_dispatch.Plan") -> None: dict(batch_splitk=False, attn_block_splitk=False), ), # Attention. + ( + "ops.attention.swa_block", + dict(single_early_loads=on("swa_block_single_early_loads")), + dict(single_early_loads=False), + ), ( "ops.attention.multi_request_split_kv", dict(combine_gate=on("multi_request_combine_gate")), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py index 97fca5d58c5e..bf8174fc70d6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py @@ -140,6 +140,9 @@ "TRTLLM_MUSESPARK_DRAFTER_FINISH_NORM": PREFIX + "DRAFTER_FINISH_NORM=on", "TRTLLM_MUSESPARK_DRAFTER_SILU_EPILOGUE": PREFIX + "DRAFTER_SILU_EPILOGUE=on", "TRTLLM_MUSESPARK_TEP8_ATTN_BLOCK_SPLITK": PREFIX + "TEP8_ATTN_BLOCK_SPLITK=on", + "TRTLLM_MUSESPARK_SWA_BLOCK_SINGLE_EARLY_LOADS": ( + PREFIX + "SWA_BLOCK_SINGLE_EARLY_LOADS=on" + ), "TRTLLM_MUSESPARK_DRAFTER_MULTI_REQUEST_CELLS": (PREFIX + "DRAFTER_MULTI_REQUEST_CELLS=on"), } ) @@ -725,6 +728,7 @@ def _vocab_parallel_head(d: Deployment, t: Resolving, p: Mapping[str, object]) - "drafter_finish_norm": _ported(Kind.BITWISE), "drafter_silu_epilogue": _ported(Kind.BITWISE), "tep8_attn_block_splitk": _ported(Kind.BITWISE), + "swa_block_single_early_loads": _ported(Kind.BITWISE), # The qkv / gate_up / coefficient cells keep their row's split_k and are bitwise the # table's row; the o_proj and down_proj cells reproduced cuBLAS's bits on every checked # shape, which cannot be proven against a closed library. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py index 185cb6278eb3..e0ad144cfa60 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/swa_block.py @@ -1135,6 +1135,24 @@ def gate_head_stride(gate: torch.Tensor) -> int: return gate.stride(1) if gate.ndim == 3 else HEAD_DIM +#: ``swa_block_single_early_loads``, pushed here by ``lever_settings.apply`` (default off): +#: ``swa_block_attention`` (decode steps of one sequence) builds its split kernel with +#: ``EARLY_LOADS`` -- with the fused prep, every projection load the program makes after its +#: wait is issued at once (the query rows, the key block's step K/V rows and, in program +#: (0, h, 0), the step's K/V rows it writes to the cache) before any prep arithmetic: one +#: dependent round trip, not one per prep -- and its combine with ``EARLY_LOADS``: the partial O +#: is loaded with the (max, sum) words and, with the gate, the gate head before its wait (the +#: projection that wrote it completed before the split, which launches the combine, passed its +#: own wait). The same statements on the same values: the same bits. +_SINGLE_EARLY_LOADS = False + + +def configure(*, single_early_loads: bool = False) -> None: + """Set this module's lever-backed settings (``lever_settings.apply``), once per process.""" + global _SINGLE_EARLY_LOADS + _SINGLE_EARLY_LOADS = bool(single_early_loads) + + class SwaBlockAttention(OpWrapper): """A sliding-window layer's decode attention block, with fewer launches, the same bits. @@ -1303,6 +1321,9 @@ def __call__( part_ml = torch.empty( (num_splits * num_kv_heads * group_rows, 2), dtype=torch.float32, device=qkv.device ) + # ``swa_block_single_early_loads``: the one-request block's two launches issue their + # operand loads at once. The blocks-per-program split keeps the served launch. + early = _SINGLE_EARLY_LOADS and blocks_per_program <= 1 if blocks_per_program > 1: launch_split_bpp( qkv, kv_pool, block_tables, block_tables.stride(0), seq_lens, kv_scale_orig_quant, @@ -1313,6 +1334,17 @@ def __call__( blocks_per_program=blocks_per_program, max_registers=split_max_registers, ) # fmt: skip else: + if early: + from tensorrt_llm.logger import logger + + logger.info_once( + "[musespark_open] swa_block_single_early_loads: early loads on the " + "one-request block (split: one round trip for its projection loads" + f"{'' if fused_prep else ' -- off without the fused prep'}; combine: O with " + f"the (max, sum) words" + f"{', the gate before its wait' if gate is not None else ''})", + key="musespark_open.swa_block.single_early_loads", + ) _swa_split_kernel[ (num_splits, num_kv_heads, triton.cdiv(group_rows, ROWS_PER_PROGRAM)) ]( @@ -1347,6 +1379,9 @@ def __call__( num_warps=NUM_WARPS, num_stages=NUM_STAGES, launch_pdl=enable_pdl, + # Passed only when on: the default launch is the kernel as before. The split's + # early loads are the fused prep's operands, so they need it. + **({"EARLY_LOADS": True} if early and fused_prep else {}), ) _swa_combine_kernel[(group_rows, num_kv_heads)]( part_o, @@ -1367,6 +1402,7 @@ def __call__( LAUNCH_WITH_PDL=enable_pdl, num_warps=4, launch_pdl=enable_pdl, + **({"EARLY_LOADS": True} if early else {}), ) return out From e333e0373d9039f17333cc1cd2db16fc5ac177e6 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Tue, 10 Nov 2026 00:03:33 -0800 Subject: [PATCH 8/9] [None][fix] MuseSpark GB300: tep8_attn_block_splitk is numeric, not bitwise A batch-4 run on GB300 showed this lever changing the decoded text: against an otherwise identical arm it more than doubled the decode steps for the same token budget (287 -> 641), while the five other levers of this series reproduced the step structure exactly. Greedy decoding, same prompts, one knob apart. The cause is a premise that does not hold at TEP8. The lever replaces the 9..32-row QKV call with a split_k 4 cell, which keeps the served bits only if the served call at those rows is already that cell. It is not here: the path that arms a wide table at those rows, wide_projection_gathers, is TEP4-only on sm103, so the served 9..32-row QKV call is the built-in cuBLAS one, which sums K in a single slice. Four slices instead of one moves the last bits, and the sampled text follows. The o_proj half is unaffected and stays bitwise: its split_k 1 cells reproduce cuBLAS exactly, checked here at 9, 17 and 25 rows. So the lever becomes lossless-numeric, which also means auto needs a gate for it rather than treating it as free. It stays default off, and the module now says where its QKV cells are bitwise and where they are not, instead of claiming every row keeps its bits. Measured alongside: this lever costs 0.265 ms on the bs4 decode step (12.440 -> 12.705 ms, ~270 steps per arm), and is the whole of the 0.195 ms the six levers cost together. The other five sit within 0.021 ms of the baseline. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 28fb09a9b987a12c8cc60f66bf25a58ad6192f88) --- .../kernels/dense_gemms/dispatcher.py | 13 ++++++++++--- .../levers.py | 9 ++++++++- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py index 8399cebef02e..9f9296c7ac82 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/dense_gemms/dispatcher.py @@ -40,7 +40,13 @@ are loaded into shared memory before its ``griddepcontrol.wait``, under the attention that launches it early (the sliding-window cluster block, the full layers' combine); and QKV at split_k 4 (flashinfer's K slices and order) with a deeper ring where the copy's reduce-scatter mailbox -leaves room. Every row keeps its bits; the forks and joins are the served ones. +leaves room. The forks and joins are the served ones. + +Bits: the o_proj cells reproduce the built-in cuBLAS call. **The QKV cells do not, wherever the +served 9..32-row call is that same cuBLAS call** -- they sum K in four slices instead of one. They +keep the served bits only where a wide table is armed at those rows with a split_k 4 cell, and on +sm103 the path that arms one (``wide_projection_gathers``) is TEP4-only, so at TEP8 this lever +changes the decoded text. Its lever is lossless-numeric for that reason, not bitwise. With ``dense_gemm_prefetch`` a projection that ran also forks a side stream that asks the L2 to fetch the next projection's weight (QKV -> the same layer's attention output -> @@ -1391,7 +1397,8 @@ def install( logger.info( f"[musespark_open] tep8_attn_block_splitk: TP8 calls of {WIDE_MIN_ROWS}.." f"{WIDE_MAX_ROWS} rows by {steps} (o_proj split_k 1: the built-in cuBLAS call's bits, " - "its weight loaded before its dependency wait under the attention; QKV split_k 4: the " - "served cell's bits)" + "its weight loaded before its dependency wait under the attention; QKV split_k 4: " + "the served cell's bits only where a wide table is armed at these rows, else four K " + "slices where the built-in call uses one)" ) return plan diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py index bf8174fc70d6..25b2a626de67 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py @@ -727,7 +727,14 @@ def _vocab_parallel_head(d: Deployment, t: Resolving, p: Mapping[str, object]) - "drafter_conv_epilogue": _ported(Kind.BITWISE), "drafter_finish_norm": _ported(Kind.BITWISE), "drafter_silu_epilogue": _ported(Kind.BITWISE), - "tep8_attn_block_splitk": _ported(Kind.BITWISE), + # The o_proj half is bitwise: split_k 1 reproduces the built-in cuBLAS call, checked + # on this target. The QKV half is not, at TEP8. It replaces the 9..32-row call with a + # split_k 4 cell on the premise that the served call there is already that cell -- true + # only where a wide table is armed, and the path that arms one + # (``wide_projection_gathers``) is TEP4-only on sm103. At TEP8 the served call is + # cuBLAS, so the summation order changes and the decoded text moves with it, which a + # batch-4 run makes plain. Lossless-numeric, and it needs a gate before it is on. + "tep8_attn_block_splitk": _ported(Kind.NUMERIC), "swa_block_single_early_loads": _ported(Kind.BITWISE), # The qkv / gate_up / coefficient cells keep their row's split_k and are bitwise the # table's row; the o_proj and down_proj cells reproduced cuBLAS's bits on every checked From 0b10da0660158d400a641f848eb5713492437747 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian Date: Tue, 10 Nov 2026 10:10:08 -0800 Subject: [PATCH 9/9] [None][fix] test: widen the paged-context off-lever set for the ported levers test_a_paged_context_plans_every_fast_path_but_the_window_compaction asserts the exact set of levers that resolve off on a sm103 TP4 paged-context deployment. The decode-stack port adds five more levers, so that set grows from two to seven. All five are correctly off there: four are default off everywhere, and finalize_entry_trigger carries its default at TP8 only. No production code changes; the assertion was enumerating a set the port legitimately extends. Signed-off-by: Iman Tabrizian --- .../test_modeling_v2_drafter_context.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py index 41ca17f3b767..54788923af7d 100644 --- a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py @@ -263,5 +263,15 @@ def test_a_paged_context_plans_every_fast_path_but_the_window_compaction(monkeyp assert structural == {"drafter_window_compact"} off = {name for name, d in decisions.items() if not d.on} # The half expert window's rule keeps it to max_batch_size 8; nothing builds on it. - assert off == {"drafter_window_compact", "half_window_experts"} + # The ported decode-stack levers are default off, and the ones that do carry a + # default carry it at TP8 only, so on this TP4 context deployment none is armed. + assert off == { + "drafter_window_compact", + "half_window_experts", + "drafter_fused_prep", + "finalize_entry_trigger", + "multi_request_combine_gate", + "router_staged_gemm", + "tep8_wide_direct_transports", + } assert plan.paths == modeling.TARGET_PATHS[103, 4].served - off