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..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 @@ -116,6 +116,12 @@ class FastPath: "decode FMHA of steps with two or more requests as balanced page-split partials", "construct", ), + FastPath( + "multi_request_combine_gate", + "the multi-request decode FMHA combine storing the attention output's norm-gate", + "construct", + ("multi_request_split_kv",), + ), FastPath( "multi_request_prep_append", "every layer's K/V append folded into the decode prep of steps with two or more requests", @@ -129,6 +135,12 @@ class FastPath: ), # Dense projections and prefetch. FastPath("dense_gemms", "dense projections on one named kernel per shape", "construct"), + FastPath( + "tep8_batch_splitk", + "TP8 decode projections of 33..160 rows on batched split-K kernels", + "construct", + ("dense_gemms",), + ), FastPath( "dense_gemm_prefetch", "L2 prefetch of the next dense projection's weight", @@ -195,6 +207,13 @@ class FastPath: "construct", ("projection_gathers", "dense_gemms"), ), + FastPath( + "tep8_wide_direct_transports", + "pool and unpool slices published from their split-K GEMMs into 32-row multicast " + "transports, 16/24/32-row decode steps", + "construct", + ("wide_projection_gathers",), + ), FastPath( "projection_publish", "sharded pool/unpool GEMMs publishing into multicast transports their consumers poll, " @@ -209,22 +228,58 @@ class FastPath: "construct", ("projection_gathers", "post_attention_hop"), ), + FastPath( + "tep8_router_shard", + "the TP8 MoE prologue's router shard and fused unpool gather on 16..32-row decode steps", + "construct", + ("moe_prologue",), + ), + FastPath( + "tep8_router_shard_wide", + "that router shard on the built-in wide bits instead of each row's one-request bits", + "construct", + ("tep8_router_shard",), + ), FastPath( "pool_into_route", "pool projection and router shard writing into the routed kernel's transport", "construct", ("moe_prologue", "dense_gemms"), ), + FastPath( + "tep8_wide_direct_pool", + "the pool GEMM publishing into the TP8 wide route transport, 16..32-row decode steps", + "load", + ("pool_into_route", "tep8_router_shard"), + ), FastPath( "finalize_transport", "expert finalize all-reduce on the fused hop's transport", "construct", ), + FastPath( + "finalize_entry_trigger", + "the fused finalize releasing its dependent at its own dependency wait", + "construct", + ("finalize_transport",), + ), + FastPath( + "tep8_finalize_transport_wide", + "the TP8 finalize transport laid out for calls of 17..40 tokens", + "construct", + ("finalize_transport",), + ), FastPath( "router_gemm", "router logits of 9..512-row calls as a BF16x3 split-K tensor-core GEMM", "construct", ), + FastPath( + "router_staged_gemm", + "router logits of 129..256-row calls on the staged BF16x3 GEMM", + "construct", + ("router_gemm",), + ), FastPath( "wide_moe_allreduce", "MoE output all-reduce of 129..1024-row decode steps on the MNNVL two-shot", @@ -249,9 +304,52 @@ class FastPath: ("preattn_prep",), ), FastPath("replicated_experts", "ring-buddy expert replication, balanced", "load"), + FastPath( + "tep8_ring_buddy_wide", + "the balanced tep8 stacked path on decode calls above its 16-row cap", + "load", + ("replicated_experts",), + ), + FastPath( + "tep8_wide_balance", + "the NVFP4 layers' balanced placement above the stacked cap, on the wide assign program", + "load", + ("replicated_experts",), + ), + FastPath( + "ring_buddy_stats", + "a placement census per call, reported on an interval by TP rank 0", + "load", + ("replicated_experts",), + ), # persistent_experts and bf16_routed_experts run over the ring-buddy stack with # replicated_experts, else over the rank's own expert block (the owner-only window). FastPath("persistent_experts", "NVFP4 experts as one persistent launch", "load"), + FastPath( + "tep8_persistent_batch_experts", + "stacked NVFP4 expert calls of 9..32 rows on the persistent kernel's batch builds", + "load", + ("persistent_experts", "replicated_experts"), + ), + FastPath( + "tep8_experts_in_kernel_placement", + "the batch builds placing the step's selection themselves, with no assign launch", + "load", + ("tep8_persistent_batch_experts",), + ), + FastPath( + "tep8_experts_interleaved_fetch", + "the batch builds' two fetch warps moving the FC2 operand's k-blocks alternately", + "load", + ("tep8_persistent_batch_experts",), + ), + FastPath( + "tep8_experts_early_start", + "the expert assignment triggering the persistent kernel's batch builds at its own " + "dependency wait", + "load", + ("tep8_persistent_batch_experts",), + ), FastPath( "half_window_experts", "the ring predecessor's upper expert half on each TP4 rank, balanced per step", @@ -305,9 +403,60 @@ class FastPath: ("draft_head",), ), FastPath("drafter_fc_shard", "drafter fc projection column-sharded", "draft"), + FastPath( + "tep8_fc_shard_wide", + "the sharded drafter fc projection on TP8 calls of 9..128 rows", + "draft", + ("drafter_fc_shard",), + ), FastPath("fused_accept", "acceptance and forced acceptance in one launch", "draft"), FastPath("drafter_fixed_gemms", "drafter GEMMs on fixed low-M tactics", "draft"), + FastPath( + "drafter_conv_epilogue", + "the drafter coefficient GEMM of 9..32-row blocks with the input-side convolution in " + "its epilogue", + "draft", + ("drafter_block_conv", "drafter_fixed_gemms"), + ), + FastPath( + "drafter_finish_norm", + "each drafter hop's finish convolution, residual add and RMSNorm as one launch", + "draft", + ("drafter_block_decode",), + ), + FastPath( + "drafter_silu_epilogue", + "the drafter gate_up projection of 9..32-row blocks with silu(gate) * up in its epilogue", + "draft", + ("drafter_block_decode",), + ), + FastPath( + "swa_block_single_early_loads", + "the one-request sliding-window block's split and combine issuing their operand loads " + "at once", + "construct", + ("swa_block",), + ), + FastPath( + "tep8_attn_block_splitk", + "the TP8 attention block's QKV and o_proj of 9..32-row decode calls on batch split-K " + "cells, each staging its weight tile before its dependency wait", + "construct", + ("dense_gemms",), + ), + FastPath( + "drafter_multi_request_cells", + "the drafter GEMMs of 16, 24 and 32-row blocks on measured TP8 split-K cells", + "draft", + ("drafter_fixed_gemms",), + ), FastPath("context_kv_fork", "drafter context-K/V store on a side stream", "draft"), + FastPath( + "drafter_fused_prep", + "the drafter's page-table refresh, context K/V store and commit as fused launches", + "draft", + ("context_kv_fork",), + ), FastPath( "drafter_window_compact", "windowed drafter arena compaction in place, one read and write per window", 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..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 @@ -22,7 +22,14 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None: if block_conv or block_decode: from .kernels.drafter_block_decode import block_decode as _draft - _draft.install(shell.llm.draft_model, block_conv, block_decode) + _draft.install( + shell.llm.draft_model, + block_conv, + block_decode, + conv_epilogue=plan.enabled("drafter_conv_epilogue"), + finish_norm=plan.enabled("drafter_finish_norm"), + silu_epilogue=plan.enabled("drafter_silu_epilogue"), + ) if plan.enabled("wide_drafter_window_attention"): # On the block decode drafter_block_decode just bound. from .kernels.wide_drafter_window_attention import attention as _wide_window @@ -66,6 +73,8 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None: shell.llm.model, declined=plan.declined, paged_commit=plan.enabled("paged_context_commit"), + fused_prep=plan.enabled("drafter_fused_prep"), + fused_prep_min_requests=plan.param("drafter_fused_prep", "MIN_REQUESTS"), ) verify_fold = plan.enabled("verify_fold") path_walk_tail = plan.enabled("path_walk_tail") 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..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 @@ -27,6 +27,26 @@ 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 ``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. 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 -> @@ -75,6 +95,8 @@ CUBLAS, DIRECT, MAX_ROWS, + O_PROJ_SHAPE, + QKV_SHAPE, SPLITK, WIDE_MAX_ROWS, dense_gemm, @@ -146,6 +168,245 @@ } +#: ``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, attn_block_splitk: bool = False) -> None: + """Set whether ``install`` arms the batched split-K tables. + + ``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, _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]``, +#: 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} + + +#: ``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], ...]]: + """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: @@ -365,8 +626,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 @@ -374,6 +639,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 +679,35 @@ 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) + 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]``). @@ -647,6 +945,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. @@ -687,6 +989,106 @@ 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, + 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 +1177,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 +1218,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 +1283,36 @@ def install( "(MUSESPARK_OPEN_SHARDED_PROJECTIONS=1); without them the pooler GEMM reads the " "whole replicated weight" ) - plan = DenseGemmPlan(kernels, prefetch=prefetch, enable_pdl=LAUNCH_PDL, capability=capability) + 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" + ) + 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) layers = list(model.layers) planned = [] @@ -899,6 +1345,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 +1378,27 @@ 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" + ) + 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 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/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..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,6 +41,23 @@ attention and append on a multi-request split geometry) instead of ``dflash_window_attention``. +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 + 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. +* ``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``, which call ``_convolve``; with ``drafter_block_conv`` installed that is the fused convolution @@ -50,12 +67,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 +83,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 +97,315 @@ _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) + + +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 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. @@ -88,6 +417,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 +673,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 +761,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: @@ -425,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. @@ -435,7 +795,7 @@ def windowed_block_decode( ) prenormed = hops.mlp( layer_idx, - layer.mlp, + mlp, hidden_out, res_flat, mlp_conv_coefficients, @@ -444,7 +804,25 @@ 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. + 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 +831,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 +858,15 @@ 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, + silu_epilogue: 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 +882,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, silu_epilogue) def _builtin_block_convs(draft_model: nn.Module) -> int: @@ -513,7 +899,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 +910,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. -def _install_windowed_attention(draft_model: nn.Module | None, block_decode: bool) -> None: + 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, finish_norm: bool, silu_epilogue: bool +) -> None: enable_pdl = LAUNCH_PDL wrapped = isinstance( getattr(draft_model, "_dflash_flash_attention", None), WindowedBlockAttention @@ -549,3 +976,16 @@ 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)" + ) + 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/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..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,13 +66,15 @@ 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 +86,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 +354,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..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 @@ -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,13 +99,16 @@ 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 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 @@ -104,15 +117,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 +454,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 +478,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 +498,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 +623,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 +653,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 +667,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 +684,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 +709,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 +756,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 +782,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 +803,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 +825,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 +909,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 +1060,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 +1118,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 +1150,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 +1209,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 +1231,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 +1270,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 +1290,7 @@ def install( static_from = block failed: Optional[str] = None installed, refreshed = [], [] + widened_rows: List[int] = [] widened = [] halved = [] skipped = {} @@ -850,6 +1308,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 +1327,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 +1390,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 +1400,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 +1489,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 = maybe_pin_memory(torch.zeros(2 * PLACEMENT_STATS_FIELDS, dtype=torch.int64)) + + 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..140e02909959 --- /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,495 @@ +# 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..8fbd14b36b2b --- /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,311 @@ +# 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..96b10dbc8508 --- /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,186 @@ +# 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"), + attn_block_splitk=on("tep8_attn_block_splitk"), + ), + 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")), + 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.gemm.fixed_low_m_gemm", + dict(multi_request_cells=on("drafter_multi_request_cells")), + dict(multi_request_cells=False), + ), + ( + "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..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 @@ -82,7 +82,70 @@ 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" + ), + "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_SWA_BLOCK_SINGLE_EARLY_LOADS": ( + PREFIX + "SWA_BLOCK_SINGLE_EARLY_LOADS=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. RETIRED_PREFIXES: Mapping[str, str] = types.MappingProxyType({}) @@ -384,6 +447,70 @@ 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 _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. + + 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_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. + + 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 +525,65 @@ 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 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 + def _tep4_path_rule() -> Rule: return Rule( @@ -459,6 +635,111 @@ 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_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_wide_direct_pool": _ported_tep8( + Kind.FLAGGED, "pool_into_route's TP8 wide direct sites" + ), + "tep8_wide_direct_transports": _ported( + Kind.FLAGGED, (Param("SITES", WIDE_DIRECT_SITES, choices=WIDE_DIRECT_SITES),) + ), + "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_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( + Kind.FLAGGED, + "the 256-row wide balance", + ( + 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 + # 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_tep8( + Kind.FLAGGED, "placement computed inside the expert kernel" + ), + "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), + "drafter_silu_epilogue": _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 + # 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/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..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 @@ -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,35 @@ 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 +385,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 +405,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 +441,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 +468,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 +1060,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 +1073,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 +1095,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 +1112,12 @@ 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) @@ -1022,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. @@ -1190,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, @@ -1200,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)) ]( @@ -1234,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, @@ -1254,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 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..39109d885dae --- /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,1665 @@ +# 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/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/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 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..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 @@ -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,56 @@ 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..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 + + +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``..``WIDE_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 <= 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}..{WIDE_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 +153,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 +261,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 +385,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 +552,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 +585,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 +610,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 +652,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 +670,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 +733,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)), ( @@ -496,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" 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_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() 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..c2187c07062f --- /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,506 @@ +# 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() 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