From 50220877d7e14b4bd896791de344ee2a9e9e2685 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 15:05:38 -0700 Subject: [PATCH 01/62] feat: support PD decode and expert-parallel megakernels --- sgl_nvidia_mega/qwen3_5/draft/__init__.py | 10 ++++++- sgl_nvidia_mega/qwen3_5/draft/loop.py | 12 ++++++++ sgl_nvidia_mega/qwen3_5/moe/mkm.py | 30 ++++++++++++++++++++ sgl_nvidia_mega/qwen3_5/moe/route.py | 5 +++- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 31 +++++++++++++++++++-- sgl_nvidia_mega/qwen3_5/spec.py | 3 ++ test/unit/test_qwen3_5_runtime.py | 20 ++++++++++++- test/unit/test_qwen3_5_spec.py | 12 ++++++++ test/unit/test_qwen3_5_steps.py | 5 ++++ 9 files changed, 122 insertions(+), 6 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/draft/__init__.py b/sgl_nvidia_mega/qwen3_5/draft/__init__.py index 3923642e..f05deafc 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/draft/__init__.py @@ -70,6 +70,7 @@ "bind_extend_graphs", "begin_draft_round", ), + "srt/runtime_context.py": ("get_parallel", '"moe_ep_size"'), } ARCHITECTURE = "Qwen3_5ForCausalLMMTP" @@ -200,6 +201,11 @@ def _requested() -> frozenset: def _build(model, hf_config, ctx) -> MtpDraft: + from sglang.srt.runtime_context import get_parallel + + from sgl_nvidia_mega.qwen3_5 import runtime + + runtime.apply_expert_parallel(get_parallel().moe_ep_size) return MtpDraft( model, binding=ctx.binding, @@ -218,7 +224,9 @@ def _build(model, hf_config, ctx) -> MtpDraft: dtypes=target.REQUIREMENTS.dtypes, decode_graph_backends=target.REQUIREMENTS.decode_graph_backends, ), - features=frozenset({Feature.SPECULATIVE_DECODING}), + # under expert parallelism the draft MoE and the slots on it stay SGLang's + # (runtime.EP_STOCK_SLOTS) + features=frozenset({Feature.SPECULATIVE_DECODING, Feature.EXPERT_PARALLEL}), numerics=numerics(_requested()), # the draft's forward stays SGLang's: no width is built max_rows=0, diff --git a/sgl_nvidia_mega/qwen3_5/draft/loop.py b/sgl_nvidia_mega/qwen3_5/draft/loop.py index 905a8b47..f24e1992 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/loop.py +++ b/sgl_nvidia_mega/qwen3_5/draft/loop.py @@ -23,11 +23,15 @@ from __future__ import annotations +import logging + import torch from sgl_nvidia_mega.qwen3_5 import plan from sgl_nvidia_mega.qwen3_5.draft import LOOP_BS, SERVED_BS, layer +logger = logging.getLogger(__name__) + # The SGLang names this module reads (test/unit/test_contract.py pins them) SGLANG_READS = { "srt/kernel_plugins/nvidia/eagle_binding.py|srt/megakernels/eagle_binding.py": ( @@ -55,6 +59,7 @@ class DraftLoop: def __init__(self, parts, *, num_steps: int): self.parts = parts self.num_steps = num_steps + self._filled = False if "draft_q1" in parts.slots and parts.attention.q1 is None: from sgl_nvidia_mega.qwen3_5.draft.attention import build_q1 @@ -79,6 +84,13 @@ def begin_round(self, io) -> None: and not io.from_served_extend ): mkd.fill_fc(io.topk_index, io.hidden_states) + if not self._filled: + # the first round after a prefill, or on a PD decode after the transfer + self._filled = True + logger.info( + "NVIDIA megakernel: draft_loop fills its first step's fc from the " + "round's input (no served draft extend before it)" + ) @torch.no_grad() def step(self, io) -> None: diff --git a/sgl_nvidia_mega/qwen3_5/moe/mkm.py b/sgl_nvidia_mega/qwen3_5/moe/mkm.py index f8962bf6..56cfb2b8 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/mkm.py +++ b/sgl_nvidia_mega/qwen3_5/moe/mkm.py @@ -131,6 +131,36 @@ def unsupported(mlp) -> str: return "" +def routing_unsupported(mlp) -> str: + """Why SGLang's trtllm-gen runner cannot take slot route_topk's materialized routing + under expert parallelism ("" when it can): a softmax top-10, renormalized, over 512 + global experts on the NVFP4 trtllm-gen weights; the runner keeps global ids and + reads this rank's window from local_expert_offset.""" + from sglang.srt.layers.moe.utils import RoutingMethodType + + from sgl_nvidia_mega.qwen3_5.sm103.moe import mkm + + experts, cfg = mlp.experts, mlp.topk.topk_config + if experts.w13_weight.dtype != torch.uint8 or experts.w13_weight.shape[0] != ( + experts.num_local_experts + ): + return ( + f"w13_weight {experts.w13_weight.dtype} {tuple(experts.w13_weight.shape)}" + ) + if not ( + experts.num_experts == mkm.EXPERTS + and experts.top_k == cfg.top_k == mkm.TOP_K + and cfg.renormalize + and cfg.scoring_func == "softmax" + and cfg.correction_bias is None + and not cfg.use_grouped_topk + and shared_topk_slots(cfg) == 0 + and experts.routing_method_type == RoutingMethodType.RenormalizeNaive + ): + return "routing other than a softmax top-10, renormalized, over 512 experts" + return "" + + def routed(mlp, maps: tuple, identity: torch.Tensor, logits, x_fp4, x_sf): """The deferred-finalize output of ``mlp``'s routed experts from the router front's logits and quantized input (``maps``: the layer's weight TMA maps; ``identity``: diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index 9b1362d6..f7d169d6 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -129,7 +129,10 @@ def build( if not fused: return Bound(route=True, shared=shared) for i, mlp in zip(range(m.start_layer, m.end_layer), layers): - reason = mkm.unsupported(mlp) + expert_parallel = mlp.experts.num_local_experts != mlp.experts.num_experts + reason = ( + mkm.routing_unsupported(mlp) if expert_parallel else mkm.unsupported(mlp) + ) if reason: return Bound(route=True, reason=f"mkm: layer {i}: {reason}") intermediate = layers[0].experts.w2_weight.shape[-1] * 2 diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 3116165d..5c31a181 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -7,10 +7,13 @@ from __future__ import annotations +import logging import os import msgspec +logger = logging.getLogger(__name__) + # The kernel slots that run by default, and the ones that run only when requested # (until step-level timing and accuracy runs move them to the default list) SHIPPED_SLOTS: tuple[str, ...] = () @@ -44,6 +47,9 @@ "draft_skv2", ) SLOTS_ENV = "SGLANG_NVIDIA_MEGAKERNEL_SLOTS" +# These slots run all 512 experts themselves. Under expert parallelism, keep the +# route-only and boundary slots but let SGLang's EP path run the routed experts. +EP_STOCK_SLOTS = frozenset({"mkm", "mkm_wide", "draft_moe"}) def _requested_slots() -> frozenset: @@ -58,6 +64,18 @@ def _requested_slots() -> frozenset: SLOTS = _requested_slots() +def apply_expert_parallel(ep_size: int) -> frozenset: + """Drop ``EP_STOCK_SLOTS`` from ``SLOTS`` when the MoE runs expert-parallel + (``ep_size`` > 1); returns the dropped names. Every process calls it before reading + ``SLOTS`` (config check, target build, draft build).""" + global SLOTS + if ep_size <= 1: + return frozenset() + dropped = SLOTS & EP_STOCK_SLOTS + SLOTS = SLOTS - EP_STOCK_SLOTS + return dropped + + def slot_reasons() -> list[str]: """Why the requested slots cannot be served ([] when they can).""" known = ( @@ -128,7 +146,7 @@ def slot_reasons() -> list[str]: "deadline_s", "binding", ), - "srt/runtime_context.py": ("get_parallel", '"tp_size"'), + "srt/runtime_context.py": ("get_parallel", '"tp_size"', '"moe_ep_size"'), } @@ -144,16 +162,23 @@ def build(model, ctx): """Each package's build, the slot plan per width and the executor; ``ctx`` is SGLang's BuildContext.""" import torch + from sglang.srt.runtime_context import get_parallel from sgl_nvidia_mega.qwen3_5 import attention, boundary, gdn, moe from sgl_nvidia_mega.qwen3_5 import plan as slot_plan from sgl_nvidia_mega.qwen3_5.runtime import parity from sgl_nvidia_mega.qwen3_5.runtime.executor import Qwen35Verify + dropped = apply_expert_parallel(get_parallel().moe_ep_size) + if dropped: + logger.info( + "NVIDIA megakernel: Qwen3.5 under expert parallelism (moe_ep_size %d): " + "slots %s run SGLang's EP path", + get_parallel().moe_ep_size, + sorted(dropped), + ) multi = slot_plan.MULTI_REQUEST in SLOTS if multi: - from sglang.srt.runtime_context import get_parallel - multi = slot_plan.multi_request(get_parallel().tp_size, SLOTS) widths = slot_plan.candidates(ctx.widths, ctx.tokens_per_request, multi) if not widths: diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index 3d2f760b..cc99418e 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -178,6 +178,7 @@ def binding_reasons() -> list[str]: def _config_disable_reasons(cfg, hf_config) -> list[str]: + runtime.apply_expert_parallel(cfg.ep_size or 1) return text_config_reasons(hf_config) + binding_reasons() + runtime.slot_reasons() @@ -197,6 +198,8 @@ def _build(model, hf_config, ctx): Feature.MULTI_NODE_TP, # HiCache: covered steps read only running requests' pages Feature.KV_OFFLOAD, + # TEP (MoE TP x EP): the slots in runtime.EP_STOCK_SLOTS stay off + Feature.EXPERT_PARALLEL, } ), numerics=NUMERICS, diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 7d7858ae..e342fdae 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -23,6 +23,11 @@ def registered(monkeypatch): monkeypatch.setattr(registry, "_REGISTRY", {}) monkeypatch.delenv(plugin.PARITY_SORT_ENV, raising=False) plugin.register() + # the target and draft builds read moe_ep_size (expert parallelism); no process groups here + monkeypatch.setattr( + "sglang.srt.runtime_context.get_parallel", + lambda: SimpleNamespace(tp_size=8, moe_ep_size=1), + ) # ---------------------------------------------------------------- the covered steps @@ -71,7 +76,7 @@ def test_the_build_covers_more_requests_only_where_a_slot_runs( monkeypatch.setattr(runtime, "SLOTS", frozenset()) monkeypatch.setattr( "sglang.srt.runtime_context.get_parallel", - lambda: SimpleNamespace(tp_size=tp_size), + lambda: SimpleNamespace(tp_size=tp_size, moe_ep_size=1), ) model = SimpleNamespace( model=SimpleNamespace(start_layer=0, end_layer=0, layers=[]) @@ -110,6 +115,19 @@ def test_verify_sample_is_refused_with_the_switch_that_replaces_it(monkeypatch): ] +def test_expert_parallel_keeps_route_and_boundary_slots(monkeypatch): + requested = frozenset( + {"route", "mkm", "mkm_wide", "hop", "gdn_push", "draft_moe", "draft_hop"} + ) + monkeypatch.setattr(runtime, "SLOTS", requested) + assert runtime.apply_expert_parallel(1) == frozenset() + assert runtime.SLOTS == requested + assert runtime.apply_expert_parallel(8) == frozenset( + {"mkm", "mkm_wide", "draft_moe"} + ) + assert runtime.SLOTS == frozenset({"route", "hop", "gdn_push", "draft_hop"}) + + def test_the_runtime_gives_the_executor_the_build_contexts_binding(monkeypatch): monkeypatch.setattr(runtime, "SLOTS", frozenset()) model = SimpleNamespace( diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index f9256ded..18dc4c66 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -85,6 +85,18 @@ def test_tp4_on_one_tray_is_accepted(): assert _refusal(tp_size=4, nnodes=1) == "" +def test_the_decode_side_of_pd_is_accepted(): + assert _refusal(disaggregation_mode="decode") == "" + + +def test_the_prefill_side_of_pd_is_refused(): + # SGLang refuses ReplaySSM on a PD prefill server itself + message = _refusal( + disaggregation_mode="prefill", enable_linear_replayssm_spec=False + ) + assert "a PD prefill server runs no covered step" in message, message + + BREAKABLE = with_phase( default_cuda_graph_config(), Phase.DECODE, backend=Backend.BREAKABLE ) diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index cca4068c..fa98b6b2 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -589,6 +589,11 @@ def test_the_route_slot_is_not_built_for_a_layer_it_cannot_serve(monkeypatch): model=SimpleNamespace(start_layer=0, end_layer=1, layers=layers) ) monkeypatch.setattr(runtime, "SLOTS", frozenset({"route"})) + # runtime.build reads moe_ep_size (expert parallelism); no process groups here + monkeypatch.setattr( + "sglang.srt.runtime_context.get_parallel", + lambda: SimpleNamespace(tp_size=8, moe_ep_size=1), + ) ctx = SimpleNamespace(widths=(7,), tokens_per_request=7, deadline_s=60.0) want = "layer 0: router gate torch.float16 (512, 4096)" assert route.build(model, 7, fused=False).reason == want From af6dbce47acaa20e9a69ec2636359a83fd1ade22 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 15:08:16 -0700 Subject: [PATCH 02/62] feat: cover every multi-request verify and draft graph --- sgl_nvidia_mega/qwen3_5/attention/route.py | 3 + sgl_nvidia_mega/qwen3_5/draft/__init__.py | 5 + sgl_nvidia_mega/qwen3_5/draft/attention.py | 10 +- sgl_nvidia_mega/qwen3_5/draft/extend.py | 42 +++- sgl_nvidia_mega/qwen3_5/draft/layer.py | 35 +++- sgl_nvidia_mega/qwen3_5/draft/loop.py | 37 +++- sgl_nvidia_mega/qwen3_5/draft/parts.py | 30 +++ sgl_nvidia_mega/qwen3_5/plan.py | 10 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 8 + .../qwen3_5/sm103/draft/src/draft_top1.cu | 190 ++++++++++++++++++ sgl_nvidia_mega/qwen3_5/sm103/draft/top1.py | 187 +++++++++++++++++ test/kernels/test_qwen3_5_draft_xchg.py | 132 +++++++++++- test/kernels/test_qwen3_5_gdn.py | 12 +- test/unit/test_qwen3_5_runtime.py | 104 +++++++++- test/unit/test_qwen3_5_spec.py | 10 +- test/unit/test_qwen3_5_steps.py | 6 +- 16 files changed, 793 insertions(+), 28 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_top1.cu create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/draft/top1.py diff --git a/sgl_nvidia_mega/qwen3_5/attention/route.py b/sgl_nvidia_mega/qwen3_5/attention/route.py index bf7595f1..eb4acd11 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/route.py +++ b/sgl_nvidia_mega/qwen3_5/attention/route.py @@ -45,6 +45,9 @@ class SplitKVGraphVariants(msgspec.Struct, frozen=True): min_length: int # capture order matters where the variants share a graph memory pool capture_labels: ClassVar[tuple[str, ...]] = (SPLIT_KV_STOCK, SPLIT_KV_LONG) + # select picks the stock label past one request (and the steps read the label only + # for one), so SGLang captures one graph at every larger batch size + max_batch_size: ClassVar[int] = 1 def select(self, forward_batch) -> str: if forward_batch.batch_size != 1: diff --git a/sgl_nvidia_mega/qwen3_5/draft/__init__.py b/sgl_nvidia_mega/qwen3_5/draft/__init__.py index f05deafc..00415257 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/draft/__init__.py @@ -99,6 +99,11 @@ def numerics(slots: frozenset) -> Numerics: # draft steps on 2.2's two-row kernels (the split fc, the hop's all-reduce points, the # two-token MoE), the rest SGLang's ops LOOP_BS = (1, 2) +# with slot draft_multi, the draft graphs of two or more requests (every batch size up to +# 128, as plan.REQUESTS) whose top-1 is the sharded exchange in place of the logits +# all-gather; the steps of 3 to 128 run SGLang's ops but for the split fc (up to its widest +# batch) +MULTI_LOOP_BS = tuple(range(2, 129)) class MtpDraft: diff --git a/sgl_nvidia_mega/qwen3_5/draft/attention.py b/sgl_nvidia_mega/qwen3_5/draft/attention.py index 9297bde7..a1d7da2f 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/attention.py +++ b/sgl_nvidia_mega/qwen3_5/draft/attention.py @@ -47,6 +47,9 @@ class Q1Variants: # stock first, as 2.2 captures them (variants share a graph memory pool) capture_labels = (Q1_STOCK, Q1_LONG) + # select picks Q1_STOCK past one request (q1 decodes one row), so SGLang captures + # one graph at every larger batch size + max_batch_size = 1 def __init__(self, min_length: int = Q1_MIN_LENGTH): self.min_length = min_length @@ -307,7 +310,12 @@ def decode(self, q8, meta, kv, rows_per_request: int): and _q1_label() != Q1_STOCK ): return self.q1(q8, kv, meta).view(-1, self.heads * self.head_dim) - if self.skv2 is not None and rows_per_request > 1 and self.skv2.long(): + if ( + self.skv2 is not None + and rows_per_request > 1 + and q8.shape[0] == rows_per_request + and self.skv2.long() + ): return self.skv2.run(attn, q8, kv, meta) kwargs = {} if rows_per_request != 1: diff --git a/sgl_nvidia_mega/qwen3_5/draft/extend.py b/sgl_nvidia_mega/qwen3_5/draft/extend.py index de103c0b..d13669ff 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/extend.py +++ b/sgl_nvidia_mega/qwen3_5/draft/extend.py @@ -23,7 +23,8 @@ import torch -from sgl_nvidia_mega.qwen3_5.draft import SERVED_BS, layer +from sgl_nvidia_mega.qwen3_5 import plan +from sgl_nvidia_mega.qwen3_5.draft import MULTI_LOOP_BS, SERVED_BS, layer # The SGLang names this module reads (test/unit/test_contract.py pins them) SGLANG_READS = { @@ -53,12 +54,24 @@ def __init__(self, parts, *, tokens_per_request: int): self.width = tokens_per_request def step_for(self, bs: int): + # more requests with slot draft_multi where the verify steps of several are covered + from sgl_nvidia_mega.qwen3_5.runtime import SLOTS + + if ( + bs in MULTI_LOOP_BS + and self.parts.top1x is not None + and plan.multi_request(self.parts.binding.tp_size, SLOTS) + ): + return self.step return self.step if bs == SERVED_BS else None @torch.no_grad() def step(self, io) -> None: parts = self.parts parts.note_engaged("draft_extend", io.num_requests) + if io.num_requests > SERVED_BS: + self._multi_step(io) + return mkd = parts.mkd if mkd is not None: embeds = parts.model.embed_tokens(io.input_ids) @@ -84,3 +97,30 @@ def step(self, io) -> None: else: io.topk_index_out.copy_(layer.top1(parts, selected).view(-1, 1)) io.topk_p_out.fill_(1.0) + + def _multi_step(self, io) -> None: + """More requests (slot draft_multi): SGLang's model step on every verify row, + then the selected rows' sharded top-1 (torch.argmax's choice, a NaN first) + straight into SGLang's outputs.""" + from sgl_nvidia_mega.qwen3_5.sm103.draft.top1 import NAN_KEY_TORCH_ARGMAX + + parts = self.parts + x = layer.fc_input(parts, io.input_ids, io.hidden_states) + hidden = layer.layer( + parts, + x, + io.positions, + io.out_cache_loc, + io.attn, + io.kv, + rows_per_request=io.tokens_per_request, + ) + selected = torch.index_select(hidden, 0, io.select_index, out=io.hidden_out) + layer.top1_into( + parts, + selected, + io.topk_index_out, + 0, + nan_key=NAN_KEY_TORCH_ARGMAX, + topk_p=io.topk_p_out, + ) diff --git a/sgl_nvidia_mega/qwen3_5/draft/layer.py b/sgl_nvidia_mega/qwen3_5/draft/layer.py index 547dc46c..33dc81ea 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/layer.py +++ b/sgl_nvidia_mega/qwen3_5/draft/layer.py @@ -74,11 +74,27 @@ def layer(parts, x, positions, cache_loc, meta, kv, *, rows_per_request: int): lay = parts.layer residual = x - hop = parts.hop if rows_per_request == 1 else None + rows = x.shape[0] + # the draft hop takes the steps whose rows fit the target hop's mailbox; the + # single-launch MoE the steps of at most its token count + hop = ( + parts.hop + if rows_per_request == 1 + and parts.hop is not None + and parts.hop.hop.supports(rows) + else None + ) # a draft step of one request pushes its o_proj into the hop and releases the next # launch early; of two requests it takes only the hop's all-reduce points push = hop if x.shape[0] == 1 else None - extend_hop = parts.hop if rows_per_request > 1 and EXTEND_MNNVL else None + extend_hop = ( + parts.hop + if rows_per_request > 1 + and EXTEND_MNNVL + and parts.hop is not None + and parts.hop.hop.supports(rows) + else None + ) attn_out = parts.attention( lay.input_layernorm(x), positions, @@ -92,7 +108,11 @@ def layer(parts, x, positions, cache_loc, meta, kv, *, rows_per_request: int): else None ), ) - moe = parts.moe if rows_per_request == 1 else None + moe = ( + parts.moe + if rows_per_request == 1 and parts.moe is not None and parts.moe.ready(rows) + else None + ) if hop is not None: moe_in, residual = hop.ffn_input( attn_out, @@ -134,3 +154,12 @@ def top1(parts, hidden): logits = tensor_model_parallel_all_gather(logits) logits = logits[:, : lm.logits_processor.vocab_size] return torch.argmax(logits, dim=-1) + + +def top1_into(parts, hidden, tokens, column: int, *, nan_key: int, **outputs) -> None: + """Slot draft_multi's top-1 of the final-norm rows ``hidden``: the LM head on this + rank's vocabulary shard, then the sharded argmax exchange, which writes the tokens + into ``tokens[:, column]`` (and ``outputs``: topk_p, positions, mrope_positions).""" + lm = parts.lm + logits = lm.logits_processor._compute_lm_head(hidden, lm.lm_head, None) + parts.top1x(logits, tokens, column, nan_key=nan_key, **outputs) diff --git a/sgl_nvidia_mega/qwen3_5/draft/loop.py b/sgl_nvidia_mega/qwen3_5/draft/loop.py index f24e1992..4defa7bd 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/loop.py +++ b/sgl_nvidia_mega/qwen3_5/draft/loop.py @@ -28,7 +28,7 @@ import torch from sgl_nvidia_mega.qwen3_5 import plan -from sgl_nvidia_mega.qwen3_5.draft import LOOP_BS, SERVED_BS, layer +from sgl_nvidia_mega.qwen3_5.draft import LOOP_BS, MULTI_LOOP_BS, SERVED_BS, layer logger = logging.getLogger(__name__) @@ -71,7 +71,8 @@ def step_for(self, bs: int): if not plan.multi_request(self.parts.binding.tp_size, SLOTS): return self.step if bs == SERVED_BS else None - return self.step if bs in LOOP_BS else None + multi = MULTI_LOOP_BS if "draft_multi" in self.parts.slots else () + return self.step if bs in LOOP_BS + multi else None def begin_round(self, io) -> None: """On the host before a served draft graph replays: the first step's fc when @@ -102,6 +103,8 @@ def step(self, io) -> None: positions, mrope = io.positions, io.mrope_positions if parts.mkd is not None and io.num_requests == SERVED_BS: self._mkd_steps(io, steps) + elif parts.top1x is not None and io.num_requests in MULTI_LOOP_BS: + self._multi_steps(io, steps) else: token, hidden = io.topk_index.view(-1), io.hidden_states for i in range(steps): @@ -116,10 +119,12 @@ def step(self, io) -> None: mrope.sub_(steps) def _fc(self, token, hidden): - """A step's fc output outside the exchange: slot draft_xchg's row-split fc, else - SGLang's.""" + """A step's fc output outside the exchange: slot draft_xchg's row-split fc up to + its widest batch, else SGLang's.""" + from sgl_nvidia_mega.qwen3_5.draft.xchg import FC_SPLIT_MAX_TOKENS + mkd = self.parts.mkd - if mkd is None: + if mkd is None or token.shape[0] > FC_SPLIT_MAX_TOKENS: return layer.fc_input(self.parts, token, hidden) return mkd.split_fc(self.parts.model.embed_tokens(token), hidden) @@ -140,6 +145,28 @@ def _model_step(self, io, i: int, x): rows_per_request=1, ) + def _multi_steps(self, io, steps: int) -> None: + """The steps of more requests with slot draft_multi: SGLang's model step (with + the split fc where slot draft_xchg runs), then the sharded top-1, which writes + the token column and advances the positions in one launch; the next step embeds + that column. Drafts SGLang's split argmax would choose (a NaN as -1e30).""" + from sgl_nvidia_mega.qwen3_5.sm103.draft.top1 import NAN_KEY_SPLIT_ARGMAX + + tokens = io.draft_tokens_out + token, hidden = io.topk_index.view(-1), io.hidden_states + for i in range(steps): + hidden = self._model_step(io, i, self._fc(token, hidden)) + layer.top1_into( + self.parts, + hidden, + tokens, + i + 1, + nan_key=NAN_KEY_SPLIT_ARGMAX, + positions=io.positions, + mrope_positions=io.mrope_positions, + ) + token = tokens[:, i + 1] + def _mkd_steps(self, io, steps: int) -> None: """The steps with slot draft_xchg, in 2.2's order: fc from the exchange, the layer, the hidden half of the next fc on a side stream, the shard head, the diff --git a/sgl_nvidia_mega/qwen3_5/draft/parts.py b/sgl_nvidia_mega/qwen3_5/draft/parts.py index 513f6a57..f9e59ddf 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/parts.py +++ b/sgl_nvidia_mega/qwen3_5/draft/parts.py @@ -64,6 +64,9 @@ def __init__(self, draft, layer): self.moe: Optional[Any] = None # draft.moe.DraftMoe (slot draft_moe) self.hop: Optional[Any] = None # draft.hop.DraftHop (slot draft_hop) self.mkd: Optional[Any] = None # draft.xchg.MkdBoundary (slot draft_xchg) + # sm103.draft.top1.Top1Exchange: the multi-request steps' sharded top-1 (slot + # draft_multi) + self.top1x: Optional[Any] = None # (entry, bs) pairs whose capture logged the engagement line self.engaged: set = set() @@ -86,6 +89,8 @@ def build(cls, draft) -> Parts: parts.hop = parts._build_hop() if "draft_xchg" in slots: parts.mkd = parts._build_mkd() + if "draft_multi" in slots: + parts.top1x = parts._build_top1x() logger.info( "NVIDIA megakernel: Qwen3.5 MTP draft parts built (slots %s; MoE %s, hop " "%s, exchange %s)", @@ -175,6 +180,31 @@ def _build_mkd(self): ) return mkd + def _build_top1x(self): + from sgl_nvidia_mega.qwen3_5.sm103.draft.top1 import Top1Exchange + + lm, b = self.lm, self.binding + rows = int(lm.lm_head.weight.shape[0]) + start = int(lm.lm_head.shard_indices.org_vocab_start_index) + vocab = int(lm.logits_processor.vocab_size) + # the shard layout the all-gathered logits have: rank-major, no padding rows + if start != b.tp_rank * rows or rows * b.tp_size != vocab: + raise RuntimeError( + "NVIDIA megakernel: slot draft_multi cannot run: the LM head's " + f"{rows}-row shard at {start} is not rank {b.tp_rank}'s slice of " + f"{vocab} ids" + ) + return Top1Exchange( + world=b.tp_size, + rank=b.tp_rank, + vocab_local=rows, + vocab_start=start, + vocab_size=vocab, + group=b.tp_device_group, + cpu_group=b.tp_cpu_group, + device=self.device, + ) + def note_engaged(self, entry: str, bs: int) -> None: """Log once per (entry, bs) captured: the engagement line every rank prints.""" if torch.cuda.is_current_stream_capturing() and (entry, bs) not in self.engaged: diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index 89c52f91..154cf1ff 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -4,8 +4,8 @@ A slot runs only where it is requested (``SGLANG_NVIDIA_MEGAKERNEL_SLOTS``), built and inside its row bounds; every other op is SGLang's own. The covered steps are the -target verify of 1, 2 or 4 requests (``REQUESTS``: SGLang's decode graphs of batch -sizes 1, 2 and 4); a width of 2 or 4 requests is covered only where a slot runs there. +target verify of 1 to 128 requests (``REQUESTS``: every decode graph batch size SGLang +captures up to 128); a width of more than one request is covered only where a slot runs. ``route`` (the router GEMM and the routed experts' NVFP4 input in one launch) runs up to 16 rows, ``mkm`` (on its outputs, the top-10 and both expert GEMMs in one launch) up to 7, and with ``mkm_wide`` (its 16-token engine) up to 16. ``gdn_verify`` (the GatedDeltaNet verify chain with its launch choices changed) covers @@ -25,8 +25,10 @@ import msgspec -MAX_ROWS = 28 # the widest covered step: four requests of 7 verify rows (TP8 bs4) -REQUESTS = (1, 2, 4) # the requests of a covered step (SGLang's decode graph sizes) +MAX_ROWS = 896 # the widest covered step: 128 requests of 7 verify rows (TP8 bs128) +# Every decode graph size, so no batch pads to the next power of two to reach a covered +# one; each slot's own row bounds still determine whether the plugin serves the step. +REQUESTS = tuple(range(1, 129)) # the TP sizes whose steps of more than one request are covered when requested (checked at # TP8); at the others a covered step is one request MULTI_REQUEST_TP = (8,) diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 5c31a181..52ec6ca0 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -45,6 +45,9 @@ "draft_hop", "draft_front", "draft_skv2", + # draft_loop and draft_extend also serve TP8's draft graphs of 2 to 128 requests, + # and those of two or more take the sharded top-1 (with multi_request) + "draft_multi", ) SLOTS_ENV = "SGLANG_NVIDIA_MEGAKERNEL_SLOTS" # These slots run all 512 experts themselves. Under expert parallelism, keep the @@ -126,6 +129,11 @@ def slot_reasons() -> list[str]: reasons.append(f"{SLOTS_ENV}: slot {name!r} runs inside entry 'draft_loop'") if "draft_hop" in SLOTS and "hop" not in SLOTS: reasons.append(f"{SLOTS_ENV}: slot 'draft_hop' runs on slot 'hop''s mailbox") + if "draft_multi" in SLOTS and "multi_request" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'draft_multi' serves the draft steps of two or more " + "requests ('multi_request')" + ) if "draft_skv2" in SLOTS and "draft_extend" not in SLOTS: reasons.append( f"{SLOTS_ENV}: slot 'draft_skv2' runs inside entry 'draft_extend'" diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_top1.cu b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_top1.cu new file mode 100644 index 00000000..4ef844fb --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_top1.cu @@ -0,0 +1,190 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// The sharded top-1 of a multi-row MTP draft step (draft_top1_xchg): every rank holds its +// vocabulary shard's bf16 logits [m, V_local]; the global argmax of each row is chosen without +// gathering the logits. Every CTA scans one chunk of one row into an order key +// (value key << 32 | ~global id: the larger value, then the lower id, wins), folds it into the +// row's key with atomicMax, and the last CTA to arrive exchanges the m row keys with the other +// ranks through the multicast mapping (mkd_xchg's transport: epoch-tagged 64-bit words, three +// slots rotating by epoch, bounded waits that trap), takes each row's largest key and writes the +// token into its column of the draft tokens, 1.0 into the top-k probability, and advances the +// positions (and the three M-RoPE rows) by one. +// +// The choice equals the argmax of the all-gathered logits: the value order is the float order +// with -0.0 as +0.0; ties go to the lowest vocabulary id, which is rank-major; a NaN maps to +// nan_key (0xFFFFFFFF: the first NaN wins, as torch.argmax; the key of -1e30: as the split +// argmax of SGLang's draft_topk1_postprocess). +// +// Built with NVRTC (MKD_WORLD = the TP size, 4 or 8). + +#include "mkd_common.cuh" + +#ifndef MKD_WORLD +#define MKD_WORLD 8 +#endif + +namespace sglang { + +namespace draft_top1 { + +inline constexpr int kSlots = 3; +inline constexpr int kMaxRows = 128; // the draft graphs' widest batch +inline constexpr int kThreads = 256; +inline constexpr int kVec = 8; // bf16 per 16-byte load +inline constexpr uint32_t kSpinLimit = 1u << 27; + +SGL_DEVICE uint32_t value_key(float v, uint32_t nan_key) { + if (v != v) return nan_key; + if (v == 0.0f) v = 0.0f; // -0.0 orders as +0.0 + const uint32_t bits = __float_as_uint(v); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +SGL_DEVICE uint64_t pack(uint32_t vkey, uint32_t id) { + return (static_cast(vkey) << 32) | static_cast(0xFFFFFFFFu - id); +} + +SGL_DEVICE uint64_t max64(uint64_t a, uint64_t b) { + return a > b ? a : b; +} + +SGL_DEVICE void mc_store(uint64_t* addr, uint32_t payload, uint32_t tag) { + const uint64_t w = static_cast(payload) | (static_cast(tag) << 32); + asm volatile("multimem.st.relaxed.sys.global.b64 [%0], %1;" ::"l"(addr), "l"(w) : "memory"); +} + +SGL_DEVICE uint32_t poll(const uint64_t* src, uint32_t tag) { + uint64_t w; + uint32_t spins = 0; + do { + asm volatile("ld.relaxed.sys.global.u64 %0, [%1];" : "=l"(w) : "l"(src) : "memory"); + if (++spins > kSpinLimit) __trap(); + } while (static_cast(w >> 32) != tag); + return static_cast(w); +} + +} // namespace draft_top1 + +struct DraftTop1Params { + const bf16_t* logits; // [m, row_stride] this rank's shard logits, written by an earlier grid + unsigned long long* keys; // [kMaxRows] per-row fold, zero between calls + uint32_t* epoch; // [2]: call counter shared by all ranks (from 1), CTA arrivals + uint64_t* mc; // multicast base of [kSlots, world, kMaxRows, 2] words + const uint64_t* uc; // this rank's own copy of it + int64_t* tokens; // tokens[row * tokens_stride + column] = the global argmax + float* topk_p; // [m] 1.0, or null + int64_t* positions; // [m] += 1, or null + int64_t* mrope; // [3, mrope_stride] += 1 in columns [0, m), or null + int64_t tokens_stride; + int64_t mrope_stride; + int64_t row_stride; // elements between logits rows + int32_t m; + int32_t vocab_local; // V_local: columns of a row this rank scans + int32_t vocab_start; // global id of this rank's column 0 + int32_t vocab_size; // ids >= vocab_size are padding and never chosen + int32_t column; + int32_t rank; + uint32_t nan_key; +}; + +// grid (chunks, m): chunk c of row r scans columns [c * kThreads * kVec, ...) of the row. +template +SGL_DEVICE void draft_top1_xchg_kernel(const DraftTop1Params& p) { + using namespace device; + using namespace draft_top1; + using vec_t = AlignedVector; + constexpr int kWarps = kThreads / kWarpThreads; + __shared__ uint64_t s_red[kWarps]; + __shared__ bool s_last; + + const int row = static_cast(blockIdx.y); + const int chunk = static_cast(blockIdx.x); + const int span = kThreads * kVec; + const int col0 = chunk * span + static_cast(threadIdx.x) * kVec; + const int ids_end = min(p.vocab_local, p.vocab_size - p.vocab_start); + const bf16_t* src = p.logits + static_cast(row) * p.row_stride; + + PDLWaitPrimary(); + uint64_t best = 0; + if (col0 + kVec <= ids_end && (p.row_stride % kVec) == 0) { + vec_t v; + v.load(src + col0); +#pragma unroll + for (int j = 0; j < kVec; ++j) { + best = max64(best, pack(value_key(cast(v[j]), p.nan_key), + static_cast(p.vocab_start + col0 + j))); + } + } else { + for (int j = 0; j < kVec && col0 + j < ids_end; ++j) { + best = max64(best, pack(value_key(cast(src[col0 + j]), p.nan_key), + static_cast(p.vocab_start + col0 + j))); + } + } +#pragma unroll + for (uint32_t mask = kWarpThreads / 2; mask >= 1; mask >>= 1) { + best = max64(best, __shfl_xor_sync(0xffffffffu, best, mask, kWarpThreads)); + } + const int warp = static_cast(threadIdx.x / kWarpThreads); + if (get_lane_id() == 0) s_red[warp] = best; + __syncthreads(); + if (threadIdx.x == 0) { + uint64_t b = s_red[0]; +#pragma unroll + for (int w = 1; w < kWarps; ++w) b = max64(b, s_red[w]); + atomicMax(p.keys + row, static_cast(b)); + __threadfence(); + const uint32_t total = gridDim.x * gridDim.y; + s_last = atomicAdd(p.epoch + 1, 1u) == total - 1; + } + __syncthreads(); + if (!s_last) return; + + // the last CTA: this rank's row keys -> every rank, then the largest key of each row + __shared__ uint32_t s_tag; + if (threadIdx.x == 0) { + __threadfence(); + s_tag = *reinterpret_cast(p.epoch); + } + __syncthreads(); + const uint32_t tag = s_tag; + const int64_t slot = static_cast(tag % kSlots); + const int t = static_cast(threadIdx.x); + if (t < p.m) { + const uint64_t k = atomicExch(p.keys + t, 0ull); // read the fold and leave it zero + uint64_t* dst = p.mc + ((slot * kWorld + p.rank) * kMaxRows + t) * 2; + mc_store(dst, static_cast(k), tag); + mc_store(dst + 1, static_cast(k >> 32), tag); + } + if (t < p.m) { + uint64_t win = 0; +#pragma unroll + for (int r = 0; r < kWorld; ++r) { + const uint64_t* w = p.uc + ((slot * kWorld + r) * kMaxRows + t) * 2; + const uint32_t lo = poll(w, tag); + const uint32_t hi = poll(w + 1, tag); + win = max64(win, (static_cast(hi) << 32) | lo); + } + const int64_t id = static_cast(0xFFFFFFFFu - static_cast(win)); + p.tokens[t * p.tokens_stride + p.column] = id; + if (p.topk_p != nullptr) p.topk_p[t] = 1.0f; + if (p.positions != nullptr) p.positions[t] += 1; + if (p.mrope != nullptr) { +#pragma unroll + for (int a = 0; a < 3; ++a) p.mrope[a * p.mrope_stride + t] += 1; + } + } + __syncthreads(); + if (threadIdx.x == 0) { + p.epoch[1] = 0u; + p.epoch[0] = tag + 1; + } +} + +} // namespace sglang + +using sglang::DraftTop1Params; + +extern "C" __global__ void __launch_bounds__(256) draft_top1_xchg(const DraftTop1Params __grid_constant__ p) { + sglang::draft_top1_xchg_kernel(p); +} diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/top1.py b/sgl_nvidia_mega/qwen3_5/sm103/draft/top1.py new file mode 100644 index 00000000..9b14eb2a --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/top1.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""The sharded top-1 of a multi-row MTP draft step (``src/draft_top1.cu``): the argmax of +every row over the TP group's vocabulary shards in one launch, without the logits' +all-gather. It writes the token into a column of the draft tokens, 1.0 into the top-k +probability, and advances the positions (and M-RoPE positions) by one.""" + +from __future__ import annotations + +import ctypes +from functools import cache +from typing import Optional + +import torch + +from sgl_nvidia_mega.common import jit +from sgl_nvidia_mega.qwen3_5.sm103.draft.xchg import ( + NAN_KEY_SPLIT_ARGMAX, + NAN_KEY_TORCH_ARGMAX, + SRC, + WORLDS, + symmetric_words, +) + +__all__ = ["NAN_KEY_SPLIT_ARGMAX", "NAN_KEY_TORCH_ARGMAX", "Top1Exchange"] + +# draft_top1::kSlots / kMaxRows / kThreads * kVec +SLOTS = 3 +MAX_ROWS = 128 +THREADS = 256 +CHUNK = THREADS * 8 + + +class DraftTop1Params(ctypes.Structure): + """sglang::DraftTop1Params.""" + + _fields_ = [ + (name, ctypes.c_void_p) + for name in ( + "logits", + "keys", + "epoch", + "mc", + "uc", + "tokens", + "topk_p", + "positions", + "mrope", + ) + ] + [ + ("tokens_stride", ctypes.c_int64), + ("mrope_stride", ctypes.c_int64), + ("row_stride", ctypes.c_int64), + ("m", ctypes.c_int32), + ("vocab_local", ctypes.c_int32), + ("vocab_start", ctypes.c_int32), + ("vocab_size", ctypes.c_int32), + ("column", ctypes.c_int32), + ("rank", ctypes.c_int32), + ("nan_key", ctypes.c_uint32), + ] + + +@cache +def _kernel(device_index: int, world: int) -> jit.Kernel: + major, minor = torch.cuda.get_device_capability(device_index) + cubin = jit.compile_source( + (SRC / "draft_top1.cu").read_text(), + f"q35_draft_top1_{world}", + f"sm_{major}{minor}a", + ("-std=c++17", "-default-device", f"-DMKD_WORLD={world}"), + headers={"mkd_common.cuh": (SRC / "mkd_common.cuh").read_text()}, + ) + return jit.load(cubin, "draft_top1_xchg", (THREADS, 1, 1), 0) + + +def build(device_index: int, world: int) -> None: + """Compile the kernel (no compile inside a capture).""" + if world not in WORLDS: + raise ValueError(f"the sharded top-1 serves TP {WORLDS}, not {world}") + _kernel(device_index, world) + + +class Top1Exchange: + """The exchange buffer, epoch and per-row fold of one rank. Construct on every rank of + ``group`` (a collective), before any capture.""" + + def __init__( + self, + *, + world: int, + rank: int, + vocab_local: int, + vocab_start: int, + vocab_size: int, + group, + cpu_group, + device: torch.device, + ): + import torch.distributed as dist + + build(device.index, world) + self.world, self.rank, self.device = world, rank, device + self.vocab_local, self.vocab_start = vocab_local, vocab_start + self.vocab_size = vocab_size + self._buf, self._handle, self.mc = symmetric_words( + SLOTS * world * MAX_ROWS * 2, group, device + ) + self.uc = self._buf.data_ptr() + # [epoch from 1, CTA arrivals] + self.epoch = torch.tensor([1, 0], dtype=torch.int32, device=device) + self.keys = torch.zeros(MAX_ROWS, dtype=torch.int64, device=device) + torch.cuda.synchronize(device) + dist.barrier(group=cpu_group) + + def __call__( + self, + local_logits: torch.Tensor, + tokens: torch.Tensor, + column: int, + *, + nan_key: int, + topk_p: Optional[torch.Tensor] = None, + positions: Optional[torch.Tensor] = None, + mrope_positions: Optional[torch.Tensor] = None, + ) -> None: + """``tokens[:, column]`` = the global argmax of each row of ``local_logits`` + [m, V_local] (this rank's shard, bf16); ``topk_p`` [m] (any trailing 1) = 1.0; + ``positions`` [m] and the three rows of ``mrope_positions`` [3, >= m] += 1.""" + m = int(local_logits.shape[0]) + if not 1 <= m <= MAX_ROWS: + raise ValueError(f"sharded top-1: {m} rows (1-{MAX_ROWS})") + if ( + local_logits.dtype != torch.bfloat16 + or local_logits.shape[1] != self.vocab_local + or local_logits.stride(1) != 1 + ): + raise ValueError( + f"sharded top-1: logits {local_logits.dtype} " + f"{tuple(local_logits.shape)} strides {tuple(local_logits.stride())} " + f"(want bf16 [m, {self.vocab_local}], unit column stride)" + ) + if tokens.dtype != torch.int64 or tokens.dim() != 2 or tokens.shape[0] < m: + raise ValueError( + f"sharded top-1: tokens {tokens.dtype} {tuple(tokens.shape)}" + ) + for name, t in (("topk_p", topk_p), ("positions", positions)): + if t is not None and (t.numel() < m or not t.is_contiguous()): + raise ValueError(f"sharded top-1: {name} {tuple(t.shape)}") + if positions is not None and positions.dtype != torch.int64: + raise ValueError(f"sharded top-1: positions {positions.dtype}") + if mrope_positions is not None and ( + mrope_positions.dtype != torch.int64 + or mrope_positions.shape[0] != 3 + or mrope_positions.stride(1) != 1 + ): + raise ValueError(f"sharded top-1: M-RoPE {tuple(mrope_positions.shape)}") + p = DraftTop1Params( + logits=local_logits.data_ptr(), + keys=self.keys.data_ptr(), + epoch=self.epoch.data_ptr(), + mc=self.mc, + uc=self.uc, + tokens=tokens.data_ptr(), + topk_p=0 if topk_p is None else topk_p.data_ptr(), + positions=0 if positions is None else positions.data_ptr(), + mrope=0 if mrope_positions is None else mrope_positions.data_ptr(), + tokens_stride=tokens.stride(0), + mrope_stride=0 if mrope_positions is None else mrope_positions.stride(0), + row_stride=local_logits.stride(0), + m=m, + vocab_local=self.vocab_local, + vocab_start=self.vocab_start, + vocab_size=self.vocab_size, + column=column, + rank=self.rank, + nan_key=nan_key, + ) + kernel = _kernel(self.device.index, self.world) + chunks = (self.vocab_local + CHUNK - 1) // CHUNK + _launch_2d(kernel, (chunks, m), p, self.device) + + +def _launch_2d(kernel: jit.Kernel, grid: tuple[int, int], params, device) -> None: + """``jit.launch`` of a (x, y) grid on the current stream, as a programmatic dependent.""" + stream = torch.cuda.current_stream(device).cuda_stream + jit.launch(kernel, (grid[0], grid[1], 1), [params], stream) diff --git a/test/kernels/test_qwen3_5_draft_xchg.py b/test/kernels/test_qwen3_5_draft_xchg.py index 943e3035..63a799e4 100644 --- a/test/kernels/test_qwen3_5_draft_xchg.py +++ b/test/kernels/test_qwen3_5_draft_xchg.py @@ -11,7 +11,10 @@ - the memo rows equal the split fc's partials summed in rank order (checked at construction), and the split fc's gather and rank-order sum are exact; - a whole draft round (five steps, the extend exchange with extend fc) replays from a CUDA - graph with fresh inputs, equal to its eager run. + graph with fresh inputs, equal to its eager run; +- the sharded top-1 of slot draft_multi (``sm103/draft/top1``) picks, for 1-16 rows, the + argmax of the all-gathered logits under both NaN rules (ties to the lowest id, -0.0 as + +0.0), writes the token column, 1.0 and positions + 1, and replays from a CUDA graph. The bitwise comparison with the kernels of SGLang branch nvidia/qwen35-agentx-gb300 needs that branch installed as ``sglang`` and is not part of this test. @@ -299,5 +302,132 @@ def run_round(st): dist.barrier(group=cpu_group) +V_TOP1 = 31040 # TP8's LM head shard rows + + +@pytest.fixture(scope="module") +def top1x(ranks): + from sgl_nvidia_mega.qwen3_5.sm103.draft.top1 import Top1Exchange + + rank, world, group, cpu_group = ranks + return Top1Exchange( + world=world, + rank=rank, + vocab_local=V_TOP1, + vocab_start=rank * V_TOP1, + vocab_size=world * V_TOP1, + group=group, + cpu_group=cpu_group, + device=torch.device("cuda", torch.cuda.current_device()), + ) + + +def _top1_logits(rank, m, case): + g = _gen(7000 + 31 * rank + m) + logits = _randn(m, V_TOP1, g=g) + if case == "ties": # equal maxima inside a shard and across ranks + logits = (torch.round(logits.float() * 2) / 2).to(torch.bfloat16) + elif case == "nan": + if rank == 1: + logits[0, 17] = float("nan") + logits[m - 1, 5] = float("nan") + if rank == 2: + logits[0, 3] = float("nan") + elif case == "zero": # -0.0 ties +0.0: the lower id wins + logits = torch.full_like(logits, -1.0) + if rank == 2: + logits[:, 40] = -0.0 + if rank == 3: + logits[:, 10] = 0.0 + return logits + + +def _top1_want(logits, group, nan): + full = _all_gather(logits, group).permute(1, 0, 2).reshape(logits.shape[0], -1) + full = full.float() + if nan == "split": + full = torch.where(torch.isnan(full), torch.full_like(full, -1e30), full) + return torch.argmax( + full, dim=1 + ) # the first maximum (a NaN first under torch's rule) + + +@pytest.mark.parametrize("case", ["random", "ties", "nan", "zero"]) +@pytest.mark.parametrize("nan", ["split", "torch"]) +@pytest.mark.parametrize("m", [1, 2, 4, 8, 16, 32, 64, 128]) +def test_sharded_top1_is_the_gathered_argmax(ranks, top1x, m, nan, case): + from sgl_nvidia_mega.qwen3_5.sm103.draft import top1 + + rank, world, group, _ = ranks + key = top1.NAN_KEY_SPLIT_ARGMAX if nan == "split" else top1.NAN_KEY_TORCH_ARGMAX + with torch.inference_mode(): + logits = _top1_logits(rank, m, case) + tokens = torch.full((m, NUM_STEPS), -1, dtype=torch.int64, device="cuda") + topk_p = torch.zeros(m, 1, dtype=torch.float32, device="cuda") + positions = torch.arange(m, dtype=torch.int64, device="cuda") * 10 + mrope = torch.arange(3 * (m + 1), dtype=torch.int64, device="cuda").view(3, -1) + mrope0 = mrope.clone() + top1x( + logits, + tokens, + 3, + nan_key=key, + topk_p=topk_p, + positions=positions, + mrope_positions=mrope[:, :m], + ) + torch.cuda.synchronize() + want = _top1_want(logits, group, nan) + assert torch.equal(tokens[:, 3], want) + assert (tokens[:, [0, 1, 2, 4, 5]] == -1).all() + assert (topk_p == 1.0).all() + assert torch.equal(positions, torch.arange(m, device="cuda") * 10 + 1) + assert torch.equal(mrope[:, :m], mrope0[:, :m] + 1) + assert torch.equal(mrope[:, m:], mrope0[:, m:]) + + +def test_sharded_top1_replays_from_a_cuda_graph(ranks, top1x): + from sgl_nvidia_mega.qwen3_5.sm103.draft import top1 + + rank, world, group, cpu_group = ranks + m = 8 + with torch.inference_mode(): + logits = _top1_logits(rank, m, "random") + tokens = torch.zeros(m, NUM_STEPS, dtype=torch.int64, device="cuda") + positions = torch.zeros(m, dtype=torch.int64, device="cuda") + + def run(): + for col in range(1, NUM_STEPS): + top1x( + logits, + tokens, + col, + nan_key=top1.NAN_KEY_SPLIT_ARGMAX, + positions=positions, + ) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + dist.barrier(group=cpu_group) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + torch.cuda.synchronize() + dist.barrier(group=cpu_group) + for case in ("random", "ties", "nan"): + logits.copy_(_top1_logits(rank, m, case)) + positions.zero_() + graph.replay() + torch.cuda.synchronize() + want = _top1_want(logits, group, "split") + for col in range(1, NUM_STEPS): + assert torch.equal(tokens[:, col], want) + assert (positions == NUM_STEPS - 1).all() + dist.barrier(group=cpu_group) + + if __name__ == "__main__": multigpu_pytest_main(__name__, __file__, num_gpus=(4,)) diff --git a/test/kernels/test_qwen3_5_gdn.py b/test/kernels/test_qwen3_5_gdn.py index c76b9af5..ff3cf7d7 100644 --- a/test/kernels/test_qwen3_5_gdn.py +++ b/test/kernels/test_qwen3_5_gdn.py @@ -31,7 +31,7 @@ H, K, HV, V = 2, 128, 8, 128 # per rank at TP8: 16 / 64 heads over 8 D = 2 * H * K + HV * V # the conv's q / k / v channels -DRAFT, WIDTH, SLOTS = 7, 4, 9 +DRAFT, WIDTH, SLOTS = 7, 4, 136 def _gpu_ready() -> bool: @@ -56,11 +56,11 @@ def _projection(bs, views, g): return qkvz[:, :D] if views else qkvz[:, :D].contiguous() -# a covered step's requests (slot gdn_verify covers B1, B2 and B4) and their slots -REQUESTS = (1, 2, 4) -STATE_SLOTS = (4, 7, 0, 8) -RING_SLOTS = (3, 5, 0, 8) -REPLAY_SLOTS = (1, 2, 4, 6) +# a covered step's requests (slot gdn_verify covers B1 to B128) and their slots +REQUESTS = (1, 2, 4, 8, 16, 32, 64, 128) +STATE_SLOTS = (4, 7, 0, 8, *range(39, 11, -1), *range(135, 39, -1)) +RING_SLOTS = (3, 5, 0, 8, *range(12, 136)) +REPLAY_SLOTS = (1, 2, 4, 6, *range(10, 134)) @pytest.mark.parametrize("bs", REQUESTS) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index e342fdae..535b9dad 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -205,6 +205,12 @@ def test_the_draft_loop_routes_q1_by_kv_length_as_2_2( built = draft.SPEC.build("lm", SimpleNamespace(context_len=4096), _draft_ctx()) variants = built.graph_variants["draft_loop"] assert variants.capture_labels == ("q1_stock", "q1_long") + # past max_batch_size every batch takes q1_stock, so SGLang captures one graph there + assert variants.max_batch_size == 1 + long = SimpleNamespace(attention_kv_len_lower_bound=10**6, seq_lens_cpu=None) + for bs in range(variants.max_batch_size + 1, 129): + long.batch_size = bs + assert variants.select(long) == "q1_stock" batch = SimpleNamespace( batch_size=batch_size, attention_kv_len_lower_bound=lower_bound, @@ -245,6 +251,11 @@ def test_the_draft_extend_takes_the_split_kv_variants_with_slot_draft_skv2( {"draft_loop", "draft_skv2"}, "slot 'draft_skv2' runs inside entry 'draft_extend'", ), + ( + {"draft_loop", "draft_multi"}, + "slot 'draft_multi' serves the draft steps of two or more requests " + "('multi_request')", + ), ], ) def test_the_draft_slots_name_what_they_need(monkeypatch, slots, want): @@ -256,7 +267,7 @@ class _DraftParts: """Parts with no slot built: the stock model step.""" def __init__(self, slots=frozenset(), tp_size=8): - self.slots, self.mkd, self.engaged = slots, None, [] + self.slots, self.mkd, self.top1x, self.engaged = slots, None, None, [] self.binding = SimpleNamespace(tp_size=tp_size) def note_engaged(self, entry, bs): @@ -334,6 +345,8 @@ def test_the_draft_loop_writes_each_steps_token_and_restores_the_positions( (8, {"multi_request"}, (1, 2)), (8, set(), (1,)), (4, {"multi_request"}, (1,)), + (8, {"multi_request", "draft_multi"}, tuple(range(1, 129))), + (4, {"multi_request", "draft_multi"}, (1,)), ], ) def test_the_draft_loop_serves_two_requests_at_tp8_on_request( @@ -342,8 +355,93 @@ def test_the_draft_loop_serves_two_requests_at_tp8_on_request( from sgl_nvidia_mega.qwen3_5.draft.loop import DraftLoop monkeypatch.setattr(runtime, "SLOTS", frozenset(slots)) - loop = DraftLoop(_DraftParts(tp_size=tp_size), num_steps=4) - assert tuple(bs for bs in (1, 2, 4) if loop.step_for(bs) == loop.step) == served + loop = DraftLoop(_DraftParts(slots=frozenset(slots), tp_size=tp_size), num_steps=4) + sizes = tuple(range(1, 130)) + assert tuple(bs for bs in sizes if loop.step_for(bs) == loop.step) == served + + +@pytest.mark.parametrize( + "tp_size, slots, built, served", + [ + (8, {"multi_request", "draft_multi"}, True, tuple(range(1, 129))), + (8, {"multi_request", "draft_multi"}, False, (1,)), + (4, {"multi_request", "draft_multi"}, True, (1,)), + (8, {"draft_multi"}, True, (1,)), + ], +) +def test_the_draft_extend_serves_more_requests_with_draft_multi( + monkeypatch, tp_size, slots, built, served +): + from sgl_nvidia_mega.qwen3_5.draft.extend import DraftExtend + + monkeypatch.setattr(runtime, "SLOTS", frozenset(slots)) + parts = _DraftParts(slots=frozenset(slots), tp_size=tp_size) + parts.top1x = object() if built else None + extend = DraftExtend(parts, tokens_per_request=7) + sizes = tuple(range(1, 130)) + assert tuple(bs for bs in sizes if extend.step_for(bs) == extend.step) == served + + +def test_a_multi_request_draft_loop_takes_the_sharded_top1(monkeypatch): + from sgl_nvidia_mega.qwen3_5.draft import layer + from sgl_nvidia_mega.qwen3_5.draft.loop import DraftLoop + + calls = [] + _stub_model_step(monkeypatch, calls) + + def top1_into( + parts, hidden, tokens, column, *, nan_key, positions, mrope_positions + ): + calls.append(("top1x", column, nan_key)) + tokens[:, column] = hidden[:, 0].to(torch.int64) % 1000 + positions.add_(1) + + monkeypatch.setattr(layer, "top1_into", top1_into) + parts = _DraftParts(slots=frozenset({"draft_multi"})) + parts.top1x = object() + loop = DraftLoop(parts, num_steps=3) + io = SimpleNamespace( + num_requests=4, + num_steps=3, + topk_index=torch.tensor([[1], [2], [3], [4]]), + hidden_states=torch.tensor([[1.0], [2.0], [3.0], [4.0]]), + positions=torch.tensor([10, 20, 30, 40]), + mrope_positions=None, + out_cache_loc=torch.arange(12), + steps=("m0", "m1"), + kv="kv", + draft_tokens_out=torch.full((4, 3), -1, dtype=torch.int64), + ) + loop.step(io) + assert parts.engaged == [("draft_loop", 4)] + assert [c[0] for c in calls if c[0] in ("top1", "top1x")] == ["top1x", "top1x"] + # each step embeds the column the previous exchange wrote + assert [c[1] for c in calls if c[0] == "fc"] == [[1, 2, 3, 4], [20, 30, 40, 50]] + assert io.draft_tokens_out.tolist() == [ + [1, 20, 210], + [2, 30, 310], + [3, 40, 410], + [4, 50, 510], + ] + # the request-major slots of step i, made dense + assert [c[2] for c in calls if c[0] == "layer"] == [[0, 3, 6, 9], [1, 4, 7, 10]] + assert io.positions.tolist() == [10, 20, 30, 40] + + +@pytest.mark.parametrize("requests, split", [(16, True), (32, False)]) +def test_the_split_fc_takes_the_draft_steps_up_to_its_widest_batch( + monkeypatch, requests, split +): + from sgl_nvidia_mega.qwen3_5.draft.loop import DraftLoop + + calls = [] + _stub_model_step(monkeypatch, calls) + parts = _DraftParts(slots=frozenset({"draft_multi"})) + parts.model = SimpleNamespace(embed_tokens=lambda t: t) + parts.mkd = SimpleNamespace(split_fc=lambda e, h: calls.append(("split",)) or h) + loop = DraftLoop(parts, num_steps=2) + loop._fc(torch.zeros(requests, dtype=torch.int64), torch.zeros(requests, 1)) + assert [c[0] for c in calls] == (["split"] if split else ["fc"]) @pytest.mark.parametrize("split_fc", [False, True]) diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 18dc4c66..5a8ff30d 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -214,11 +214,14 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): assert _on(plan.slots(7, every)) == sorted(plan.Slots.__struct_fields__) # two requests: the 16-row router front, the GDN chain and the hop's one-shot rows assert _on(plan.slots(14, every)) == ["gdn_pdl", "gdn_verify", "hop", "route"] - # four requests: past the router front's and the one-shot all-reduce's rows + # four requests and wider: past the router front's and one-shot all-reduce's rows assert _on(plan.slots(28, every)) == ["gdn_pdl", "gdn_verify"] - # SGLang offers every multiple of 7 up to MAX_ROWS; three requests are not covered + for rows in (56, 112, 224, 448, 896): + assert _on(plan.slots(rows, every)) == ["gdn_pdl", "gdn_verify"] + # SGLang offers every multiple of 7 up to MAX_ROWS; every request count is a candidate offered = tuple(range(7, plan.MAX_ROWS + 1, 7)) - assert plan.candidates(offered, 7) == (7, 14, 28) + assert plan.candidates(offered, 7) == offered + assert plan.candidates(offered + (903,), 7) == offered assert plan.candidates(offered, 7, multi=False) == (7,) # more than one request only when requested, and only at TP8; TP4 keeps the # one-request step it is checked at @@ -226,6 +229,7 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): assert not plan.multi_request(8, frozenset()) assert not plan.multi_request(4, frozenset({"multi_request"})) assert plan.covered((7, 14, 28), every, 7) == (7, 14, 28) + assert plan.covered((7, 448, 896), every, 7) == (7, 448, 896) # one request is always covered, more only where a slot runs assert plan.covered((7, 14, 28), plan.Facts(), 7) == (7,) assert plan.covered((7, 14, 28), plan.Facts(route=True, mkm=True), 7) == (7, 14) diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index fa98b6b2..1ceb4531 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -1598,7 +1598,11 @@ def batch(bs=1, bound=None, lens=None): assert route.VERIFY.select(batch(bound=8192)) == "split_kv_long" assert route.VERIFY.select(batch(bound=8191)) == "split_kv_stock" assert route.EXTEND.select(batch(bound=32767)) == "split_kv_stock" - assert route.VERIFY.select(batch(bs=2, bound=10**6)) == "split_kv_stock" + # past max_batch_size every batch takes the stock label, so SGLang captures one graph there + assert route.SplitKVGraphVariants.max_batch_size == 1 + for bs in range(route.SplitKVGraphVariants.max_batch_size + 1, 129): + assert route.VERIFY.select(batch(bs=bs, bound=10**6)) == "split_kv_stock" + assert route.EXTEND.select(batch(bs=bs, bound=10**6)) == "split_kv_stock" assert route.VERIFY.select(batch(lens=torch.tensor([8192]))) == "split_kv_long" assert route.VERIFY.select(batch(lens=torch.tensor([9000, 9000]))) == ( "split_kv_stock" From ccdd552bd32fbdfc27ebb009b16c700dcb62a14b Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 15:12:22 -0700 Subject: [PATCH 03/62] feat: extend fused attention across verify widths --- sgl_nvidia_mega/qwen3_5/attention/__init__.py | 15 +- sgl_nvidia_mega/qwen3_5/attention/mka.py | 85 ++++- sgl_nvidia_mega/qwen3_5/attention/route.py | 117 +++++-- sgl_nvidia_mega/qwen3_5/attention/split_kv.py | 112 +++++- sgl_nvidia_mega/qwen3_5/draft/__init__.py | 7 +- sgl_nvidia_mega/qwen3_5/draft/attention.py | 86 ++++- sgl_nvidia_mega/qwen3_5/plan.py | 10 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 31 +- sgl_nvidia_mega/qwen3_5/runtime/executor.py | 2 +- .../qwen3_5/sm103/attention/gate_quant.py | 50 +++ .../qwen3_5/sm103/attention/mka.py | 55 ++- .../qwen3_5/sm103/attention/src/mka.cu | 328 ++++++++++++++++++ sgl_nvidia_mega/qwen3_5/spec.py | 4 + test/kernels/test_qwen3_5_attention.py | 119 +++++++ test/unit/test_qwen3_5_runtime.py | 116 +++++++ test/unit/test_qwen3_5_spec.py | 31 +- test/unit/test_qwen3_5_steps.py | 96 +++++ 17 files changed, 1193 insertions(+), 71 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/attention/gate_quant.py diff --git a/sgl_nvidia_mega/qwen3_5/attention/__init__.py b/sgl_nvidia_mega/qwen3_5/attention/__init__.py index 52b66d69..288d37f9 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/attention/__init__.py @@ -47,7 +47,7 @@ "srt/models/qwen3_5.py:Qwen3_5AttentionDecoderLayer.self_attention": "db0f4d98e0fb1abf", } -ATTENTION_SLOTS = ("skv2", "mka", "mka_push") +ATTENTION_SLOTS = ("skv2", "skv2_cp", "mka", "mka_push") class Bound(msgspec.Struct, frozen=True): @@ -55,6 +55,7 @@ class Bound(msgspec.Struct, frozen=True): skv2: bool = False # slot skv2 serves every full-attention layer mka: bool = False + mka_wide: bool = False # mka's wide entries serve the steps of several requests mka_push: bool = False reason: str = "" # why a requested slot is not built @@ -87,10 +88,20 @@ def build(model, ctx) -> Bound: ) if reason: return Bound(reason=f"layer {i}: {reason}") + if "skv2_cp" in SLOTS: + reason = split_kv.cp_reason() + if reason: + return Bound(reason=f"slot skv2_cp: {reason}") if "mka" in SLOTS: for _, layer in layers: mka.pack(layer) - return Bound(skv2="skv2" in SLOTS, mka="mka" in SLOTS, mka_push="mka_push" in SLOTS) + mka_on = "mka" in SLOTS + return Bound( + skv2="skv2" in SLOTS, + mka=mka_on, + mka_wide=mka_on and all(layer.num_heads == 4 for _, layer in layers), + mka_push="mka_push" in SLOTS, + ) def pool_reasons(view, binding, skv2: bool = True, mka: bool = False) -> list[str]: diff --git a/sgl_nvidia_mega/qwen3_5/attention/mka.py b/sgl_nvidia_mega/qwen3_5/attention/mka.py index 03d620ce..cdf61a8d 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/mka.py +++ b/sgl_nvidia_mega/qwen3_5/attention/mka.py @@ -15,6 +15,12 @@ slot ``mka_push`` the epilogue also pushes its rows into the decode hop's mailbox where the MoE input's all-reduce will run the hop for them, and that all-reduce only consumes (the branch's MKA_PUSH). + +Steps of several requests (TP8, with ``multi_request``): up to ``FUSED_MAX_ROWS`` rows +the wide entries of the same kernels; past it the front path, which keeps SGLang's qkv +and o_proj GEMMs and fuses the rest bit for bit: the q / k norm, RoPE and gate split with +the FP8 q and K / V write in one launch (the attention front of slot ``draft_front``), +and the output gate with o_proj's static FP8 input quantization in one launch. """ from __future__ import annotations @@ -66,8 +72,13 @@ "k", "v", ), + "srt/layers/quantization/fp8_utils.py": ("flashinfer_bmm_fp8",), } +# the widest step the fused kernels serve; wider steps take the front path, which measured +# faster from four requests on (TP8, results/trt-1022/l2a micro e58d29c) +FUSED_MAX_ROWS = 14 + # the packed (qkv, o_proj) weights per full-attention layer id, filled at build WEIGHTS: dict = {} # the rows the slot has run at (each logged once, as engagement evidence) @@ -145,7 +156,13 @@ def step(layer, x, st): m = x.shape[0] if m not in _ENGAGED: _ENGAGED.add(m) - logger.info("NVIDIA megakernel: Qwen3.5 slot mka runs at %d rows", m) + logger.info( + "NVIDIA megakernel: Qwen3.5 slot mka runs at %d rows (%s)", + m, + "fused" if m <= FUSED_MAX_ROWS else "front path", + ) + if m > FUSED_MAX_ROWS: + return _front_step(layer, x, st) attn = layer.attn width = layer.num_heads * layer.head_dim q = torch.empty((m, width), dtype=torch.float8_e4m3fn, device=x.device) @@ -199,6 +216,72 @@ def step(layer, x, st): return out +def _front_step(layer, x, st): + """The front path of a wide step: SGLang's qkv GEMM, the front with the FP8 KV write, + the attention, the gate with o_proj's input quantization, SGLang's o_proj GEMM.""" + from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, + ) + from sglang.srt.layers.dp_attention import is_allocation_symmetric + from sglang.srt.layers.quantization.fp8_utils import flashinfer_bmm_fp8 + from sglang.srt.runtime_context import get_parallel + + from sgl_nvidia_mega.qwen3_5.attention import split_kv + from sgl_nvidia_mega.qwen3_5.sm103.attention.gate_quant import gate_quant + from sgl_nvidia_mega.qwen3_5.sm103.draft.front import ( + fused_qk_gemma_rmsnorm_rope_gate, + ) + + m = x.shape[0] + attn = layer.attn + width = layer.num_heads * layer.head_dim + qkv, _ = layer.qkv_proj(x) + q_gate, k, v = qkv.split( + [ + 2 * width, + layer.num_kv_heads * layer.head_dim, + layer.num_kv_heads * layer.head_dim, + ], + dim=-1, + ) + i = st.view.kv.layer_ids.index(attn.layer_id) + k_scale, v_scale = split_kv.kv_scales(attn, x.device) + rope = layer.rotary_emb + q, _, gate = fused_qk_gemma_rmsnorm_rope_gate( + q_gate, + k, + layer.q_norm.weight.data, + layer.k_norm.weight.data, + rope.cos_sin_cache, + st.positions, + layer.q_norm.variance_epsilon, + layer.num_heads, + layer.num_kv_heads, + layer.head_dim, + rope.rotary_dim, + has_gate=True, + mrope_axis_map=rope.axis_map, + v=v, + kv_cache=(st.view.kv.k[i], st.view.kv.v[i]), + cache_loc=st.fb.out_cache_loc, + k_scale=k_scale, + v_scale=v_scale, + ) + o = attn(q.reshape(m, width), None, None, st.fb).reshape(m, width) + oq = gate_quant(o.contiguous(), gate.reshape(m, width), layer.o_proj.input_scale) + # the same allocation as RowParallelLinear's output, which the boundary all-reduces + with use_symmetric_memory( + get_parallel().tp_group, disabled=not is_allocation_symmetric() + ): + return flashinfer_bmm_fp8( + oq, + layer.o_proj.weight, + layer.o_proj.input_scale, + layer.o_proj.weight_scale, + torch.bfloat16, + ) + + def _push_hop(layer, st, m: int): """Slot mka_push: the hop the epilogue pushes its rows into, where the MoE input's all-reduce runs the hop for them (it then only consumes); else None.""" diff --git a/sgl_nvidia_mega/qwen3_5/attention/route.py b/sgl_nvidia_mega/qwen3_5/attention/route.py index eb4acd11..3339bcfb 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/route.py +++ b/sgl_nvidia_mega/qwen3_5/attention/route.py @@ -15,11 +15,17 @@ ``capture_labels`` in order and replays the one ``select`` picks. A step reads the label being captured; where SGLang captures none (no runner takes the variants yet), every covered step takes split-KV attention. + +Slots ``skv2_cp`` / ``draft_skv2_cp`` (and ``draft_q1_cp`` for the draft loop) add a +third label, ``SPLIT_KV_CP``: one request past the measured per-site crossover takes +CP-lite (SGLang's read-split attention over ranks that hold the same KV head) in place +of SKV2, in a graph of its own, so the shorter steps' graphs stay as they are. """ from __future__ import annotations -from typing import ClassVar +import functools +from typing import ClassVar, Optional import msgspec @@ -36,6 +42,11 @@ SPLIT_KV_STOCK = "split_kv_stock" SPLIT_KV_LONG = "split_kv_long" +SPLIT_KV_CP = "split_kv_cp" +CP_SITES = {"verify": "skv2_cp", "draft": "draft_q1_cp", "extend": "draft_skv2_cp"} +# TP8 static A/B crossovers against SKV2 / SKV2-Q1: verify 136 Ki, draft step +# 180 Ki, draft extend 240 Ki committed KV tokens. +CP_MIN_KV = {"verify": 139264, "draft": 184320, "extend": 245760} class SplitKVGraphVariants(msgspec.Struct, frozen=True): @@ -50,17 +61,29 @@ class SplitKVGraphVariants(msgspec.Struct, frozen=True): max_batch_size: ClassVar[int] = 1 def select(self, forward_batch) -> str: - if forward_batch.batch_size != 1: - return SPLIT_KV_STOCK - length = forward_batch.attention_kv_len_lower_bound - if length is None: - lengths = forward_batch.seq_lens_cpu - if ( - lengths is not None - and lengths.device.type == "cpu" - and lengths.numel() == 1 - ): - length = int(lengths[0]) + length = one_request_length(forward_batch) + return ( + SPLIT_KV_LONG + if length is not None and length >= self.min_length + else SPLIT_KV_STOCK + ) + + +class SplitKVCPGraphVariants(SplitKVGraphVariants, frozen=True): + """``SplitKVGraphVariants`` with ``SPLIT_KV_CP`` for one request from + ``cp_min_length`` tokens of KV.""" + + cp_min_length: int = CP_MIN_KV["verify"] + capture_labels: ClassVar[tuple[str, ...]] = ( + SPLIT_KV_STOCK, + SPLIT_KV_LONG, + SPLIT_KV_CP, + ) + + def select(self, forward_batch) -> str: + length = one_request_length(forward_batch) + if length is not None and length >= self.cp_min_length: + return SPLIT_KV_CP return ( SPLIT_KV_LONG if length is not None and length >= self.min_length @@ -68,24 +91,76 @@ def select(self, forward_batch) -> str: ) +def one_request_length(forward_batch) -> Optional[int]: + """A host-side lower bound of the KV length of a one-request step, else None.""" + if forward_batch.batch_size != 1: + return None + length = forward_batch.attention_kv_len_lower_bound + if length is None: + lengths = forward_batch.seq_lens_cpu + if ( + lengths is not None + and lengths.device.type == "cpu" + and lengths.numel() == 1 + ): + length = int(lengths[0]) + return length + + VERIFY = SplitKVGraphVariants(min_length=8192) EXTEND = SplitKVGraphVariants(min_length=32768) -def split_kv(variants: SplitKVGraphVariants, forward_batch) -> bool: - """Whether this step takes split-KV attention: the label being captured, else (an - eager step) ``variants.select``; every step where SGLang captures no label of - ours or keeps no host lower bound.""" +def cp_min_length(site: str) -> Optional[int]: + """The KV length from which one request's ``site`` ("verify" / "draft" / "extend") + takes CP-lite, or None when its slot is not requested.""" + from sgl_nvidia_mega.qwen3_5.runtime import SLOTS + + if CP_SITES[site] not in SLOTS: + return None + return CP_MIN_KV[site] + + +def verify_variants() -> SplitKVGraphVariants: + """The target verify's variants: ``VERIFY``, with ``SPLIT_KV_CP`` under slot skv2_cp.""" + n = cp_min_length("verify") + if n is None: + return VERIFY + return _with_cp(VERIFY.min_length, n) + + +def extend_variants() -> SplitKVGraphVariants: + """The draft extend's variants: ``EXTEND``, with ``SPLIT_KV_CP`` under slot draft_skv2_cp.""" + n = cp_min_length("extend") + if n is None: + return EXTEND + return _with_cp(EXTEND.min_length, n) + + +@functools.cache +def _with_cp(min_length: int, cp_length: int) -> SplitKVCPGraphVariants: + return SplitKVCPGraphVariants(min_length=min_length, cp_min_length=cp_length) + + +def label(variants: SplitKVGraphVariants, forward_batch) -> str: + """This step's label of ``variants``: the one being captured, else (an eager step) + ``variants.select``; ``SPLIT_KV_LONG`` where SGLang captures no label of ours or + keeps no host lower bound.""" from sglang.srt.model_executor.runner_utils.capture_mode import ( get_capture_attention_variant, get_is_capture_mode, ) - label = get_capture_attention_variant() - if label in SplitKVGraphVariants.capture_labels: - return label == SPLIT_KV_LONG + captured = get_capture_attention_variant() + if captured in variants.capture_labels: + return captured if get_is_capture_mode() or not hasattr( forward_batch, "attention_kv_len_lower_bound" ): - return True - return variants.select(forward_batch) == SPLIT_KV_LONG + return SPLIT_KV_LONG + return variants.select(forward_batch) + + +def split_kv(variants: SplitKVGraphVariants, forward_batch) -> bool: + """Whether this step takes SKV2: ``label`` is ``SPLIT_KV_LONG``.""" + return label(variants, forward_batch) == SPLIT_KV_LONG diff --git a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py index 7005a6dd..5b960996 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py +++ b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py @@ -16,11 +16,14 @@ ``core`` takes an FP8 Q whose K / V rows are written (slot ``mka``'s prologue writes them itself) and runs SKV2 where the step takes split-KV attention (``route``), SGLang's attention otherwise. It reads the page table, the KV lengths and the pool from -the forward's view (``Step.view``). +the forward's view (``Step.view``). With slot ``skv2_cp``, a step of one request past +the measured crossover runs CP-lite (``cp_core``) instead of SKV2. """ from __future__ import annotations +import importlib +import inspect import logging import torch @@ -65,6 +68,11 @@ ), "srt/model_executor/forward_batch_info.py": ("out_cache_loc",), "srt/environ.py": ("envs", "SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR"), + "srt/layers/attention/trtllm_mha_backend.py": ( + "DEFAULT_WORKSPACE_SIZE_MB", + '"trtllm_mha_zero_workspace"', + ), + "srt/runtime_context.py": ("get_buffer",), } # The SGLang functions whose body this module mirrors, by the digest of their code @@ -170,19 +178,103 @@ def core(attn, q, st): step are written: [rows, heads * head_dim] in bf16.""" from sgl_nvidia_mega.qwen3_5.attention import route - split = route.split_kv(route.VERIFY, st.fb) - if (split, q.shape[0]) not in _ENGAGED: - _ENGAGED.add((split, q.shape[0])) - logger.info( - "NVIDIA megakernel: Qwen3.5 slot skv2 runs at %d rows (%s)", - q.shape[0], - "SKV2" if split else "SGLang's attention below the split-KV length", - ) - if not split: + label = route.label(route.verify_variants(), st.fb) + if (label, q.shape[0]) not in _ENGAGED: + _ENGAGED.add((label, q.shape[0])) + if label == route.SPLIT_KV_CP: + logger.info( + "NVIDIA megakernel: Qwen3.5 slot skv2_cp runs at %d rows (CP-lite from " + "%d tokens of KV)", + q.shape[0], + route.cp_min_length("verify"), + ) + else: + logger.info( + "NVIDIA megakernel: Qwen3.5 slot skv2 runs at %d rows (%s)", + q.shape[0], + "SKV2" + if label == route.SPLIT_KV_LONG + else "SGLang's attention below the split-KV length", + ) + if label == route.SPLIT_KV_STOCK: return attn(q, None, None, st.fb) + if label == route.SPLIT_KV_CP: + o = cp_core(attn, q, st.view.kv, st.view.attn, st.fb, ROWS) + if o is not None: + return o return skv2(attn, q, st.view.kv, st.view.attn) +def cp_lite(): + """SGLang's process CP-lite (q35_cp_lite.CpLite, built by its TRT-LLM MHA backend + when a CP graph variant is requested) when it can serve one request + (``attend(min_bs=)``), else None.""" + try: + module = importlib.import_module("sglang.srt.layers.attention.q35_cp_lite") + except ImportError: + return None + if "min_bs" not in inspect.signature(module.CpLite.attend).parameters: + return None + return module.CpLite._instance + + +def cp_reason() -> str: + """Why the CP graph variants cannot run ("" when they can).""" + try: + module = importlib.import_module("sglang.srt.layers.attention.q35_cp_lite") + except ImportError: + return "this SGLang has no CP-lite (q35_cp_lite)" + if "min_bs" not in inspect.signature(module.CpLite.attend).parameters: + return "this SGLang's CP-lite serves no single request (attend has no min_bs)" + if module.CpLite._instance is None: + return "SGLang's attention backend did not build CP-lite" + return "" + + +def cp_core(attn, q, kv, meta, forward_batch, q_len: int): + """CP-lite on one request's ``q_len`` FP8 query rows ``q`` of ``attn`` over the pool + ``kv`` at ``meta``'s pages and lengths (this step's rows written): [rows, heads * + head_dim] bf16, or None where CP-lite declines (decided by shapes alone, so every + rank of the KV group agrees).""" + from sglang.kernels.ops.attention.utils import canonicalize_stride + from sglang.srt.layers.attention.trtllm_mha_backend import ( + DEFAULT_WORKSPACE_SIZE_MB, + ) + from sglang.srt.runtime_context import get_buffer + + cp = cp_lite() + if cp is None: + return None + i = kv.layer_ids.index(attn.layer_id) + # the pool's NHD rows as HND pages, as TRT-LLM MHA hands them to CP-lite + caches = tuple( + canonicalize_stride( + buf.view(-1, PAGE_SIZE, heads, attn.head_dim).permute(0, 2, 1, 3) + ) + for buf, heads in ((kv.k[i], attn.tp_k_head_num), (kv.v[i], attn.tp_v_head_num)) + ) + workspace = get_buffer( + "trtllm_mha_zero_workspace", + lambda: torch.zeros( + DEFAULT_WORKSPACE_SIZE_MB * 1024 * 1024, dtype=torch.uint8, device=q.device + ), + ) + o = cp.attend( + q.reshape(-1, attn.tp_q_head_num, attn.head_dim), + caches, + meta.page_table, + meta.cache_seqlens, + q_len, + float(1.0 * kv.k_scale[i] * attn.scaling), + float(kv.v_scale[i]), + workspace, + attn, + forward_batch, + min_bs=1, + ) + return None if o is None else o.view(-1, attn.tp_q_head_num * attn.head_dim) + + def skv2(attn, q, kv, meta): """SKV2 on one request's 7 FP8 query rows ``q`` of ``attn`` over the pool ``kv`` (KvPoolBuffers) at ``meta``'s pages and lengths (AttnStepMetadata), whose rows for diff --git a/sgl_nvidia_mega/qwen3_5/draft/__init__.py b/sgl_nvidia_mega/qwen3_5/draft/__init__.py index 00415257..d9461613 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/draft/__init__.py @@ -144,15 +144,18 @@ def __init__( # (2.2 routes the draft steps' q1 decode by KV length) self.graph_variants = None if "draft_loop" in slots and "draft_q1" in slots: + from sgl_nvidia_mega.qwen3_5.attention import route from sgl_nvidia_mega.qwen3_5.draft.attention import Q1Variants - self.graph_variants = {"draft_loop": Q1Variants()} + self.graph_variants = { + "draft_loop": Q1Variants(cp_min_length=route.cp_min_length("draft")) + } if "draft_extend" in slots and "draft_skv2" in slots: from sgl_nvidia_mega.qwen3_5.attention import route self.graph_variants = { **(self.graph_variants or {}), - "draft_extend": route.EXTEND, + "draft_extend": route.extend_variants(), } def summary(self) -> str: diff --git a/sgl_nvidia_mega/qwen3_5/draft/attention.py b/sgl_nvidia_mega/qwen3_5/draft/attention.py index a1d7da2f..165c23ac 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/attention.py +++ b/sgl_nvidia_mega/qwen3_5/draft/attention.py @@ -29,6 +29,9 @@ # SGLANG_Q35MF_MKA_Q1_MIN_LENGTH default, which its TP8 serving configuration keeps) Q1_STOCK, Q1_LONG = "q1_stock", "q1_long" Q1_MIN_LENGTH = 16384 +# slot draft_q1_cp: one request past attention.route.CP_MIN_KV["draft"] runs its steps' +# decode as CP-lite (SGLang's read-split attention over the KV group) +Q1_CP = "q1_cp" # A draft step's q1 decode and its merge launch as programmatic dependents, the # one-row front releasing them at entry (same arithmetic) Q1_PDL = True @@ -43,16 +46,20 @@ class Q1Variants: """The draft loop's graph variants (SGLang's AttentionGraphVariants), as 2.2's - Q1GraphVariants: SKV2-Q1 from a host-only KV length lower bound.""" + Q1GraphVariants: SKV2-Q1 from a host-only KV length lower bound, and with + ``cp_min_length`` (slot draft_q1_cp) CP-lite from that length.""" - # stock first, as 2.2 captures them (variants share a graph memory pool) - capture_labels = (Q1_STOCK, Q1_LONG) # select picks Q1_STOCK past one request (q1 decodes one row), so SGLang captures # one graph at every larger batch size max_batch_size = 1 - def __init__(self, min_length: int = Q1_MIN_LENGTH): + def __init__(self, min_length: int = Q1_MIN_LENGTH, cp_min_length=None): self.min_length = min_length + self.cp_min_length = cp_min_length + # stock first, as 2.2 captures them (variants share a graph memory pool) + self.capture_labels = (Q1_STOCK, Q1_LONG) + ( + (Q1_CP,) if cp_min_length is not None else () + ) def select(self, forward_batch) -> str: if forward_batch.batch_size != 1: @@ -66,7 +73,11 @@ def select(self, forward_batch) -> str: and lengths.numel() == 1 ): length = int(lengths[0]) - return Q1_LONG if length is not None and length >= self.min_length else Q1_STOCK + if length is None: + return Q1_STOCK + if self.cp_min_length is not None and length >= self.cp_min_length: + return Q1_CP + return Q1_LONG if length >= self.min_length else Q1_STOCK # The SGLang names this module reads (test/unit/test_contract.py pins them) @@ -188,6 +199,9 @@ def __init__(self, parts): self.skv2 = ( DraftExtendSplitKV.build(parts) if "draft_skv2" in parts.slots else None ) + # slots draft_q1_cp / draft_skv2_cp: one request's steps / extend in their + # CP graph variant run CP-lite at one request (DraftCP) + self.cp_b1 = DraftCP.build(parts) def __call__( self, x, positions, cache_loc, meta, kv, *, rows_per_request: int, o_proj=None @@ -302,6 +316,14 @@ def decode(self, q8, meta, kv, rows_per_request: int): attn = self.layer.attn bmm1_scale = 1.0 * kv.k_scale[0] * attn.scaling bmm2_scale = kv.v_scale[0] + if ( + self.cp_b1 is not None + and q8.shape[0] == rows_per_request + and self.cp_b1.takes(rows_per_request) + ): + o = self.cp_b1.run(self, q8, meta, kv, rows_per_request) + if o is not None: + return o # q1 decodes one request if ( self.q1 is not None @@ -397,6 +419,60 @@ def run(self, attn, q8, kv, meta): return self.split_kv.skv2(attn, q8, kv, meta) +class DraftCP: + """Slots draft_q1_cp / draft_skv2_cp: in the draft loop's ``Q1_CP`` and the draft + extend's ``SPLIT_KV_CP`` graph variants (one request past its measured crossover), + the decode of the request's rows as CP-lite, which then serves one request + (SGLang's q35_cp_lite through ``attention.split_kv.cp_core``'s call).""" + + def __init__(self, split_kv, route, steps: bool, extend: bool): + self.split_kv, self.route = split_kv, route + self.steps, self.extend = steps, extend + self._logged: set = set() + + @classmethod + def build(cls, parts): + steps = "draft_q1_cp" in parts.slots + extend = "draft_skv2_cp" in parts.slots + if not (steps or extend): + return None + from sgl_nvidia_mega.qwen3_5.attention import route, split_kv + + why = split_kv.cp_reason() + if why: + raise RuntimeError( + f"NVIDIA megakernel: slots draft_q1_cp / draft_skv2_cp cannot run: {why}" + ) + return cls(split_kv, route, steps, extend) + + def takes(self, rows_per_request: int) -> bool: + """The graph being captured is this decode's CP variant.""" + label = _q1_label() + if rows_per_request == 1: + return self.steps and label == Q1_CP + return self.extend and label == self.route.SPLIT_KV_CP + + def run(self, attention, q8, meta, kv, rows_per_request: int): + layer = attention.layer + attn = layer.attn + o = self.split_kv.cp_core(attn, q8, kv, meta, attention, rows_per_request) + if o is not None and torch.cuda.is_current_stream_capturing(): + key = (rows_per_request, q8.shape[0]) + if key not in self._logged: + self._logged.add(key) + logger.info( + "NVIDIA megakernel: Qwen3.5 MTP draft slot %s runs the %s's %d " + "rows as CP-lite (from %d tokens of KV)", + "draft_q1_cp" if rows_per_request == 1 else "draft_skv2_cp", + "step" if rows_per_request == 1 else "extend", + q8.shape[0], + self.route.cp_min_length( + "draft" if rows_per_request == 1 else "extend" + ), + ) + return o + + def _q1_label(): """The KV-length graph variant SGLang is capturing ("q1_stock" / "q1_long"), or None where its draft runner captures one graph per batch size (then Q1 serves diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index 154cf1ff..e2dc0eee 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -17,8 +17,9 @@ expert's gate/up in route's launch) up to 8 rows. ``skv2`` runs the full attention's core of one request through the split-KV attention of SGLang branch nvidia/qwen35-agentx-gb300, ``mka`` the fused attention prologue and epilogue around it -(up to 8 rows), and ``mka_push`` (with ``hop``) has the epilogue push its rows into the -hop. +(up to 8 rows; with four query heads a rank, its wide entries up to 14 rows and past +them its front path, bit for bit SGLang's, at every width), and ``mka_push`` (with +``hop``) has the epilogue push its rows into the hop. """ from __future__ import annotations @@ -47,6 +48,8 @@ # the rows slot route_shared's shared stage fits shared memory for ROUTE_SHARED_ROWS = 8 MKA_MAX_ROWS = 8 +# with four query heads a rank: the wide entries, then the front path (attention/mka.py) +MKA_WIDE_MAX_ROWS = MAX_ROWS class Slots(msgspec.Struct, frozen=True): @@ -87,6 +90,7 @@ class Facts(msgspec.Struct, frozen=True): ) skv2: bool = False # SKV2 for every full-attention layer mka: bool = False # the fused prologue and epilogue for every full-attention layer + mka_wide: bool = False # with its wide entries (four query heads a rank) mka_push: bool = False # with the epilogue's push into the hop @@ -95,7 +99,7 @@ def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> raise ValueError(f"{rows} rows are outside the covered widths 1-{MAX_ROWS}") route = facts.route and rows <= ROUTE_MAX_ROWS hop = facts.hop and rows <= HOP_MAX_ROWS - mka = facts.mka and rows <= MKA_MAX_ROWS + mka = facts.mka and rows <= (MKA_WIDE_MAX_ROWS if facts.mka_wide else MKA_MAX_ROWS) return Slots( route=route, mkm=route diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 52ec6ca0..1e5fddbb 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -31,10 +31,14 @@ "mka", "mka_push", ) -# Steps of more than one request (TP8's decode graphs of batch sizes 2 and 4), covered only +# Steps of more than one request (TP8's decode graphs of batch sizes 2 to 128), covered only # when requested in the same variable: covering them costs the one-request step about # 0.3% at TP8 OPT_IN_COVERAGE: tuple[str, ...] = ("multi_request",) +# Graph variants of a built slot, requested in the same variable: no width of their own. +# skv2_cp: one long request past the measured crossover runs CP-lite in place of SKV2, +# in a graph variant of its own (attention/route.py) +OPT_IN_VARIANTS: tuple[str, ...] = ("skv2_cp",) # The MTP draft's EAGLE entries and kernel groups (qwen3_5/draft), requested in the same # variable; the draft spec's build reads them, the target's ignores them DRAFT_ENTRIES: tuple[str, ...] = ("draft_loop", "draft_extend") @@ -48,6 +52,9 @@ # draft_loop and draft_extend also serve TP8's draft graphs of 2 to 128 requests, # and those of two or more take the sharded top-1 (with multi_request) "draft_multi", + # One long request's steps / extend use CP-lite in graph variants of their own. + "draft_q1_cp", + "draft_skv2_cp", ) SLOTS_ENV = "SGLANG_NVIDIA_MEGAKERNEL_SLOTS" # These slots run all 512 experts themselves. Under expert parallelism, keep the @@ -85,6 +92,7 @@ def slot_reasons() -> list[str]: *SHIPPED_SLOTS, *OPT_IN_SLOTS, *OPT_IN_COVERAGE, + *OPT_IN_VARIANTS, *DRAFT_ENTRIES, *DRAFT_SLOTS, ) @@ -124,7 +132,7 @@ def slot_reasons() -> list[str]: reasons.append( f"{SLOTS_ENV}: slot 'gdn_pdl' launches slot 'gdn_verify''s verify" ) - for name in sorted(SLOTS & set(DRAFT_SLOTS) - {"draft_skv2"}): + for name in sorted(SLOTS & set(DRAFT_SLOTS) - {"draft_skv2", "draft_skv2_cp"}): if "draft_loop" not in SLOTS: reasons.append(f"{SLOTS_ENV}: slot {name!r} runs inside entry 'draft_loop'") if "draft_hop" in SLOTS and "hop" not in SLOTS: @@ -138,6 +146,16 @@ def slot_reasons() -> list[str]: reasons.append( f"{SLOTS_ENV}: slot 'draft_skv2' runs inside entry 'draft_extend'" ) + for name, base in ( + ("skv2_cp", "skv2"), + ("draft_q1_cp", "draft_q1"), + ("draft_skv2_cp", "draft_skv2"), + ): + if name in SLOTS and base not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot {name!r} takes slot {base!r}'s long-KV graph " + "variant's place from its length up" + ) if "mka_push" in SLOTS and not {"mka", "hop"} <= SLOTS: reasons.append( f"{SLOTS_ENV}: slot 'mka_push' pushes slot 'mka''s epilogue into slot " @@ -217,6 +235,7 @@ def build(model, ctx): route_shared=bound["moe"].shared, skv2=bound["attention"].skv2, mka=bound["attention"].mka, + mka_wide=bound["attention"].mka_wide, mka_push=bound["attention"].mka_push, ) covered = slot_plan.covered(pctx.widths, facts, pctx.tokens_per_request) @@ -229,7 +248,13 @@ def build(model, ctx): ) built = msgspec.structs.asdict(facts) # every package's build agrees across ranks, so this raises on every rank alike - target_slots = SLOTS - set(DRAFT_ENTRIES) - set(DRAFT_SLOTS) - set(OPT_IN_COVERAGE) + target_slots = ( + SLOTS + - set(DRAFT_ENTRIES) + - set(DRAFT_SLOTS) + - set(OPT_IN_COVERAGE) + - set(OPT_IN_VARIANTS) + ) # mkm's 16-token engine has nothing to serve where no covered step is wider than 7 rows vacuous = ( {"mkm_wide"} diff --git a/sgl_nvidia_mega/qwen3_5/runtime/executor.py b/sgl_nvidia_mega/qwen3_5/runtime/executor.py index 9045a715..4a77a28d 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/executor.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/executor.py @@ -122,7 +122,7 @@ def __init__( # the attention graph variants per SGLang runner ("target_verify"), which it # captures once it takes a megakernel's (attention/route.py) self.graph_variants = ( - {"target_verify": attention.route.VERIFY} + {"target_verify": attention.route.verify_variants()} if any(s.skv2 for s in slots.values()) else {} ) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/gate_quant.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/gate_quant.py new file mode 100644 index 00000000..d4d863c6 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/gate_quant.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Adapted from SGLang (Apache-2.0, SGLang Team): +# python/sglang/kernels/ops/elementwise/elementwise.py (_fused_sigmoid_mul_kernel) and +# python/sglang/kernels/ops/quantization/fp8_kernel.py (_static_quant_fp8). +"""The attention output gate and ``o_proj``'s static FP8 input quantization in one Triton +launch: the bf16 ``attn * sigmoid(gate)`` of SGLang's ``fused_sigmoid_mul``, then +``static_quant_fp8``'s ``clamp(y * (1 / scale))``, bit for bit the two launches.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _gate_quant_kernel( + attn_ptr, + gate_ptr, + scale_ptr, + out_ptr, + hidden: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = cols < hidden + tl.extra.cuda.gdc_wait() + attn = tl.load(attn_ptr + row * hidden + cols, mask=mask, other=0.0).to(tl.float32) + g = tl.load(gate_ptr + row * hidden + cols, mask=mask, other=0.0).to(tl.float32) + y = (attn * tl.sigmoid(g)).to(tl.bfloat16).to(tl.float32) + s_inv = 1.0 / tl.load(scale_ptr).to(tl.float32) + tl.extra.cuda.gdc_launch_dependents() + q = tl.clamp(y * s_inv, -448.0, 448.0).to(tl.float8e4nv) + tl.store(out_ptr + row * hidden + cols, q, mask=mask) + + +def gate_quant( + attn: torch.Tensor, gate: torch.Tensor, scale: torch.Tensor +) -> torch.Tensor: + """FP8 [rows, hidden] of the gated attention rows ``attn`` [rows, hidden] (bf16, gate + the same shape) at the static per-tensor ``scale``.""" + rows, hidden = attn.shape + out = torch.empty_like(attn, dtype=torch.float8_e4m3fn) + block = 1024 + _gate_quant_kernel[(rows, triton.cdiv(hidden, block))]( + attn, gate, scale, out, hidden, BLOCK=block, num_warps=4, launch_pdl=True + ) + return out diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/mka.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/mka.py index e98af0d7..81e514d9 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/attention/mka.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/mka.py @@ -4,7 +4,9 @@ # python/sglang/kernels/ops/attention/mka_fused.py and mka_fused.cu (its host side). """MKA (``src/mka.cu``): the fused full-attention prologue and epilogue of one request's verify rows (at most 8) at hidden 4096 and head_dim 256, the kernels of SGLang branch -nvidia/qwen35-agentx-gb300 built with NVRTC. +nvidia/qwen35-agentx-gb300 built with NVRTC; and their wide entries for the rows of +several requests (9 to ``WIDE_MAX_ROWS``, four query heads a rank), which run the same +math on 8-row tiles in one launch, bit for bit the narrow kernels on each 8-row chunk. ``prologue``: the normed hidden rows' static FP8 quantization, the qkv (+ gate) GEMM, the q / k Gemma RMSNorm, the partial NeoX M-RoPE, the FP8 query rows, the FP8 K / V rows @@ -29,15 +31,22 @@ SOURCE = Path(__file__).parent / "src" / "mka.cu" THREADS, CLUSTER, HEAD_DIM, HIDDEN, MAX_ROWS = 256, 8, 256, 4096, 8 -ENTRIES = ( - "q35mka_prologue", - "q35mka_prologue_late", - "q35mka_gate", - "q35mka_epilogue_q4", - "q35mka_epilogue_q4_late", - "q35mka_epilogue_q8", - "q35mka_epilogue_q8_late", -) +WIDE_MAX_ROWS = 112 # src/mka.cu kWideTiles 8-row tiles +_INBOX = WIDE_MAX_ROWS // 8 * CLUSTER * 256 * 4 +# each entry's dynamic shared memory: the wide entries' inbox and operand tiles +ENTRIES = { + "q35mka_prologue": 0, + "q35mka_prologue_late": 0, + "q35mka_gate": 0, + "q35mka_epilogue_q4": 0, + "q35mka_epilogue_q4_late": 0, + "q35mka_epilogue_q8": 0, + "q35mka_epilogue_q8_late": 0, + "q35mka_prologue_wide": _INBOX + WIDE_MAX_ROWS // 8 * 16 * 32 * 8, + "q35mka_prologue_wide_late": _INBOX + WIDE_MAX_ROWS // 8 * 16 * 32 * 8, + "q35mka_epilogue_wide_q4": _INBOX + WIDE_MAX_ROWS // 8 * 4 * 32 * 8, + "q35mka_epilogue_wide_q4_late": _INBOX + WIDE_MAX_ROWS // 8 * 4 * 32 * 8, +} class ProParams(ctypes.Structure): @@ -89,7 +98,7 @@ def _kernels(device_index: int) -> dict[str, jit.Kernel]: ("-std=c++17", "-default-device"), headers=hop.headers(), ) - return {n: jit.load(cubin, n, (THREADS, 1, 1), 0) for n in ENTRIES} + return {n: jit.load(cubin, n, (THREADS, 1, 1), b) for n, b in ENTRIES.items()} def build(device) -> None: @@ -173,12 +182,17 @@ def prologue( ) -> None: """The branch's ``prologue`` (its host function's checks and launches).""" w_main, w_gate = w_packed + wide = x.shape[0] > MAX_ROWS if ( - not (x.dim() == 2 and x.shape[0] <= MAX_ROWS and x.shape[1] == HIDDEN) + not (x.dim() == 2 and x.shape[0] <= WIDE_MAX_ROWS and x.shape[1] == HIDDEN) or not x.is_contiguous() ): raise ValueError( - f"MKA prologue: x {tuple(x.shape)} is not [<= 8, 4096] contiguous" + f"MKA prologue: x {tuple(x.shape)} is not [<= {WIDE_MAX_ROWS}, 4096] contiguous" + ) + if wide and w_gate.numel(): + raise ValueError( + f"MKA prologue: {x.shape[0]} rows need four query heads a rank" ) if positions.dim() != 2 or positions.shape[0] != 3 or positions.stride(1) != 1: raise ValueError( @@ -220,7 +234,11 @@ def prologue( wg=w_gate.data_ptr() if split else None, ) if not split: - name = "q35mka_prologue" if early_trigger else "q35mka_prologue_late" + name = ( + "q35mka_prologue" + + ("_wide" if wide else "") + + ("" if early_trigger else "_late") + ) _launch(name, (2 * nq + 2) * CLUSTER, p, x.device) return # The gate launch must start while the prologue runs, so the prologue triggers right after its wait. @@ -242,11 +260,12 @@ def epilogue( """The branch's ``epilogue``; ``site``: the hop mailbox to push the rows into.""" if not (attn.is_contiguous() and gate.is_contiguous() and out.is_contiguous()): raise ValueError("MKA epilogue: attn, gate and out must be contiguous") - if attn.shape[0] > MAX_ROWS or attn.shape[1] != nq * HEAD_DIM: + m = int(attn.shape[0]) + wide = m > MAX_ROWS + if m > (WIDE_MAX_ROWS if nq == 4 else MAX_ROWS) or attn.shape[1] != nq * HEAD_DIM: raise ValueError(f"MKA epilogue: attn {tuple(attn.shape)} for {nq} heads") if nq not in (4, 8): raise ValueError("MKA epilogue: nq must be 4 (TP8) or 8 (TP4)") - m = int(attn.shape[0]) if site is not None and m > site.m_cap: raise ValueError(f"MKA epilogue: {m} rows past the hop's {site.m_cap}") p = EpiParams( @@ -260,5 +279,7 @@ def epilogue( push=int(site is not None), site=site if site is not None else Site(), ) - name = f"q35mka_epilogue_q{nq}" + ("" if early_trigger else "_late") + name = f"q35mka_epilogue{'_wide' if wide else ''}_q{nq}" + ( + "" if early_trigger else "_late" + ) _launch(name, 16 * CLUSTER, p, attn.device) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu b/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu index 5531d774..95eb922a 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu @@ -144,6 +144,27 @@ __device__ __forceinline__ uint32_t gate_quant4(const __nv_bfloat16* o, const __ << 16); } +// The same roundings on operands already in registers (the wide kernels batch their loads). +__device__ __forceinline__ uint32_t quant4_static_r(const uint2 raw, float inv) { + const __nv_bfloat162 v01 = *reinterpret_cast(&raw.x); + const __nv_bfloat162 v23 = *reinterpret_cast(&raw.y); + float f[4] = {__low2float(v01), __high2float(v01), __low2float(v23), __high2float(v23)}; +#pragma unroll + for (int i = 0; i < 4; ++i) f[i] = fminf(fmaxf(__fmul_rn(f[i], inv), -448.f), 448.f); + return e4m3x2(f[0], f[1]) | (e4m3x2(f[2], f[3]) << 16); +} +__device__ __forceinline__ uint32_t gate_quant4_r(const uint2 ro, const uint2 rg, float inv) { + const __nv_bfloat162 o01 = *reinterpret_cast(&ro.x); + const __nv_bfloat162 o23 = *reinterpret_cast(&ro.y); + const __nv_bfloat162 g01 = *reinterpret_cast(&rg.x); + const __nv_bfloat162 g23 = *reinterpret_cast(&rg.y); + return e4m3x2(gate_quant1(__low2float(o01), __low2float(g01), inv), + gate_quant1(__high2float(o01), __high2float(g01), inv)) | + (e4m3x2(gate_quant1(__low2float(o23), __low2float(g23), inv), + gate_quant1(__high2float(o23), __high2float(g23), inv)) + << 16); +} + // Inbox slot s of a 256-row tile -> row: warp w = s/32 holds rows 32w + {g, g+8, g+16, g+24}. __device__ __forceinline__ int slot_row(int s) { return 32 * (s >> 5) + ((s >> 2) & 7) + 8 * (s & 3); } @@ -512,6 +533,301 @@ __device__ __forceinline__ void mka_epilogue_kernel(const EpiParams p) { if (!kEarlyTrigger) pdl_trigger(); } + +// Wide steps (9-112 rows, several requests' verify rows; TP8, no split gate launch): the same per-CTA math as +// the kernels above on ceil(M/8) token tiles in one epoch. Weights load once; every tile's MMA reuses them, and +// each tile's partial rows go straight to their owner (token t to rank t mod 8, inbox slot t / 8) with no +// barrier between tiles. Tile t's arithmetic equals the narrow kernel's on rows [8t, 8t + 8), bit for bit. +constexpr int kWideTiles = 14; // 112 rows +constexpr int kInboxBytes = kCluster * 256 * 4; // one owned token's inbox + +template +__device__ __forceinline__ void mka_prologue_wide_kernel(const ProParams p) { + extern __shared__ __align__(16) unsigned char wsm[]; + float* inbox = reinterpret_cast(wsm); // [tile][rank][256] + uint2* bfr = reinterpret_cast(wsm + kWideTiles * kInboxBytes); // [tile][16][32] + __shared__ float red8[8]; + __shared__ float nrm[256]; + __shared__ float tcs[kWideTiles][2 * kHalfRot]; + __shared__ int64_t tslot[kWideTiles]; + __shared__ __align__(8) uint64_t ibar; + const int rank = (int)cluster_ctarank(); + const int nq = p.nq; + const int unit = blockIdx.x / kCluster; + const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + const int g = lane >> 2, t4 = lane & 3; + if (tid == 0) { + mbar_init(&ibar, 1); + mbar_init_fence(); + } + __syncthreads(); + cluster_arrive_relaxed(); + + uint4 a[16][2]; + { + const uint4* wb = p.w + (size_t)(unit * kCluster + rank) * (8 * 16 * 2 * 32) + warp * 1024 + lane; +#pragma unroll + for (int s = 0; s < 16; ++s) +#pragma unroll + for (int b = 0; b < 2; ++b) a[s][b] = ldg_nc_v4(wb + (s * 2 + b) * 32); + } + const bool is_q = unit < 2 * nq && (unit & 1) == 0; + const bool is_gate = unit < 2 * nq && (unit & 1) == 1; + const bool is_k = unit == 2 * nq; + const bool needs_norm = is_q || is_k; + const int f = slot_row(tid); + const bool rot = needs_norm && f < 2 * kHalfRot; + const int rl = f & (kHalfRot - 1); + const float sx = __ldg(p.sx), sw = __ldg(p.sw); + __nv_bfloat16 nwr = __float2bfloat16_rn(0.f); + if (needs_norm) nwr = (is_q ? p.qnw : p.knw)[f]; + float kvs = 1.f; + if (is_k) kvs = __ldg(p.k_scale); + if (!needs_norm && !is_gate) kvs = __ldg(p.v_scale); + + pdl_wait(); + if (kEarlyTrigger) pdl_trigger(); + const float inv_sx = div_full(1.0f, sx); + const float alpha = __fmul_rn(sx, sw); + const int m = p.m; + const int tiles = (m + 7) >> 3; + const int own = rank < m ? (m - rank + 7) >> 3 : 0; // tokens rank, rank + 8, ... + // the owned tokens' cache slots and cos / sin, all in flight at once + if (tid < own) { + const int tok = 8 * tid + rank; + tslot[tid] = p.loc_i32 ? (int64_t)__ldg(static_cast(p.loc) + tok) + : __ldg(static_cast(p.loc) + tok); + } + if (needs_norm) { + for (int i = tid; i < own * kHalfRot; i += kThreads) { + const int o = i / kHalfRot, l = i % kHalfRot; + const int64_t ps = __ldg(p.pos + __ldg(p.axis + l) * p.pos_stride + 8 * o + rank); + tcs[o][l] = __ldg(p.cos_sin + ps * p.cs_stride + l); + tcs[o][kHalfRot + l] = __ldg(p.cos_sin + ps * p.cs_stride + kHalfRot + l); + } + } + // x's tiles in batches of 4, every load of a batch in flight before its first use + for (int t0 = 0; t0 < tiles; t0 += 4) { + uint2 raw[4][2][2]; +#pragma unroll + for (int u = 0; u < 4; ++u) { + const bool tok_ok = 8 * (t0 + u) + g < m; + const __nv_bfloat16* xr = p.x + (size_t)(8 * (t0 + u) + g) * 4096 + rank * 512 + 4 * t4; +#pragma unroll + for (int i = 0; i < 2; ++i) +#pragma unroll + for (int hh = 0; hh < 2; ++hh) + raw[u][i][hh] = tok_ok ? *reinterpret_cast(xr + (2 * warp + i) * 32 + hh * 16) + : make_uint2(0u, 0u); + } +#pragma unroll + for (int u = 0; u < 4; ++u) { + if (t0 + u >= tiles) break; + const bool tok_ok = 8 * (t0 + u) + g < m; +#pragma unroll + for (int i = 0; i < 2; ++i) + bfr[((t0 + u) * 16 + 2 * warp + i) * 32 + lane] = + tok_ok ? make_uint2(quant4_static_r(raw[u][i][0], inv_sx), quant4_static_r(raw[u][i][1], inv_sx)) + : make_uint2(0u, 0u); + } + } + __syncthreads(); + if (tid == 0 && own) mbar_expect_tx(&ibar, own * kInboxBytes); + cluster_wait(); + const uint32_t bar = smem_u32(&ibar); + for (int t = 0; t < tiles; ++t) { + float ch[2][4][4]; +#pragma unroll + for (int b = 0; b < 2; ++b) +#pragma unroll + for (int c = 0; c < 4; ++c) +#pragma unroll + for (int i = 0; i < 4; ++i) ch[b][c][i] = 0.f; +#pragma unroll + for (int s = 0; s < 16; ++s) { + const uint2 bb = bfr[(t * 16 + s) * 32 + lane]; + mma_e4m3(ch[0][s & 3], a[s][0], bb.x, bb.y); + mma_e4m3(ch[1][s & 3], a[s][1], bb.x, bb.y); + } + float acc[2][4]; +#pragma unroll + for (int b = 0; b < 2; ++b) +#pragma unroll + for (int i = 0; i < 4; ++i) + acc[b][i] = __fadd_rn(__fadd_rn(ch[b][0][i], ch[b][1][i]), __fadd_rn(ch[b][2][i], ch[b][3][i])); + const uint32_t src = smem_u32(&inbox[(t * kCluster + rank) * 256 + warp * 32 + g * 4]); +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int tok = 2 * t4 + j; + if (8 * t + tok < m) + st_async_v4(mapa(src, tok), acc[0][j], acc[0][2 + j], acc[1][j], acc[1][2 + j], mapa(bar, tok)); + } + } + if (!own) { + if (!kEarlyTrigger) pdl_trigger(); + return; + } + mbar_wait(&ibar, 0); // also keeps this CTA alive until every peer's push has landed + + for (int o = 0; o < own; ++o) { + const int tok = 8 * o + rank; + float v = 0.f; +#pragma unroll + for (int r = 0; r < kCluster; ++r) v = __fadd_rn(v, inbox[(o * kCluster + r) * 256 + tid]); + v = bf16_round(__fmul_rn(v, alpha)); + const int64_t slot = tslot[o]; + if (is_gate) { + p.gate_out[(size_t)tok * nq * kHeadDim + (unit >> 1) * kHeadDim + f] = __float2bfloat16_rn(v); + } else if (!needs_norm) { + p.v_cache[(size_t)slot * kHeadDim + f] = e4m3(__fmul_rn(__fdiv_rn(1.0f, kvs), v)); + } else { + float ss = __fmul_rn(v, v); +#pragma unroll + for (int sh = 16; sh > 0; sh >>= 1) ss = __fadd_rn(ss, __shfl_xor_sync(0xffffffffu, ss, sh)); + if (lane == 0) red8[warp] = ss; + __syncthreads(); + float tot = 0.f; +#pragma unroll + for (int i = 0; i < 8; ++i) tot = __fadd_rn(tot, red8[i]); + const float inv = rsqrt_approx_ftz(__fadd_rn(div_full(tot, 256.0f), p.eps)); + float n = bf16_round(__fmul_rn(__fadd_rn(__bfloat162float(nwr), 1.0f), __fmul_rn(inv, v))); + nrm[f] = n; + __syncthreads(); + if (rot) { + const float cs_c = tcs[o][rl], cs_s = tcs[o][kHalfRot + rl]; + const float x1 = f < kHalfRot ? n : nrm[f - kHalfRot]; + const float x2 = f < kHalfRot ? nrm[f + kHalfRot] : n; + n = f < kHalfRot ? bf16_round(__fmaf_rn(cs_c, x1, -__fmul_rn(cs_s, x2))) + : bf16_round(__fmaf_rn(x1, cs_s, __fmul_rn(cs_c, x2))); + } + if (is_k) + p.k_cache[(size_t)slot * kHeadDim + f] = e4m3(__fmul_rn(__fdiv_rn(1.0f, kvs), n)); + else + p.q_out[(size_t)tok * nq * kHeadDim + (unit >> 1) * kHeadDim + f] = e4m3(n); + __syncthreads(); // red8 / nrm are reused by the next token + } + } + if (!kEarlyTrigger) pdl_trigger(); +} + +template +__device__ __forceinline__ void mka_epilogue_wide_kernel(const EpiParams p) { + constexpr int kSteps = NQ; + constexpr int kK = NQ * kHeadDim; + constexpr int kWarpU4 = kSteps * 2 * 32; + extern __shared__ __align__(16) unsigned char wsm[]; + float* inbox = reinterpret_cast(wsm); // [tile][rank][256] + uint2* bfr = reinterpret_cast(wsm + kWideTiles * kInboxBytes); // [tile][kSteps][32] + __shared__ __align__(16) __nv_bfloat16 ob[256]; + __shared__ __align__(8) uint64_t ibar; + const int rank = (int)cluster_ctarank(); + const int unit = blockIdx.x / kCluster; + const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + const int g = lane >> 2, t4 = lane & 3; + if (tid == 0) { + mbar_init(&ibar, 1); + mbar_init_fence(); + } + __syncthreads(); + cluster_arrive_relaxed(); + + uint4 a[kSteps][2]; + { + const uint4* wb = p.w + (size_t)(unit * kCluster + rank) * (8 * kWarpU4) + warp * kWarpU4 + lane; +#pragma unroll + for (int s = 0; s < kSteps; ++s) +#pragma unroll + for (int b = 0; b < 2; ++b) a[s][b] = ldg_nc_v4(wb + (s * 2 + b) * 32); + } + const float sx = __ldg(p.sx), sw = __ldg(p.sw); + + pdl_wait(); + if (kEarlyTrigger) pdl_trigger(); + const float inv_sx = div_full(1.0f, sx); + const float alpha = __fmul_rn(sx, sw); + const int m = p.m; + const int tiles = (m + 7) >> 3; + const int own = rank < m ? (m - rank + 7) >> 3 : 0; + static_assert(2 * kSteps == 8, "one (step, half) per warp"); + { + // every tile's operands in flight before the first use + const int s = warp >> 1, h = warp & 1; + uint2 ro[kWideTiles], rg[kWideTiles]; +#pragma unroll + for (int t = 0; t < kWideTiles; ++t) { + const bool ok = 8 * t + g < m; + const size_t off = (size_t)(8 * t + g) * kK + rank * (kSteps * 32) + 4 * t4 + s * 32 + h * 16; + ro[t] = ok ? *reinterpret_cast(p.attn + off) : make_uint2(0u, 0u); + rg[t] = ok ? *reinterpret_cast(p.gate + off) : make_uint2(0u, 0u); + } +#pragma unroll + for (int t = 0; t < kWideTiles; ++t) { + if (t >= tiles) break; + uint32_t* bw = reinterpret_cast(bfr + t * kSteps * 32); + bw[(s * 32 + lane) * 2 + h] = 8 * t + g < m ? gate_quant4_r(ro[t], rg[t], inv_sx) : 0u; + } + } + __syncthreads(); + if (tid == 0 && own) mbar_expect_tx(&ibar, own * kInboxBytes); + cluster_wait(); + const uint32_t bar = smem_u32(&ibar); + constexpr int kCh = kSteps < 4 ? kSteps : 4; + for (int t = 0; t < tiles; ++t) { + float ch[2][kCh][4]; +#pragma unroll + for (int b = 0; b < 2; ++b) +#pragma unroll + for (int c = 0; c < kCh; ++c) +#pragma unroll + for (int i = 0; i < 4; ++i) ch[b][c][i] = 0.f; +#pragma unroll + for (int s = 0; s < kSteps; ++s) { + const uint2 bb = bfr[(t * kSteps + s) * 32 + lane]; + mma_e4m3(ch[0][s % kCh], a[s][0], bb.x, bb.y); + mma_e4m3(ch[1][s % kCh], a[s][1], bb.x, bb.y); + } + float acc[2][4]; +#pragma unroll + for (int b = 0; b < 2; ++b) +#pragma unroll + for (int i = 0; i < 4; ++i) { + float x = ch[b][0][i]; +#pragma unroll + for (int c = 1; c < kCh; ++c) x = __fadd_rn(x, ch[b][c][i]); + acc[b][i] = x; + } + const uint32_t src = smem_u32(&inbox[(t * kCluster + rank) * 256 + warp * 32 + g * 4]); +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int tok = 2 * t4 + j; + if (8 * t + tok < m) + st_async_v4(mapa(src, tok), acc[0][j], acc[0][2 + j], acc[1][j], acc[1][2 + j], mapa(bar, tok)); + } + } + if (!own) { + if (!kEarlyTrigger) pdl_trigger(); + return; + } + mbar_wait(&ibar, 0); + const uint32_t stage = p.push ? q35hop::stage_peek(p.site) : 0u; + for (int o = 0; o < own; ++o) { + const int tok = 8 * o + rank; + float v = 0.f; +#pragma unroll + for (int r = 0; r < kCluster; ++r) v = __fadd_rn(v, inbox[(o * kCluster + r) * 256 + tid]); + const __nv_bfloat16 out = __float2bfloat16_rn(__fmul_rn(v, alpha)); + if (p.push) { + ob[slot_row(tid)] = out; + __syncthreads(); + if (tid < 32) + q35hop::push_frag(p.site, stage, tok, unit * 32 + tid, *reinterpret_cast(&ob[8 * tid])); + __syncthreads(); // ob is reused by the next token + } + p.out[(size_t)tok * 4096 + unit * kHeadDim + slot_row(tid)] = out; + } + if (!kEarlyTrigger) pdl_trigger(); +} + } // namespace q35mka // C-linkage entries, one per instantiation the launcher uses (mka.py) @@ -536,3 +852,15 @@ extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_epilogu extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_epilogue_q8_late(const q35mka::EpiParams p) { q35mka::mka_epilogue_kernel<8, false>(p); } +extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_prologue_wide(const q35mka::ProParams p) { + q35mka::mka_prologue_wide_kernel(p); +} +extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_prologue_wide_late(const q35mka::ProParams p) { + q35mka::mka_prologue_wide_kernel(p); +} +extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_epilogue_wide_q4(const q35mka::EpiParams p) { + q35mka::mka_epilogue_wide_kernel<4, true>(p); +} +extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_epilogue_wide_q4_late(const q35mka::EpiParams p) { + q35mka::mka_epilogue_wide_kernel<4, false>(p); +} diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index cc99418e..8d79db35 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -126,6 +126,10 @@ }, # SKV2 against SGLang's TRT-LLM MHA decode (test_skv2_is_within_twice_sglangs_error) "skv2": {"skv2.attention.rel_l2_vs_stock": 2.0}, + # CP-lite's exact path (SGLang's q35_cp_lite) at one request, as skv2's bound + "skv2_cp": {"skv2_cp.attention.rel_l2_vs_stock": 2.0}, + "draft_q1_cp": {"draft_q1_cp.attention.rel_l2_vs_stock": 2.0}, + "draft_skv2_cp": {"draft_skv2_cp.attention.rel_l2_vs_stock": 2.0}, # MKA's own qkv / o_proj GEMMs; the GEMM slots' bounds, no fp64 kernel test yet "mka": { "mka.qkv_proj.rel_l2_vs_stock": 1.05, diff --git a/test/kernels/test_qwen3_5_attention.py b/test/kernels/test_qwen3_5_attention.py index 14aaa4d5..823c6b99 100644 --- a/test/kernels/test_qwen3_5_attention.py +++ b/test/kernels/test_qwen3_5_attention.py @@ -297,6 +297,69 @@ def test_mka_replays_in_a_cuda_graph_bit_for_bit(): assert torch.isfinite(step.out.float()).all() +@pytest.mark.parametrize("m", [9, 14, 28, 56, 112]) +def test_wide_mka_is_the_narrow_mka_on_each_8_row_chunk(m): + """The wide entries (several requests' verify rows in one launch) give the bytes of + the narrow kernels called on rows [8t, 8t + 8), eagerly and replayed in a CUDA + graph on new inputs.""" + from sgl_nvidia_mega.qwen3_5.sm103.attention import mka + + g = torch.Generator(device="cuda").manual_seed(m) + ops = _mka_operands(g) + step = _MkaStep(m, g) + + def chunked(): + k, v = step.k.clone(), step.v.clone() + q, gate, out = (torch.empty_like(t) for t in (step.q, step.gate, step.out)) + for s in range(0, m, 8): + e = min(s + 8, m) + mka.prologue( + x=step.x[s:e], + w_packed=ops["w_qkv"], + sx=ops["qkv_sx"], + sw=ops["qkv_sw"], + q_norm_w=ops["qnw"], + k_norm_w=ops["knw"], + cos_sin=ops["cos_sin"], + positions=step.pos[:, s:e], + axis_map=ops["axis"], + cache_loc=step.loc[s:e], + k_cache=k, + v_cache=v, + k_scale=ops["k_scale"], + v_scale=ops["v_scale"], + q_out=q[s:e], + gate_out=gate[s:e], + nq=HEADS, + eps=1e-6, + ) + mka.epilogue( + attn=step.attn[s:e], + gate=gate[s:e], + w_packed=ops["w_o"], + sx=ops["o_sx"], + sw=ops["o_sw"], + out=out[s:e], + nq=HEADS, + ) + torch.cuda.synchronize() + rows = (k[step.loc], v[step.loc]) + return [t.clone().view(torch.uint8) for t in (q, gate, *rows, out)] + + want = chunked() + step.run(ops) + assert _same(want, step.bits()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step.run(ops) + for _ in range(2): + step.fill(g) + graph.replay() + replayed = step.bits() + assert _same(chunked(), replayed) + assert torch.isfinite(step.out.float()).all() + + def _bits(a, b): """Whether two tensors (of any dtype, FP8 too) hold the same bytes.""" if a is None or b is None: @@ -381,3 +444,59 @@ def run(): assert _bits(a, b) for a, b in zip(eager[2:], pools): assert _bits(a[rows], b[rows]) + + +@pytest.mark.parametrize("rows", [7, 14, 28, 56, 112, 224, 448, 896]) +def test_front_with_the_cache_write_is_sglangs_front_then_fp8_kv_write(rows): + """The front with its FP8 KV write (one launch) gives the bytes of SGLang's front + followed by its fused FP8 q / k / v + KV write, at every covered width.""" + from sglang.kernels.ops.attention.fused_qk_rmsnorm_rope_gate import ( + fused_qk_gemma_rmsnorm_rope_gate as sglang_front, + ) + from sglang.kernels.ops.kvcache.fused_fp8_qkv_kv_cache import ( + fused_fp8_qkv_kv_cache, + ) + + from sgl_nvidia_mega.qwen3_5.sm103.draft.front import ( + fused_qk_gemma_rmsnorm_rope_gate, + ) + + g = torch.Generator(device="cuda").manual_seed(300 + rows) + args, axis, v, pools, loc = _front_operands(rows, g, slots=64 * 32) + scales = [torch.full((1,), s, device="cuda") for s in (0.8, 1.3)] + want_pools = tuple(p.clone() for p in pools) + q, k, gate = sglang_front(*args, has_gate=True, mrope_axis_map=axis) + q8 = fused_fp8_qkv_kv_cache( + q.reshape(rows, -1), k.reshape(rows, -1), v, *want_pools, loc, *scales + ) + got8, _, got_gate = fused_qk_gemma_rmsnorm_rope_gate( + *args, + has_gate=True, + mrope_axis_map=axis, + v=v, + kv_cache=pools, + cache_loc=loc, + k_scale=scales[0], + v_scale=scales[1], + ) + torch.cuda.synchronize() + assert _bits(q8.reshape(rows, -1), got8.reshape(rows, -1)) + assert _bits(gate.reshape(rows, -1), got_gate.reshape(rows, -1)) + rows_written = loc.long() + for a, b in zip(want_pools, pools): + assert _bits(a[rows_written], b[rows_written]) + + +@pytest.mark.parametrize("rows", [7, 14, 28, 56, 112, 224, 448, 896]) +def test_gate_quant_is_sglangs_sigmoid_mul_then_static_quant(rows): + from sglang.kernels.ops.elementwise.elementwise import fused_sigmoid_mul + from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 + + from sgl_nvidia_mega.qwen3_5.sm103.attention.gate_quant import gate_quant + + g = torch.Generator(device="cuda").manual_seed(400 + rows) + attn = (torch.randn(rows, HEADS * DIM, device="cuda", generator=g) * 4).bfloat16() + gate = (torch.randn(rows, HEADS * DIM, device="cuda", generator=g) * 3).bfloat16() + scale = torch.full((1,), 0.013, device="cuda") + want, _ = static_quant_fp8(fused_sigmoid_mul(attn, gate), scale) + assert _bits(want, gate_quant(attn, gate, scale)) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 535b9dad..117fab47 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -140,6 +140,23 @@ def test_the_runtime_gives_the_executor_the_build_contexts_binding(monkeypatch): assert built.binding == "EagleBinding" and built.summary() == "slots none (stock)" +def test_slot_skv2_cp_builds_as_a_graph_variant_of_slot_skv2(monkeypatch): + from sgl_nvidia_mega.qwen3_5.attention import split_kv + + monkeypatch.setattr(runtime, "SLOTS", frozenset({"skv2", "skv2_cp"})) + monkeypatch.setattr(split_kv, "cp_reason", lambda: "") + model = SimpleNamespace( + model=SimpleNamespace(start_layer=0, end_layer=0, layers=[]) + ) + ctx = SimpleNamespace( + widths=(7,), tokens_per_request=7, deadline_s=60.0, binding="EagleBinding" + ) + assert runtime.slot_reasons() == [] + built = runtime.build(model, ctx) + variants = built.graph_variants["target_verify"] + assert variants.capture_labels == ("split_kv_stock", "split_kv_long", "split_kv_cp") + + def _draft_ctx(): return SimpleNamespace( widths=(), tokens_per_request=7, deadline_s=60.0, binding="EagleBinding" @@ -242,6 +259,105 @@ def test_the_draft_extend_takes_the_split_kv_variants_with_slot_draft_skv2( assert built.graph_variants["draft_extend"] is route.EXTEND +def test_the_cp_slots_offer_the_draft_runners_a_cp_variant(monkeypatch): + from sgl_nvidia_mega.qwen3_5 import draft + from sgl_nvidia_mega.qwen3_5.attention import route + + slots = {"draft_loop", "draft_q1", "draft_q1_cp", "draft_extend", "draft_skv2"} + monkeypatch.setattr(runtime, "SLOTS", frozenset(slots | {"draft_skv2_cp"})) + assert runtime.slot_reasons() == [] + built = draft.SPEC.build("lm", SimpleNamespace(context_len=4096), _draft_ctx()) + loop = built.graph_variants["draft_loop"] + extend = built.graph_variants["draft_extend"] + assert loop.capture_labels == ("q1_stock", "q1_long", "q1_cp") + assert extend.capture_labels == ( + "split_kv_stock", + "split_kv_long", + "split_kv_cp", + ) + + def one(n): + return SimpleNamespace( + batch_size=1, attention_kv_len_lower_bound=n, seq_lens_cpu=None + ) + + draft_cross = route.CP_MIN_KV["draft"] + extend_cross = route.CP_MIN_KV["extend"] + assert [ + loop.select(one(n)) for n in (16383, 16384, draft_cross - 1, draft_cross) + ] == [ + "q1_stock", + "q1_long", + "q1_long", + "q1_cp", + ] + assert [ + extend.select(one(n)) for n in (32767, 32768, extend_cross - 1, extend_cross) + ] == ["split_kv_stock", "split_kv_long", "split_kv_long", "split_kv_cp"] + + monkeypatch.setattr( + runtime, "SLOTS", frozenset({"draft_loop", "draft_q1_cp", "skv2_cp"}) + ) + assert sorted(r.split(": ", 1)[1] for r in runtime.slot_reasons()) == [ + "slot 'draft_q1_cp' takes slot 'draft_q1''s long-KV graph variant's place " + "from its length up", + "slot 'skv2_cp' takes slot 'skv2''s long-KV graph variant's place from its " + "length up", + ] + + +@pytest.mark.parametrize("rows, label", [(1, "q1_cp"), (7, "split_kv_cp")]) +def test_the_draft_decode_takes_cp_lite_for_one_request_in_its_cp_variant( + monkeypatch, rows, label +): + from sgl_nvidia_mega.qwen3_5.attention import route, split_kv + from sgl_nvidia_mega.qwen3_5.draft import attention as draft_attention + + heads, dim, page = 4, 256, 64 + calls = [] + + class CpLite: + def attend( + self, q, kv, pages, lens, q_len, bmm1, bmm2, ws, layer, owner, min_bs=None + ): + calls.append((tuple(q.shape), q_len, bmm1, bmm2, min_bs)) + return torch.ones(q.shape[0], heads, dim, dtype=torch.bfloat16) + + monkeypatch.setattr(split_kv, "cp_lite", CpLite) + monkeypatch.setattr(draft_attention, "_q1_label", lambda: label) + monkeypatch.setattr( + "sglang.srt.runtime_context.get_buffer", lambda name, make: None + ) + attn = draft_attention.Attention.__new__(draft_attention.Attention) + attn.layer = SimpleNamespace( + attn=SimpleNamespace( + layer_id=5, + scaling=0.5, + sliding_window_size=-1, + tp_q_head_num=heads, + tp_k_head_num=1, + tp_v_head_num=1, + head_dim=dim, + ) + ) + attn.heads, attn.head_dim, attn.page_size = heads, dim, page + attn.q1 = attn.skv2 = None + attn._noted = set() + attn.cp_b1 = draft_attention.DraftCP(split_kv, route, steps=True, extend=True) + cache = torch.zeros(4 * page, 1, dim, dtype=torch.float8_e4m3fn) + kv = SimpleNamespace( + layer_ids=(5,), k=(cache,), v=(cache,), k_scale=(2.0,), v_scale=(3.0,) + ) + meta = SimpleNamespace( + page_table=torch.zeros(1, 4, dtype=torch.int32), + cache_seqlens=torch.full((1,), 100, dtype=torch.int32), + ) + q8 = torch.zeros(rows, heads * dim, dtype=torch.float8_e4m3fn) + out = attn.decode(q8, meta, kv, rows_per_request=rows) + assert out.shape == (rows, heads * dim) and bool((out == 1).all()) + assert calls == [((rows, heads, dim), rows, 1.0, 3.0, 1)] + + @pytest.mark.parametrize( "slots, want", [ diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 5a8ff30d..9c5e61f0 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -212,12 +212,18 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): **{**dict.fromkeys(plan.Facts.__struct_fields__, True), "mkm_wide": False} ) assert _on(plan.slots(7, every)) == sorted(plan.Slots.__struct_fields__) - # two requests: the 16-row router front, the GDN chain and the hop's one-shot rows - assert _on(plan.slots(14, every)) == ["gdn_pdl", "gdn_verify", "hop", "route"] - # four requests and wider: past the router front's and one-shot all-reduce's rows - assert _on(plan.slots(28, every)) == ["gdn_pdl", "gdn_verify"] - for rows in (56, 112, 224, 448, 896): - assert _on(plan.slots(rows, every)) == ["gdn_pdl", "gdn_verify"] + # two requests: also MKA's wide entries and their push into the hop + assert _on(plan.slots(14, every)) == [ + "gdn_pdl", + "gdn_verify", + "hop", + "mka", + "mka_push", + "route", + ] + # Wider requests take MKA's bitwise front path. + for rows in (28, 56, 112, 224, 448, 896): + assert _on(plan.slots(rows, every)) == ["gdn_pdl", "gdn_verify", "mka"] # SGLang offers every multiple of 7 up to MAX_ROWS; every request count is a candidate offered = tuple(range(7, plan.MAX_ROWS + 1, 7)) assert plan.candidates(offered, 7) == offered @@ -236,6 +242,15 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): attention = plan.Facts(skv2=True, mka=True, mka_push=True, hop=True) assert plan.covered((7, 14, 28), attention, 7) == (7, 14) assert _on(plan.slots(14, attention)) == ["hop"] + wide_attention = plan.Facts(mka=True, mka_wide=True) + assert plan.covered((7, 14, 28, 56, 112, 224), wide_attention, 7) == ( + 7, + 14, + 28, + 56, + 112, + 224, + ) # mkm_wide: mkm's 16-token engine at two requests, not past the router front's rows wide = plan.Facts(route=True, mkm=True, mkm_wide=True) assert _on(plan.slots(14, wide)) == ["mkm", "route"] @@ -258,6 +273,10 @@ def test_the_slot_plans_row_bounds_are_the_kernels(): assert plan.MKM_WIDE_MAX_ROWS == mkm.M16_MAX_TOKENS assert plan.GDN_PUSH_MAX_ROWS == hop.OPROJ_MAX_M assert plan.MKA_MAX_ROWS == mka.MAX_ROWS + from sgl_nvidia_mega.qwen3_5.attention import mka as mka_step + + assert plan.MKA_WIDE_MAX_ROWS == plan.MAX_ROWS + assert mka_step.FUSED_MAX_ROWS <= mka.WIDE_MAX_ROWS def test_a_numeric_slot_turns_the_claim_into_its_tolerances(): diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 1ceb4531..17794b2e 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -1610,6 +1610,95 @@ def batch(bs=1, bound=None, lens=None): assert route.VERIFY.select(batch()) == "split_kv_stock" +def test_slot_skv2_cp_adds_a_cp_variant_from_its_kv_length(monkeypatch): + from sgl_nvidia_mega.qwen3_5 import runtime + from sgl_nvidia_mega.qwen3_5.attention import route + + def batch(bs=1, bound=None): + return SimpleNamespace( + batch_size=bs, attention_kv_len_lower_bound=bound, seq_lens_cpu=None + ) + + monkeypatch.setattr(runtime, "SLOTS", frozenset({"skv2"})) + assert route.verify_variants() is route.VERIFY + monkeypatch.setattr(runtime, "SLOTS", frozenset({"skv2", "skv2_cp"})) + variants = route.verify_variants() + crossover = route.CP_MIN_KV["verify"] + assert variants.capture_labels == ("split_kv_stock", "split_kv_long", "split_kv_cp") + assert (variants.min_length, variants.cp_min_length) == (8192, crossover) + assert variants.max_batch_size == 1 + assert [ + variants.select(batch(bound=n)) for n in (8191, 8192, crossover - 1, crossover) + ] == [ + "split_kv_stock", + "split_kv_long", + "split_kv_long", + "split_kv_cp", + ] + assert variants.select(batch(bs=2, bound=10**6)) == "split_kv_stock" + # slot draft_skv2_cp is not requested + assert route.extend_variants() is route.EXTEND + assert route.CP_MIN_KV == { + "verify": 139264, + "draft": 184320, + "extend": 245760, + } + + +@pytest.mark.parametrize("served", [True, False]) +def test_the_skv2_cp_slot_runs_cp_lite_for_one_request_in_its_variant( + monkeypatch, served +): + from sgl_nvidia_mega.qwen3_5 import runtime + from sgl_nvidia_mega.qwen3_5.attention import split_kv + + monkeypatch.setattr(runtime, "SLOTS", frozenset({"skv2", "skv2_cp"})) + log = _log() + attn, st = _skv2_setup(log) + q8 = torch.zeros(7, 1024, dtype=torch.float8_e4m3fn) + log.known["q8"] = q8 + out = torch.ones(7, 4, 256, dtype=torch.bfloat16) + calls = [] + + class CpLite: + def attend( + self, q, kv, pages, lens, q_len, bmm1, bmm2, ws, layer, fb, min_bs=None + ): + calls.append( + ( + tuple(q.shape), + tuple(kv[0].shape), + kv[0].stride(), + pages is st.view.attn.page_table, + lens is st.view.attn.cache_seqlens, + q_len, + bmm1, + bmm2, + layer is attn, + fb is st.fb, + min_bs, + ) + ) + return out if served else None + + monkeypatch.setattr(split_kv, "cp_lite", CpLite) + with ( + _skv2_kernels(log, label="split_kv_cp"), + patch("sglang.srt.runtime_context.get_buffer", lambda name, make: None), + ): + o = split_kv.core(attn, q8, st) + # the request's 7 rows on the view's pool pages, as SKV2 takes them, at one request + assert calls == [ + ((7, 4, 256), (3, 1, 64, 256), (16384, 16384, 256, 1), True, True, 7) + + (0.75 * 0.0625, 0.5, True, True, 1) + ] + if served: + assert o.shape == (7, 1024) and bool((o == 1).all()) and len(log) == 0 + else: + # CP-lite declines by shape: SKV2 runs as in the long variant + assert log[-1][0] == "skv2" + + @pytest.mark.parametrize( "kw, want", [ @@ -1973,7 +2062,14 @@ def layer(kind): monkeypatch.setattr(mka, "unsupported", lambda l: why.get(id(l), "")) monkeypatch.setattr(mka, "pack", packed.append) ctx = SimpleNamespace(tokens_per_request=7) + for made in layers[1:]: + made.num_heads = 8 + # TP4's eight query heads a rank: the narrow kernels only assert attention.build(model, ctx) == attention.Bound(mka=True) + packed.clear() + for made in layers[1:]: + made.num_heads = 4 + assert attention.build(model, ctx) == attention.Bound(mka=True, mka_wide=True) assert packed == layers[1:] # a layer the slot cannot serve leaves the slot unbuilt, packs nothing more why[id(layers[2])] = "no attention output gate" From 0e41d5fe6310ac6e2f65668f635c10526d2239b3 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 15:14:18 -0700 Subject: [PATCH 04/62] feat: fuse GDN verify cores across request widths --- sgl_nvidia_mega/qwen3_5/gdn/__init__.py | 26 + sgl_nvidia_mega/qwen3_5/gdn/conv.py | 268 ++++++++ sgl_nvidia_mega/qwen3_5/gdn/front.py | 13 + sgl_nvidia_mega/qwen3_5/gdn/verify.py | 123 ++++ sgl_nvidia_mega/qwen3_5/plan.py | 15 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 11 + sgl_nvidia_mega/qwen3_5/runtime/executor.py | 14 + sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py | 280 ++++++++ sgl_nvidia_mega/qwen3_5/sm103/gdn/src/core.cu | 598 ++++++++++++++++++ test/kernels/test_qwen3_5_gdn.py | 220 +++++++ test/unit/test_qwen3_5_spec.py | 21 +- 11 files changed, 1584 insertions(+), 5 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/gdn/src/core.cu diff --git a/sgl_nvidia_mega/qwen3_5/gdn/__init__.py b/sgl_nvidia_mega/qwen3_5/gdn/__init__.py index 05f0ea73..1a653697 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/__init__.py @@ -62,6 +62,10 @@ class Bound(msgspec.Struct, frozen=True): gdn_verify: bool = False # slot gdn_verify serves every GDN layer gdn_pdl: bool = False # and launches its verify dependent on the conv + gdn_core: bool = False # and runs its verify, gated norm and FP8 in one launch + gdn_conv: bool = ( + False # and its conv with the x rows loaded first and one window store + ) # slot gdn_push: layer id -> out_proj's weights packed for the hop's push GEMM packed: dict = msgspec.field(default_factory=dict) # slot gdn_front: layer id -> in_proj_qkvz's weights packed for the front @@ -113,9 +117,25 @@ def build(model, ctx) -> Bound: return Bound(reason=f"gdn_front: layer {i}: {reason}") front_packed = front.build(model, gdns) error = torch.zeros(1, dtype=torch.int32, device="cuda") + core = "gdn_core" in SLOTS + if core: + import torch + + from sgl_nvidia_mega.qwen3_5.sm103.gdn import core as core_kernel + + attn = gdns[0][1].attn + core_kernel.build( + torch.cuda.current_device(), + ctx.tokens_per_request, + attn.num_k_heads, + attn.num_v_heads, + gdns[0][1].norm.weight.dtype == torch.float32, + ) return Bound( gdn_verify=True, gdn_pdl="gdn_pdl" in SLOTS, + gdn_core=core, + gdn_conv="gdn_conv" in SLOTS, packed=packed, front=front_packed, error=error, @@ -130,6 +150,12 @@ def step(layer, x, st): from sgl_nvidia_mega.qwen3_5.gdn import verify opts = verify.SLOT_PDL if st.slots.gdn_pdl else verify.SLOT + if st.slots.gdn_core: + opts = msgspec.structs.replace(opts, core=True) + verify.engaged("gdn_core", x.shape[0]) + if st.slots.gdn_conv: + opts = msgspec.structs.replace(opts, conv_rows=True) + verify.engaged("gdn_conv", x.shape[0]) # push only where the MoE input's all-reduce will consume on the hop if st.slots.gdn_push and hop.takes_moe_input( st.bound["boundary"].hop, layer, st, x.shape[0] diff --git a/sgl_nvidia_mega/qwen3_5/gdn/conv.py b/sgl_nvidia_mega/qwen3_5/gdn/conv.py index b8687c1c..387fe5fc 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/conv.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/conv.py @@ -50,6 +50,9 @@ def conv_update( *, block_n: int = 256, token_major_out: bool = False, + hoist: bool = False, + num_warps: int = 4, + window_mode: int = 0, ) -> torch.Tensor: """``causal_conv1d_update`` of x (batch, dim, seqlen) for a verify step, with the channel tile ``block_n`` and, with ``token_major_out``, an output whose @@ -79,6 +82,49 @@ def conv_update( def grid(meta): return (batch, triton.cdiv(dim, meta["BLOCK_N"])) + if hoist: + if width != 4 or seqlen > 8 or not seqlen > state_len: + raise ValueError( + f"the hoisted conv takes width 4 and {state_len} < rows <= 8, not {width} / {seqlen}" + ) + st = intermediate_conv_window.stride() + if window_mode == 1 and not (st[1] == 1 and st[3] == 1 and st[2] >= seqlen + 2): + raise ValueError( + f"window rows need the deduplicated window layout, not strides {st}" + ) + _HOISTED[0] = _HOISTED[0] or _hoisted_kernel() + _HOISTED[0][grid]( + x, + weight, + bias, + conv_state, + conv_state_indices, + intermediate_conv_window, + intermediate_state_indices, + out, + dim, + seqlen, + state_len, + num_cache_lines, + *x.stride(), + *weight.stride(), + *conv_state.stride(), + conv_state_indices.stride(0), + *intermediate_conv_window.stride(), + intermediate_state_indices.stride(0), + *out.stride(), + PAD_SLOT_ID, + HAS_BIAS=bias is not None, + SILU_ACTIVATION=activation in ["silu", "swish"], + NP2_STATELEN=triton.next_power_of_2(state_len), + BLOCK_N=block_n, + USE_GDC=bool(pdl_kwargs), + WINDOW_MODE=window_mode, + num_warps=num_warps, + **({"launch_pdl": True} if pdl_kwargs else {}), + ) + return out + _causal_conv1d_update_kernel[grid]( x, weight, @@ -123,6 +169,228 @@ def grid(meta): BLOCK_N=block_n, SAVE_INTERMEDIATE=True, HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=False, + num_warps=num_warps, **pdl_kwargs, ) return out + + +_HOISTED = [None] + + +def _hoisted_kernel(): + """SGLang's ``_causal_conv1d_update_kernel`` for this verify case (continuous batching, pad slots, saved + windows, no tree, width 4) with every token's x row loaded before the first store: the stores of a token + no longer serialize the next token's load (the compiler cannot move a load above a possibly aliasing + store). Same values, same arithmetic order: bit for bit SGLang's kernel.""" + import triton + import triton.language as tl + + @triton.jit + def kernel( + x_ptr, + w_ptr, + bias_ptr, + conv_state_ptr, + conv_state_indices_ptr, + inter_ptr, + inter_indices_ptr, + o_ptr, + dim: tl.constexpr, + seqlen: tl.constexpr, + state_len: tl.constexpr, + num_cache_lines: tl.constexpr, + stride_x_seq: tl.constexpr, + stride_x_dim: tl.constexpr, + stride_x_token: tl.constexpr, + stride_w_dim: tl.constexpr, + stride_w_width: tl.constexpr, + stride_conv_state_seq: tl.constexpr, + stride_conv_state_dim: tl.constexpr, + stride_conv_state_tok: tl.constexpr, + stride_state_indices: tl.constexpr, + stride_inter_seq: tl.constexpr, + stride_inter_step: tl.constexpr, + stride_inter_dim: tl.constexpr, + stride_inter_win: tl.constexpr, + stride_inter_indices: tl.constexpr, + stride_o_seq: tl.constexpr, + stride_o_dim: tl.constexpr, + stride_o_token: tl.constexpr, + pad_slot_id: tl.constexpr, + HAS_BIAS: tl.constexpr, + SILU_ACTIVATION: tl.constexpr, + NP2_STATELEN: tl.constexpr, + BLOCK_N: tl.constexpr, + USE_GDC: tl.constexpr, + WINDOW_MODE: tl.constexpr, + ): + if USE_GDC: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + idx_seq = tl.program_id(0) + idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + conv_state_batch_coord = tl.load( + conv_state_indices_ptr + idx_seq * stride_state_indices + ).to(tl.int64) + inter_coord = tl.load(inter_indices_ptr + idx_seq * stride_inter_indices).to( + tl.int64 + ) + if conv_state_batch_coord == pad_slot_id: + return + # STEP 1: the prior window + conv_states_base = ( + conv_state_ptr + + (conv_state_batch_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim) + ) + mask_w = idx_feats < dim + col0 = tl.load(conv_states_base, mask_w, 0.0) + col1 = tl.load(conv_states_base + 1 * stride_conv_state_tok, mask_w, 0.0) + col2 = tl.load(conv_states_base + 2 * stride_conv_state_tok, mask_w, 0.0) + prior1 = col1 + prior2 = col2 + # every token's x row, before any store + x_base = x_ptr + (idx_seq * stride_x_seq) + (idx_feats * stride_x_dim) + x0 = tl.load(x_base + 0 * stride_x_token, mask=mask_w) + if seqlen > 1: + x1 = tl.load(x_base + 1 * stride_x_token, mask=mask_w) + if seqlen > 2: + x2 = tl.load(x_base + 2 * stride_x_token, mask=mask_w) + if seqlen > 3: + x3 = tl.load(x_base + 3 * stride_x_token, mask=mask_w) + if seqlen > 4: + x4 = tl.load(x_base + 4 * stride_x_token, mask=mask_w) + if seqlen > 5: + x5 = tl.load(x_base + 5 * stride_x_token, mask=mask_w) + if seqlen > 6: + x6 = tl.load(x_base + 6 * stride_x_token, mask=mask_w) + if seqlen > 7: + x7 = tl.load(x_base + 7 * stride_x_token, mask=mask_w) + # STEP 2: the new conv state (state_len < seqlen: the last state_len x rows) + idx_tokens = tl.arange(0, NP2_STATELEN) + conv_state_ptrs_source = ( + conv_state_ptr + + (conv_state_batch_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim)[None, :] + + ((idx_tokens + seqlen) * stride_conv_state_tok)[:, None] + ) + mask = ( + (conv_state_batch_coord < num_cache_lines) + & ((idx_tokens + seqlen) < state_len)[:, None] + & (idx_feats < dim)[None, :] + ) + conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0) + VAL = state_len - seqlen + x_ptrs = x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None] + mask_x = ( + (idx_tokens - VAL >= 0)[:, None] + & (idx_tokens - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + tl.debug_barrier() + new_conv_state = tl.where(mask, conv_state, loaded_x) + conv_state_ptrs_target = ( + conv_state_ptr + + (conv_state_batch_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim) + )[None, :] + (idx_tokens * stride_conv_state_tok)[:, None] + tl.store( + conv_state_ptrs_target, + new_conv_state, + (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :], + ) + # STEP 3 / 4: bias and weights + if HAS_BIAS: + acc_preload = tl.load(bias_ptr + idx_feats, mask=mask_w, other=0.0).to( + tl.float32 + ) + else: + acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) + w_base = w_ptr + (idx_feats * stride_w_dim) + w_col0 = tl.load(w_base + (0 * stride_w_width), mask_w, other=0.0) + w_col1 = tl.load(w_base + (1 * stride_w_width), mask_w, other=0.0) + w_col2 = tl.load(w_base + (2 * stride_w_width), mask_w, other=0.0) + w_col3 = tl.load(w_base + (3 * stride_w_width), mask_w, other=0.0) + # STEP 5: each token + for idx_token in tl.static_range(seqlen): + acc = acc_preload + if idx_token == 0: + xt = x0 + elif idx_token == 1: + xt = x1 + elif idx_token == 2: + xt = x2 + elif idx_token == 3: + xt = x3 + elif idx_token == 4: + xt = x4 + elif idx_token == 5: + xt = x5 + elif idx_token == 6: + xt = x6 + else: + xt = x7 + acc += col0 * w_col0 + acc += col1 * w_col1 + acc += col2 * w_col2 + acc += xt * w_col3 + col0 = col1 + col1 = col2 + col2 = xt + if WINDOW_MODE == 0: + base_ptr = ( + inter_ptr + + inter_coord * stride_inter_seq + + idx_token * stride_inter_step + + idx_feats * stride_inter_dim + ) + tl.store(base_ptr + 0 * stride_inter_win, col0, mask=mask_w) + tl.store(base_ptr + 1 * stride_inter_win, col1, mask=mask_w) + tl.store(base_ptr + 2 * stride_inter_win, col2, mask=mask_w) + if SILU_ACTIVATION: + acc = acc / (1 + tl.exp(-acc)) + tl.store( + o_ptr + + idx_seq * stride_o_seq + + idx_token * stride_o_token + + (idx_feats * stride_o_dim), + acc, + mask=(idx_token < seqlen) & (idx_feats < dim), + ) + if WINDOW_MODE == 1: + # the deduplicated windows (window[t][w] = row[t + w]) as one row a channel: the prior two state + # columns, then every x row; each element once, the values the per-token stores leave there + j = tl.arange(0, 16) + row = tl.where( + j[None, :] == 0, + prior1[:, None], + tl.zeros((BLOCK_N, 16), dtype=prior1.dtype), + ) + row = tl.where(j[None, :] == 1, prior2[:, None], row) + row = tl.where(j[None, :] == 2, x0[:, None], row) + if seqlen > 1: + row = tl.where(j[None, :] == 3, x1[:, None], row) + if seqlen > 2: + row = tl.where(j[None, :] == 4, x2[:, None], row) + if seqlen > 3: + row = tl.where(j[None, :] == 5, x3[:, None], row) + if seqlen > 4: + row = tl.where(j[None, :] == 6, x4[:, None], row) + if seqlen > 5: + row = tl.where(j[None, :] == 7, x5[:, None], row) + if seqlen > 6: + row = tl.where(j[None, :] == 8, x6[:, None], row) + if seqlen > 7: + row = tl.where(j[None, :] == 9, x7[:, None], row) + tl.store( + inter_ptr + + inter_coord * stride_inter_seq + + (idx_feats * stride_inter_dim)[:, None] + + j[None, :], + row, + mask=(idx_feats < dim)[:, None] & (j < seqlen + 2)[None, :], + ) + + return kernel diff --git a/sgl_nvidia_mega/qwen3_5/gdn/front.py b/sgl_nvidia_mega/qwen3_5/gdn/front.py index 3a8fb912..f36e686f 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/front.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/front.py @@ -231,6 +231,19 @@ def step(gdn, x: torch.Tensor, st, opts) -> torch.Tensor: ring_indices=st.fb.req_pool_indices[:1], pad_slot=meta.pad_slot_id, ) + if opts.core: + from sgl_nvidia_mega.qwen3_5.gdn.verify import _fused + + # b / a in the front's b8 / a8 order + return _fused( + gdn, + conv, + projected[:, CONV_CHANNELS:].view(ROWS, attn.num_v_heads, attn.head_v_dim), + ba[:, 8:], + ba[:, :8], + st, + opts, + ) query, key, value = torch.split(conv, [attn.q_dim, attn.k_dim, attn.v_dim], dim=-1) H, K = attn.num_k_heads, attn.head_k_dim HV, V = attn.num_v_heads, attn.head_v_dim diff --git a/sgl_nvidia_mega/qwen3_5/gdn/verify.py b/sgl_nvidia_mega/qwen3_5/gdn/verify.py index 845da9c1..d28e098a 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/verify.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/verify.py @@ -167,6 +167,10 @@ class Options(msgspec.Struct, frozen=True): norm_fp8: bool # the gated norm writes out_proj's static FP8 input pdl: bool = False # the verify launches dependent on the conv (slot gdn_pdl) push: bool = False # out_proj pushes its partial into the hop (slot gdn_push) + core: bool = False # verify, gated norm and FP8 in one launch (slot gdn_core) + conv_rows: bool = ( + False # the hoisted conv writing each channel's window row once (slot gdn_conv) + ) STOCK = Options(views=False, block_n=256, block_v=64, num_warps=1, norm_fp8=False) @@ -226,6 +230,42 @@ def pool_reasons(pools, tokens_per_request: int) -> list[str]: return reasons +def conv_pool_reasons(pools, tokens_per_request: int) -> list[str]: + """Why slot gdn_conv cannot run on the pools ([] when it can): width-4 conv states and the + deduplicated windows (window[t][w] = row[t + w]) of 4 to 8 rows a request.""" + reasons = [] + conv, windows = pools.conv[0], pools.intermediate_conv_window[0] + if conv.shape[-1] != 3: + reasons.append(f"conv states {tuple(conv.shape)} are not width 4") + st = windows.stride() + if not (st[2] == 1 and st[4] == 1 and st[3] >= tokens_per_request + 2): + reasons.append(f"conv windows {tuple(windows.shape)} {st} are not deduplicated") + if not 3 < tokens_per_request <= 8: + reasons.append(f"{tokens_per_request} rows a request") + return reasons + + +def core_pool_reasons(pools) -> list[str]: + """Why slot gdn_core cannot write this step's ring entries ([] when it can): the + compact ring (update / key low parts, no exact-replay buffers) of BF16 checkpoints, + int32 cursors and int8 flush flags.""" + reasons = [] + if ( + pools.ring_beta is not None + or pools.ring_rawk is None + or pools.ring_rawv is None + ): + reasons.append("not the compact (low-part) ReplaySSM ring") + if ( + pools.ring_write_pos.dtype != torch.int32 + or pools.ring_cache_base.dtype != torch.int32 + ): + reasons.append("ring cursors are not int32") + if pools.ring_is_flush.dtype != torch.int8: + reasons.append("ring flush flags are not int8") + return reasons + + def step(gdn, x: torch.Tensor, st, opts: Options = SLOT) -> torch.Tensor: """``Qwen3_5GatedDeltaNet.forward`` of a target-verify step with the launch choices ``opts``: the TP-partial output [rows, hidden] of the normed input x.""" @@ -256,6 +296,8 @@ def step(gdn, x: torch.Tensor, st, opts: Options = SLOT) -> torch.Tensor: gdn.head_k_dim, gdn.head_v_dim, ) + if opts.core: + return _fused(gdn, _conv(gdn.attn, mixed_qkv, st, opts), z, a, b, st, opts) core = _core(gdn.attn, mixed_qkv, a, b, st, opts) core = core.reshape(-1, core.shape[-1]) nv = gdn.num_v_heads // gdn.attn_tp_size @@ -306,6 +348,86 @@ def _out_proj(out_proj, fp8: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: ) +def _conv(attn, mixed_qkv, st, opts) -> torch.Tensor: + """The causal conv update of the step's q / k / v rows: [rows, channels], token-major + with ``opts.views``.""" + from sgl_nvidia_mega.qwen3_5.gdn.conv import conv_update + + meta, pools = st.view.gdn, st.view.mamba + li = pools.layer_ids.index(attn.layer_id) + seq_len = mixed_qkv.shape[0] + draft = st.binding.tokens_per_request + bs = seq_len // draft + states = meta.state_indices + y = conv_update( + mixed_qkv.view(bs, draft, -1).transpose(1, 2), + pools.conv[0][li], + attn.conv_weights, + attn.bias, + attn.activation, + states[:bs], + pools.intermediate_conv_window[0][li], + meta.intermediate_indices[:bs], + block_n=opts.block_n, + token_major_out=opts.views, + **({"hoist": True, "window_mode": 1} if opts.conv_rows else {}), + ) + return y.transpose(1, 2).view(seq_len, -1) + + +def _fused(gdn, qkv, z, a, b, st, opts) -> torch.Tensor: + """Slot gdn_core: the ReplaySSM verify of the convolved rows ``qkv``, the gated norm + and its static FP8 output in one launch (``sm103/gdn/core.py``), then ``out_proj`` + (FlashInfer's BMM, or the hop's push GEMM with ``opts.push``).""" + from sgl_nvidia_mega.qwen3_5.sm103.gdn.core import core + + attn, pools = gdn.attn, st.view.mamba + li = pools.layer_ids.index(attn.layer_id) + rows = qkv.shape[0] + nv = gdn.num_v_heads // gdn.attn_tp_size + x = torch.empty( + (rows, nv * gdn.head_v_dim), dtype=torch.float8_e4m3fn, device=qkv.device + ) + + def ring(t): + return None if t is None else t[li] + + core( + qkv, + z, + a, + b, + attn.A_log, + attn.dt_bias, + pools.temporal[li], + pools.ring_d[li], + pools.ring_k[li], + pools.ring_g[li], + ring(pools.ring_rawv), + ring(pools.ring_rawk), + st.view.gdn.state_indices, + st.fb.req_pool_indices, + pools.ring_write_pos, + pools.ring_cache_base, + pools.ring_is_flush, + gdn.norm.weight, + gdn.out_proj.input_scale, + gdn.norm.eps, + x, + rows_per_request=st.binding.tokens_per_request, + num_k_heads=attn.num_k_heads, + early_dependents=opts.push, + ) + if opts.push: + return st.bound["boundary"].hop.out_proj_push( + x_fp8=x, + packed_w=st.bound["gdn"].packed[attn.layer_id], + input_scale=gdn.out_proj.input_scale, + weight_scale=gdn.out_proj.weight_scale, + ) + return _out_proj(gdn.out_proj, x, qkv.dtype) + + def _core(attn, mixed_qkv, a, b, st, opts): """The backend's target-verify arm for the circular ReplaySSM ring.""" from sgl_nvidia_mega.qwen3_5.gdn import replayssm @@ -328,6 +450,7 @@ def _core(attn, mixed_qkv, a, b, st, opts): meta.intermediate_indices[:bs], block_n=opts.block_n, token_major_out=opts.views, + **({"hoist": True, "window_mode": 1} if opts.conv_rows else {}), ) mixed_qkv = y.transpose(1, 2).view(seq_len, -1) query, key, value = torch.split( diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index e2dc0eee..932dd70c 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -8,8 +8,10 @@ captures up to 128); a width of more than one request is covered only where a slot runs. ``route`` (the router GEMM and the routed experts' NVFP4 input in one launch) runs up to 16 rows, ``mkm`` (on its outputs, the top-10 and both expert GEMMs in one launch) up to -7, and with ``mkm_wide`` (its 16-token engine) up to 16. ``gdn_verify`` (the GatedDeltaNet verify chain with its launch choices changed) covers -every width, and ``hop`` (the boundary's fused all-reduce points on the decode hop) the +7, and with ``mkm_wide`` (its 16-token engine) up to 16. ``gdn_verify`` (the +GatedDeltaNet verify chain with its launch choices changed) covers every width, with +``gdn_core`` (numeric: its ReplaySSM verify, gated norm and FP8 output in one launch) +up to 224 rows, and ``hop`` (the boundary's fused all-reduce points on the decode hop) the widths where FlashInfer's low-latency kernels, which it equals bit for bit, run. ``gdn_push`` (numeric: out_proj as the hop's push GEMM, on gdn_verify's FP8 norm) runs up to 8 rows. ``gdn_front`` (numeric: the GDN projections and conv in one launch, then @@ -41,6 +43,7 @@ # hop the rows up to which FlashInfer's MNNVL CuTe DSL fusion runs the one-shot # all-reduce the hop equals bit for bit (its finalize runs one-shot up to 23 rows) ROUTE_MAX_ROWS = 16 +GDN_CORE_MAX_ROWS = 224 # and gdn_core's launch MKM_MAX_ROWS = 7 MKM_WIDE_MAX_ROWS = 16 HOP_MAX_ROWS = 15 @@ -61,6 +64,10 @@ class Slots(msgspec.Struct, frozen=True): ) gdn_verify: bool = False # the GDN verify chain's launch choices (gdn/verify.py) gdn_pdl: bool = False # its ReplaySSM verify launched dependent on the conv + # its ReplaySSM verify, gated norm and FP8 output in one launch (sm103/gdn/core.py) + gdn_core: bool = False + # its conv with every x row loaded first and each window row in one store (gdn/conv.py; bitwise) + gdn_conv: bool = False hop: bool = False # the boundary's fused all-reduce points (boundary/hop.py) # out_proj pushes its partial into the hop from its GEMM (gdn/verify.py; numeric) gdn_push: bool = False @@ -82,6 +89,8 @@ class Facts(msgspec.Struct, frozen=True): mkm_wide: bool = False # with their 16-token engine gdn_verify: bool = False # the GDN verify slot built for every GDN layer gdn_pdl: bool = False # with its dependent verify launch + gdn_core: bool = False # with the one-launch verify core + gdn_conv: bool = False # with the hoisted conv writing each window row once hop: bool = False # the decode hop, built for the widest step it covers gdn_push: bool = False # out_proj packed for the push GEMM on every GDN layer gdn_front: bool = False # in_proj_qkvz packed for the front on every GDN layer @@ -107,6 +116,8 @@ def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> and rows <= (MKM_WIDE_MAX_ROWS if facts.mkm_wide else MKM_MAX_ROWS), gdn_verify=facts.gdn_verify, gdn_pdl=facts.gdn_verify and facts.gdn_pdl, + gdn_core=facts.gdn_verify and facts.gdn_core and rows <= GDN_CORE_MAX_ROWS, + gdn_conv=facts.gdn_verify and facts.gdn_conv, hop=hop, gdn_push=facts.gdn_verify and hop diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 1e5fddbb..edab5f39 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -23,6 +23,9 @@ "mkm_wide", "gdn_verify", "gdn_pdl", + "gdn_core", + # the verify conv with every x row loaded first and the windows in one store a channel (bitwise) + "gdn_conv", "hop", "gdn_push", "gdn_front", @@ -128,6 +131,12 @@ def slot_reasons() -> list[str]: reasons.append( f"{SLOTS_ENV}: slot 'gdn_front' runs slot 'gdn_pdl''s dependent verify" ) + if "gdn_core" in SLOTS and "gdn_verify" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'gdn_core' runs slot 'gdn_verify''s step in one launch" + ) + if "gdn_conv" in SLOTS and "gdn_verify" not in SLOTS: + reasons.append(f"{SLOTS_ENV}: slot 'gdn_conv' runs slot 'gdn_verify''s conv") if "gdn_pdl" in SLOTS and "gdn_verify" not in SLOTS: reasons.append( f"{SLOTS_ENV}: slot 'gdn_pdl' launches slot 'gdn_verify''s verify" @@ -229,6 +238,8 @@ def build(model, ctx): mkm_wide=bound["moe"].mkm_wide, gdn_verify=bound["gdn"].gdn_verify, gdn_pdl=bound["gdn"].gdn_pdl, + gdn_core=bound["gdn"].gdn_core, + gdn_conv=bound["gdn"].gdn_conv, hop=bound["boundary"].hop is not None, gdn_push=bool(bound["gdn"].packed), gdn_front=bool(bound["gdn"].front), diff --git a/sgl_nvidia_mega/qwen3_5/runtime/executor.py b/sgl_nvidia_mega/qwen3_5/runtime/executor.py index 4a77a28d..24d7973c 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/executor.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/executor.py @@ -185,6 +185,20 @@ def _check_pools(self, view): "NVIDIA megakernel: slot gdn_verify cannot run: " + "; ".join(reasons) ) + if any(s.gdn_conv for s in self.slots.values()): + reasons = gdn.verify.conv_pool_reasons( + view.mamba, self.binding.tokens_per_request + ) + if reasons: + raise RuntimeError( + "NVIDIA megakernel: slot gdn_conv cannot run: " + "; ".join(reasons) + ) + if any(s.gdn_core for s in self.slots.values()): + reasons = gdn.verify.core_pool_reasons(view.mamba) + if reasons: + raise RuntimeError( + "NVIDIA megakernel: slot gdn_core cannot run: " + "; ".join(reasons) + ) if any(s.gdn_front for s in self.slots.values()): from sgl_nvidia_mega.qwen3_5.gdn import front diff --git a/sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py b/sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py new file mode 100644 index 00000000..b4cbeccc --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""The GatedDeltaNet verify core of a target-verify step in one launch (``src/core.cu``): the +ReplaySSM verify on the circular ring (non-flush rows; this step's ring entries written as +SGLang's kernel writes them) and the gated RMSNorm writing out_proj's static FP8 input. One CTA +per (request, value head); numeric against slot gdn_verify's two launches (FP32 sums in other +orders, the same BF16 rounding points). + +It launches as a programmatic dependent: the checkpoint rows, ring cursors and static parameters +load before ``griddepcontrol.wait``, every activation (q / k / v, z, a, b) after it. With +``early_dependents`` the next launch may start at once (it must wait before reading the output). +""" + +from __future__ import annotations + +import ctypes +from functools import cache +from pathlib import Path +from typing import Optional + +import torch + +from sgl_nvidia_mega.common import jit + +SOURCE = Path(__file__).parent / "src" / "core.cu" +THREADS, K, V, RING, MAX_ROWS_PER_REQUEST = 128, 128, 128, 16, 8 + + +class Params(ctypes.Structure): + """Params of src/core.cu.""" + + _fields_ = [ + ("x", ctypes.c_void_p), + ("x_stride", ctypes.c_int64), + ("z", ctypes.c_void_p), + ("z_stride", ctypes.c_int64), + ("a", ctypes.c_void_p), + ("b", ctypes.c_void_p), + ("ab_stride", ctypes.c_int64), + ("A_log", ctypes.c_void_p), + ("dt_bias", ctypes.c_void_p), + ("state", ctypes.c_void_p), + ("state_slot", ctypes.c_int64), + ("ring_d", ctypes.c_void_p), + ("d_slot", ctypes.c_int64), + ("ring_k", ctypes.c_void_p), + ("k_slot", ctypes.c_int64), + ("ring_g", ctypes.c_void_p), + ("g_slot", ctypes.c_int64), + ("ring_dlo", ctypes.c_void_p), + ("dlo_slot", ctypes.c_int64), + ("ring_klo", ctypes.c_void_p), + ("klo_slot", ctypes.c_int64), + ("state_indices", ctypes.c_void_p), + ("replay_indices", ctypes.c_void_p), + ("state_index64", ctypes.c_int), + ("replay_index64", ctypes.c_int), + ("write_pos", ctypes.c_void_p), + ("cache_base", ctypes.c_void_p), + ("is_flush", ctypes.c_void_p), + ("norm_w", ctypes.c_void_p), + ("scale", ctypes.c_void_p), + ("eps", ctypes.c_float), + ("qscale", ctypes.c_float), + ("y", ctypes.c_void_p), + ("y_stride", ctypes.c_int64), + ("o_out", ctypes.c_void_p), + ("prof", ctypes.c_void_p), + ] + + +@cache +def _kernel( + device_index: int, + rows: int, + hk: int, + hv: int, + norm_f32: bool, + trigger: int, + store_o: bool, + prof: bool = False, + warps: int = 4, + min_blocks: int = 1, +) -> jit.Kernel: + major, minor = torch.cuda.get_device_capability(device_index) + defines = dict( + TT=rows, + HK=hk, + HV=hv, + NORM_F32=int(norm_f32), + TRIGGER=trigger, + STORE_O=int(store_o), + PROF=int(prof), + NW=warps, + MINB=min_blocks, + ) + source = "".join(f"#define {k} {v}\n" for k, v in defines.items()) + cubin = jit.compile_source( + source + SOURCE.read_text(), + "q35_gdn_core", + f"sm_{major}{minor}a", + ("-std=c++17", "-default-device"), + ) + return jit.load(cubin, "q35_gdn_core", (32 * warps, 1, 1), 0) + + +def build(device_index: int, rows: int, hk: int, hv: int, norm_f32: bool) -> None: + """Compile the core for a step shape, both dependent releases and both register caps + (no compile inside a step: CUDA-graph capture refuses it).""" + for trigger in (0, 2): + for min_blocks in (1, 3): + _kernel( + device_index, + rows, + hk, + hv, + norm_f32, + trigger, + False, + min_blocks=min_blocks, + ) + + +@cache +def _sms(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +def _min_blocks(device_index: int, ctas: int) -> int: + """Three CTAs an SM (168 registers) once two an SM take more than two waves (128 requests + at TP8): past that the third resident CTA wins; below it the uncapped registers do.""" + return 3 if ctas > 4 * _sms(device_index) else 1 + + +def _slot_stride(t: Optional[torch.Tensor]) -> int: + return 0 if t is None else t.stride(0) + + +def core( + qkv: torch.Tensor, + z: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + checkpoint: torch.Tensor, + ring_d: torch.Tensor, + ring_k: torch.Tensor, + ring_g: torch.Tensor, + ring_dlo: Optional[torch.Tensor], + ring_klo: Optional[torch.Tensor], + state_indices: torch.Tensor, + replay_indices: torch.Tensor, + write_pos: torch.Tensor, + cache_base: torch.Tensor, + is_flush: torch.Tensor, + norm_weight: torch.Tensor, + input_scale: torch.Tensor, + eps: float, + out: torch.Tensor, + *, + rows_per_request: int, + num_k_heads: int, + early_dependents: bool = False, + o_out: Optional[torch.Tensor] = None, + _profile: Optional[torch.Tensor] = None, + warps: int = 4, + min_blocks: Optional[int] = None, +) -> torch.Tensor: + """Write the gated RMSNorm of the verify output as FP8 into ``out`` [rows, HV * V]. + + ``qkv`` [rows, >= (2 H + HV) * 128] BF16 convolved q | k | v rows (unit column stride, + 16-byte aligned rows); ``z`` [rows, HV, V] with z.stride(1) == V; ``a`` / ``b`` [rows, HV] + row-strided views with one row stride; ``checkpoint`` [slots, HV, V, K] BF16; the ring + [slots, HV | H, 16, V | K] (low parts optional) with cursors by ``replay_indices``; the first + ``state_indices.numel()`` requests of ``rows_per_request`` rows each. ``o_out`` [rows, HV, V] + BF16 also receives the replay output (checks only).""" + rows, T = qkv.shape[0], rows_per_request + n = rows // T + HV = checkpoint.shape[1] + H = num_k_heads + if ( + qkv.dtype != torch.bfloat16 + or qkv.stride(1) != 1 + or qkv.stride(0) % 8 + or qkv.data_ptr() % 16 + or qkv.shape[1] < (2 * H + HV) * K + ): + raise ValueError( + "q | k | v rows must be BF16, unit column stride, 16-byte aligned" + ) + if not 2 <= T <= MAX_ROWS_PER_REQUEST or rows % T or state_indices.numel() < n: + raise ValueError(f"{rows} rows of {T} a request") + if tuple(checkpoint.shape[1:]) != (HV, V, K) or checkpoint.dtype != torch.bfloat16: + raise ValueError(f"checkpoint {checkpoint.dtype} {tuple(checkpoint.shape)}") + if checkpoint.stride()[1:] != (V * K, K, 1): + raise ValueError("checkpoint slots must be [HV, V, K] contiguous") + for t, heads, dim in ( + (ring_d, HV, V), + (ring_dlo, HV, V), + (ring_k, H, K), + (ring_klo, H, K), + ): + if t is not None and ( + tuple(t.shape[1:]) != (heads, RING, dim) + or t.stride()[1:] != (RING * dim, dim, 1) + ): + raise ValueError(f"ring {tuple(t.shape)} {tuple(t.stride())}") + if tuple(ring_g.shape[1:]) != (HV, RING) or ring_g.stride()[1:] != (RING, 1): + raise ValueError(f"ring g {tuple(ring_g.shape)}") + if z.shape != (rows, HV, V) or z.stride(2) != 1 or z.stride(1) != V: + raise ValueError(f"z {tuple(z.shape)} {tuple(z.stride())}") + if a.stride(1) != 1 or b.stride(1) != 1 or a.stride(0) != b.stride(0): + raise ValueError("a / b must be row-strided views with one row stride") + if ( + out.dtype != torch.float8_e4m3fn + or out.shape != (rows, HV * V) + or out.stride(1) != 1 + or out.stride(0) % 16 + or out.data_ptr() % 16 + ): + raise ValueError(f"out {out.dtype} {tuple(out.shape)}") + if write_pos.dtype != torch.int32 or cache_base.dtype != torch.int32: + raise ValueError("ring cursors must be int32") + if is_flush.dtype != torch.int8: + raise ValueError("is_flush must be int8") + p = Params( + x=qkv.data_ptr(), + x_stride=qkv.stride(0), + z=z.data_ptr(), + z_stride=z.stride(0), + a=a.data_ptr(), + b=b.data_ptr(), + ab_stride=a.stride(0), + A_log=A_log.data_ptr(), + dt_bias=dt_bias.data_ptr(), + state=checkpoint.data_ptr(), + state_slot=checkpoint.stride(0), + ring_d=ring_d.data_ptr(), + d_slot=ring_d.stride(0), + ring_k=ring_k.data_ptr(), + k_slot=ring_k.stride(0), + ring_g=ring_g.data_ptr(), + g_slot=ring_g.stride(0), + ring_dlo=0 if ring_dlo is None else ring_dlo.data_ptr(), + dlo_slot=_slot_stride(ring_dlo), + ring_klo=0 if ring_klo is None else ring_klo.data_ptr(), + klo_slot=_slot_stride(ring_klo), + state_indices=state_indices.data_ptr(), + replay_indices=replay_indices.data_ptr(), + state_index64=int(state_indices.dtype == torch.int64), + replay_index64=int(replay_indices.dtype == torch.int64), + write_pos=write_pos.data_ptr(), + cache_base=cache_base.data_ptr(), + is_flush=is_flush.data_ptr(), + norm_w=norm_weight.data_ptr(), + scale=input_scale.data_ptr(), + eps=eps, + qscale=K**-0.5, + y=out.data_ptr(), + y_stride=out.stride(0), + o_out=0 if o_out is None else o_out.data_ptr(), + prof=0 if _profile is None else _profile.data_ptr(), + ) + kernel = _kernel( + qkv.device.index, + T, + H, + HV, + norm_weight.dtype == torch.float32, + 2 if early_dependents else 0, + o_out is not None, + _profile is not None, + warps, + _min_blocks(qkv.device.index, n * HV) if min_blocks is None else min_blocks, + ) + jit.launch( + kernel, (n * HV, 1, 1), [p], torch.cuda.current_stream(qkv.device).cuda_stream + ) + return out diff --git a/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/core.cu b/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/core.cu new file mode 100644 index 00000000..c4d9f48f --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/core.cu @@ -0,0 +1,598 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Adapted from TensorRT-LLM (Apache-2.0, NVIDIA): tensorrt_llm/_torch/modules/fla/q35_gdn_core.py +// (q35_gdn_core: work split, MMA fragments and PDL placement), on SGLang's ReplaySSM circular ring +// (python/sglang/kernels/ops/attention/linear/gdn_replayssm_spec_decode.py, +// gdn_replayssm_spec_circular_kernel) and gated RMSNorm (kernels/ops/attention/fla/layernorm_gated.py). +// +// The ReplaySSM verify of a GDN target-verify step (the circular ring's non-flush rows) and its gated +// RMSNorm writing out_proj's static FP8 input, in one launch: one CTA per (request, value head), 4 +// warps, warp w owns the value columns 32w .. 32w + 31. Products run as mma.m16n8k16 (BF16 in, FP32 +// accumulate) on the normalized q / k rows and the checkpoint rows; the forward substitution, the +// intra-step output and the norm run in FP32. The BF16 roundings of SGLang's kernels are kept +// (normalized q, key high / low parts, update high / low parts, the replay output before the norm, +// the norm output before the FP8 scale); FP32 sums run in other orders (numeric). +// +// PDL: before griddepcontrol.wait the kernel reads only data of earlier steps or static data (indices, +// ring cursors, the checkpoint rows, A_log / dt_bias / norm weight / FP8 scale); every global store +// and every read of the step's activations (q / k / v, z, a, b) follows the wait. +// +// Defines: TT (rows a request), HK / HV (key / value heads), NORM_F32 (FP32 norm weight), TRIGGER (2: release dependents at entry, 0: at exit), +// STORE_O (also store the BF16 replay output, checks only). + +typedef unsigned char u8; +typedef unsigned short u16; +typedef unsigned int u32; + +#define KD 128 +#define VD 128 +#define RING 16 +#ifndef MINB +#define MINB 1 // resident CTAs an SM the registers are capped for +#endif +#ifndef NW +#define NW 4 // warps; each owns VD / NW value columns as MT m16 tiles +#endif +#define MT (VD / (16 * NW)) +#define FULL 0xffffffffu + +struct alignas(16) U4 { u32 x, y, z, w; }; + +struct Params { + const u16* x; // convolved q | k | v rows + long long x_stride; + const u16* z; // gate rows: z + row * z_stride + hv * VD + long long z_stride; + const u16* a; + const u16* b; + long long ab_stride; + const float* A_log; + const float* dt_bias; + const u16* state; // [slots, HV, V, K] + long long state_slot; + u16* ring_d; // [slots, HV, RING, V] + long long d_slot; + u16* ring_k; // [slots, HK, RING, K] + long long k_slot; + float* ring_g; // [slots, HV, RING] + long long g_slot; + u16* ring_dlo; // update low parts (layout of ring_d), or null + long long dlo_slot; + u16* ring_klo; // key low parts (layout of ring_k), or null + long long klo_slot; + const void* state_indices; + const void* replay_indices; + int state_index64; // int64 (else int32) state / replay indices + int replay_index64; + const int* write_pos; + const int* cache_base; + const signed char* is_flush; + const void* norm_w; + const float* scale; + float eps; + float qscale; + u8* y; // FP8 out_proj input [rows, HV * V] + long long y_stride; + u16* o_out; // [rows, HV, V] replay output (STORE_O) + long long* prof; // PROF: clock64 stamps [cta][warp][8] +}; + +__device__ __forceinline__ float bflo(u32 x) { return __int_as_float(x << 16); } +__device__ __forceinline__ float bfhi(u32 x) { return __int_as_float(x & 0xffff0000u); } +__device__ __forceinline__ float bf2f(u16 x) { return __int_as_float(((u32)x) << 16); } +__device__ __forceinline__ u16 f2bf(float x) { + u16 r; + asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(r) : "f"(x)); + return r; +} +__device__ __forceinline__ float bfround(float x) { return bf2f(f2bf(x)); } +__device__ __forceinline__ u32 pack2(float lo, float hi) { + u32 r; + asm("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(r) : "f"(hi), "f"(lo)); + return r; +} +__device__ __forceinline__ u32 word(const U4& v, int i) { return i == 0 ? v.x : i == 1 ? v.y : i == 2 ? v.z : v.w; } +__device__ __forceinline__ U4 ld16(const void* p) { return *reinterpret_cast(p); } +__device__ __forceinline__ void st16(void* p, U4 v) { *reinterpret_cast(p) = v; } +__device__ __forceinline__ U4 zero4() { U4 z; z.x = z.y = z.z = z.w = 0u; return z; } +__device__ __forceinline__ float rcp_rn(float x) { + float r; + asm("div.rn.f32 %0, 0f3F800000, %1;" : "=f"(r) : "f"(x)); + return r; +} +__device__ __forceinline__ u8 f2e4m3(float x) { + u16 r; + asm("cvt.rn.satfinite.e4m3x2.f32 %0, 0f00000000, %1;" : "=h"(r) : "f"(x)); + return (u8)(r & 0xff); +} +__device__ __forceinline__ float pick8(const float (&a)[8], int i) { + float r = a[0]; +#pragma unroll + for (int k = 1; k < 8; ++k) r = (i == k) ? a[k] : r; + return r; +} +__device__ __forceinline__ void mma16816(float (&d)[4], u32 a0, u32 a1, u32 a2, u32 a3, u32 b0, u32 b1) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} +// K-reduction of 128-wide rows: A = rows r (rows r + 8 zero), B = 8 rows as columns. Lane (r, q) holds the +// elements [32 c + 8 q, + 8) of its row in chunk c; k-step 2c + e takes words (2e, 2e + 1) of both operands. +__device__ __forceinline__ void kred(float (&acc)[4], const U4 (&A)[4], const U4 (&B)[4]) { +#pragma unroll + for (int c = 0; c < 4; ++c) +#pragma unroll + for (int e = 0; e < 2; ++e) + mma16816(acc, word(A[c], 2 * e), 0u, word(A[c], 2 * e + 1), 0u, word(B[c], 2 * e), word(B[c], 2 * e + 1)); +} +__device__ __forceinline__ void pdl_wait() { asm volatile("griddepcontrol.wait;" ::: "memory"); } +__device__ __forceinline__ void pdl_launch() { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); } + +__device__ __forceinline__ long long load_index(const void* p, int i, bool i64) { + return i64 ? reinterpret_cast(p)[i] : (long long)reinterpret_cast(p)[i]; +} + +#ifndef PROF +#define PROF 0 +#endif +#if PROF +#define STAMP(i, dep) do { if (lane == 0) { asm volatile("" ::"f"((float)(dep))); st_[i] = clock64(); } } while (0) +#else +#define STAMP(i, dep) do { } while (0) +#endif + +// MINB 3: three CTAs an SM (168 registers, no spill), for grids past two waves of two an SM +extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const Params p) { + constexpr int GRP = HV / HK; + const int cta = blockIdx.x; + const int n = cta / HV; + const int hv = cta - n * HV; + const int hk = hv / GRP; + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int r = lane >> 2; + const int q = lane & 3; + const int vb = warp * 16 * MT; + const int t0 = 2 * q; + const int t1 = 2 * q + 1; + const long long row0 = (long long)n * TT; + + __shared__ float s_A[8][8]; // A[s][t] = beta_s e^{G_s - G_t} k_s.k_t, s > t + __shared__ float s_F[8][8]; // F[s][t] = e^{G_t - G_s} k_s.q_t, s <= t + __shared__ u32 s_sc[2][2][32]; // history scores (q / k, history tile) as BF16 B fragments + __shared__ float s_norm[NW][8]; + +#if PROF + long long st_[8]; + st_[0] = clock64(); +#endif +#if TRIGGER == 2 + pdl_launch(); +#endif + // ---------------- before the wait: earlier-step and static data ---------------- + const long long slot = load_index(p.state_indices, n, p.state_index64 != 0); + bool live = slot >= 0; + // the checkpoint rows first: they depend on the state index only + const u16* sp = p.state + (live ? slot : 0) * p.state_slot + (long long)hv * VD * KD; + U4 H[MT][2][4]; // checkpoint rows v = vb + 16 m + 8 h + r: [m-tile][half][chunk] +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int c = 0; c < 4; ++c) + H[m][h][c] = live ? ld16(sp + (vb + 16 * m + 8 * h + r) * KD + 32 * c + 8 * q) : zero4(); + long long rid = 0; + int wp = 0, base = 0; + if (live) { + rid = load_index(p.replay_indices, n, p.replay_index64 != 0); + wp = p.write_pos[rid]; + base = p.cache_base[rid]; + live = p.is_flush[rid] == 0; + } + if (!live) { + // padding (or a flush row, which the verify never sees): zeros for out_proj + pdl_wait(); + for (int t = 0; t < TT && tid < VD; ++t) { + p.y[(row0 + t) * p.y_stride + hv * VD + tid] = 0; +#if STORE_O + p.o_out[((row0 + t) * HV + hv) * VD + tid] = 0; +#endif + } + return; + } + const float negAexp = -__expf(p.A_log[hv]); + const float dtb = p.dt_bias[hv]; + float nw[MT][2]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int v = vb + 16 * m + 8 * h + r; +#if NORM_F32 + nw[m][h] = reinterpret_cast(p.norm_w)[v]; +#else + nw[m][h] = bf2f(reinterpret_cast(p.norm_w)[v]); +#endif + } + const float rscale = rcp_rn(p.scale[0]); + + STAMP(1, bflo(H[1][1][3].w) + bflo(H[0][0][0].x) + rscale); + // ---------------- the wait: the step's activations from here on ---------------- + pdl_wait(); + STAMP(2, 0.f); + + const u16* xr = p.x + (row0 + (r < TT ? r : 0)) * p.x_stride; + const bool rowok = r < TT; + const bool ok0 = t0 < TT, ok1 = t1 < TT; + // every activation load is issued here, right after the wait + U4 QN[4], KR[4]; // normalized q (BF16) and raw k rows; rows r >= TT zero + U4 QR[4]; +#pragma unroll + for (int c = 0; c < 4; ++c) { + QR[c] = rowok ? ld16(xr + hk * KD + 32 * c + 8 * q) : zero4(); + KR[c] = rowok ? ld16(xr + HK * KD + hk * KD + 32 * c + 8 * q) : zero4(); + } + // v / z of this lane's (v, t): [m][(h, t)] = (h0, t0), (h0, t1), (h1, t0), (h1, t1) + u16 vr[MT][4], zr[MT][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int v = vb + 16 * m + 8 * h + r; + const u16* xv = p.x + 2 * HK * KD + hv * VD + v; + const u16* zv = p.z + hv * VD + v; + vr[m][2 * h] = ok0 ? xv[(row0 + t0) * p.x_stride] : (u16)0; + vr[m][2 * h + 1] = ok1 ? xv[(row0 + t1) * p.x_stride] : (u16)0; + zr[m][2 * h] = ok0 ? zv[(row0 + t0) * p.z_stride] : (u16)0; + zr[m][2 * h + 1] = ok1 ? zv[(row0 + t1) * p.z_stride] : (u16)0; + } + u16 ar = 0, br_ = 0; + if (lane < TT) { + ar = p.a[(row0 + lane) * p.ab_stride + hv]; + br_ = p.b[(row0 + lane) * p.ab_stride + hv]; + } + float kmul; // 1 / |k_r| + { + float sq[4] = {0.f, 0.f, 0.f, 0.f}, sk[4] = {0.f, 0.f, 0.f, 0.f}; +#pragma unroll + for (int c = 0; c < 4; ++c) +#pragma unroll + for (int i = 0; i < 4; ++i) { + const u32 wq = word(QR[c], i), wk = word(KR[c], i); + sq[i] += bflo(wq) * bflo(wq) + bfhi(wq) * bfhi(wq); + sk[i] += bflo(wk) * bflo(wk) + bfhi(wk) * bfhi(wk); + } + float sqs = (sq[0] + sq[1]) + (sq[2] + sq[3]), sks = (sk[0] + sk[1]) + (sk[2] + sk[3]); + sqs += __shfl_xor_sync(FULL, sqs, 1); + sks += __shfl_xor_sync(FULL, sks, 1); + sqs += __shfl_xor_sync(FULL, sqs, 2); + sks += __shfl_xor_sync(FULL, sks, 2); + const float qmul = rowok ? (1.0f / sqrtf(sqs + 1e-6f)) * p.qscale : 0.f; + kmul = rowok ? 1.0f / sqrtf(sks + 1e-6f) : 0.f; +#pragma unroll + for (int c = 0; c < 4; ++c) { + QN[c].x = pack2(bflo(QR[c].x) * qmul, bfhi(QR[c].x) * qmul); + QN[c].y = pack2(bflo(QR[c].y) * qmul, bfhi(QR[c].y) * qmul); + QN[c].z = pack2(bflo(QR[c].z) * qmul, bfhi(QR[c].z) * qmul); + QN[c].w = pack2(bflo(QR[c].w) * qmul, bfhi(QR[c].w) * qmul); + } + } + const float km0 = __shfl_sync(FULL, kmul, 4 * t0), km1 = __shfl_sync(FULL, kmul, 4 * t1); + STAMP(3, km0 + bflo(QN[3].w)); + // gates of token = lane, cumulative log-decay by an 8-lane scan + float gl = 0.f, beta = 0.f; + if (lane < TT) { + const float xg = bf2f(ar) + dtb; + const float spl = xg <= 20.f ? __logf(1.f + __expf(xg)) : xg; + gl = negAexp * spl; + beta = __fdividef(1.f, 1.f + __expf(-bf2f(br_))); + } + // the norm's gate z * sigmoid(z), here so that z arrives while the products run + float zg[MT][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float zv = bf2f(zr[m][i]); + zg[m][i] = zv * __fdividef(1.f, 1.f + __expf(-zv)); + } + float Gv = gl; +#pragma unroll + for (int d = 1; d < 8; d <<= 1) { + const float o = __shfl_up_sync(FULL, Gv, d); + if (lane >= d) Gv += o; + } + const float Gt0 = __shfl_sync(FULL, Gv, t0), Gt1 = __shfl_sync(FULL, Gv, t1); + const float bt0 = __shfl_sync(FULL, beta, t0), bt1 = __shfl_sync(FULL, beta, t1); + const float Gr = __shfl_sync(FULL, Gv, r), br = __shfl_sync(FULL, beta, r); + const float eG0 = __expf(Gt0), eG1 = __expf(Gt1); + + // checkpoint products (all warps): rows v, columns = tokens; the key product on raw rows, scaled by 1/|k_t| + float HQ[MT][4], HKs[MT][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) { +#pragma unroll + for (int i = 0; i < 4; ++i) HQ[m][i] = HKs[m][i] = 0.f; +#pragma unroll + for (int c = 0; c < 4; ++c) +#pragma unroll + for (int e = 0; e < 2; ++e) { + const u32 a0 = word(H[m][0][c], 2 * e), a1 = word(H[m][1][c], 2 * e); + const u32 a2 = word(H[m][0][c], 2 * e + 1), a3 = word(H[m][1][c], 2 * e + 1); + mma16816(HKs[m], a0, a1, a2, a3, word(KR[c], 2 * e), word(KR[c], 2 * e + 1)); + mma16816(HQ[m], a0, a1, a2, a3, word(QN[c], 2 * e), word(QN[c], 2 * e + 1)); + } + HKs[m][0] *= km0; + HKs[m][1] *= km1; + HKs[m][2] *= km0; + HKs[m][3] *= km1; + } + // this step's ring entries at (base + wp + t) mod RING, t < TT, wp + t < RING: keys (warp 2) and + // log-decays (warp 3) now, while warps 0 / 1 finish the head-common products; updates at the end + if (warp == 3 && lane < TT && wp + lane < RING) + p.ring_g[rid * p.g_slot + hv * RING + ((base + wp + lane) & (RING - 1))] = gl; + if (warp == 2 && hv % GRP == 0 && r < TT && wp + r < RING) { + const long long off = ((long long)hk * RING + ((base + wp + r) & (RING - 1))) * KD; + u16* kp = p.ring_k + rid * p.k_slot + off; + u16* klp = p.ring_klo ? p.ring_klo + rid * p.klo_slot + off : nullptr; +#pragma unroll + for (int c = 0; c < 4; ++c) { + u32 kh[4], kl[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const u32 wk = word(KR[c], i); + const float k0 = bflo(wk) * kmul, k1 = bfhi(wk) * kmul; + kh[i] = pack2(k0, k1); + kl[i] = pack2(k0 - bflo(kh[i]), k1 - bfhi(kh[i])); + } + U4 vh, vl; + vh.x = kh[0]; vh.y = kh[1]; vh.z = kh[2]; vh.w = kh[3]; + vl.x = kl[0]; vl.y = kl[1]; vl.z = kl[2]; vl.w = kl[3]; + st16(kp + 32 * c + 8 * q, vh); + if (klp) st16(klp + 32 * c + 8 * q, vl); + } + } + + // head-common products: warp 0 k k^T (A), warp 1 k q^T (F), warps 2 / 3 the history scores + if (warp == 0) { + float kk[4] = {0.f, 0.f, 0.f, 0.f}; + kred(kk, KR, KR); + const float ks = kmul; // row s = r + s_A[r][t0] = (r > t0 && r < TT) ? br * __expf(Gr - Gt0) * (kk[0] * ks * km0) : 0.f; + s_A[r][t1] = (r > t1 && r < TT) ? br * __expf(Gr - Gt1) * (kk[1] * ks * km1) : 0.f; + } else if (warp == 1) { + float kq[4] = {0.f, 0.f, 0.f, 0.f}; + kred(kq, KR, QN); + s_F[r][t0] = (r <= t0 && t0 < TT) ? __expf(Gt0 - Gr) * (kq[0] * kmul) : 0.f; + s_F[r][t1] = (r <= t1 && t1 < TT) ? __expf(Gt1 - Gr) * (kq[1] * kmul) : 0.f; + } else if ((warp == 2 || warp == 3) && wp > 0) { + // scores[j][t] = khist_j . (q_t | k_t), j = 8 jt + 2q + i (rows j >= wp zero), rounded to BF16 + U4 KHh[4]; +#pragma unroll + for (int jt = 0; jt < 2; ++jt) { + const int j = 8 * jt + r; + const u16* kh = p.ring_k + rid * p.k_slot + ((long long)hk * RING + ((base + j) & (RING - 1))) * KD; +#pragma unroll + for (int c = 0; c < 4; ++c) KHh[c] = j < wp ? ld16(kh + 32 * c + 8 * q) : zero4(); + float acc[4] = {0.f, 0.f, 0.f, 0.f}; + if (warp == 2) { + kred(acc, QN, KHh); + } else { + U4 KH[4]; +#pragma unroll + for (int c = 0; c < 4; ++c) { + KH[c].x = pack2(bflo(KR[c].x) * kmul, bfhi(KR[c].x) * kmul); + KH[c].y = pack2(bflo(KR[c].y) * kmul, bfhi(KR[c].y) * kmul); + KH[c].z = pack2(bflo(KR[c].z) * kmul, bfhi(KR[c].z) * kmul); + KH[c].w = pack2(bflo(KR[c].w) * kmul, bfhi(KR[c].w) * kmul); + } + kred(acc, KH, KHh); + } + s_sc[warp - 2][jt][lane] = pack2(acc[0], acc[1]); + } + } + STAMP(4, HQ[MT - 1][3] + HKs[MT - 1][3]); + __syncthreads(); + STAMP(5, 0.f); + + if (wp > 0) { + // committed history: hw = e^{g_total} hw + bf16(d_j e^{g_total - G_j}) . bf16(scores) + float gj = 0.f; + if (lane < RING && lane < wp) gj = p.ring_g[rid * p.g_slot + hv * RING + ((base + lane) & (RING - 1))]; +#pragma unroll + for (int d = 1; d < 32; d <<= 1) { + const float o = __shfl_up_sync(FULL, gj, d); + if (lane >= d) gj += o; + } + const float gtot = __shfl_sync(FULL, gj, 31); + const int jj[4] = {2 * q, 2 * q + 1, 2 * q + 8, 2 * q + 9}; + float dec[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float pj = __shfl_sync(FULL, gj, jj[i]); + dec[i] = jj[i] < wp ? __expf(gtot - pj) : 0.f; + } + const float td = __expf(gtot); + const u16* dp = p.ring_d + rid * p.d_slot + (long long)hv * RING * VD; + const u32 bq0 = s_sc[0][0][lane], bq1 = s_sc[0][1][lane]; + const u32 bk0 = s_sc[1][0][lane], bk1 = s_sc[1][1][lane]; +#pragma unroll + for (int m = 0; m < MT; ++m) { + float e[2][4]; // [half][j index i] +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int i = 0; i < 4; ++i) + e[h][i] = jj[i] < wp + ? bf2f(dp[((base + jj[i]) & (RING - 1)) * VD + vb + 16 * m + 8 * h + r]) * dec[i] + : 0.f; + const u32 a0 = pack2(e[0][0], e[0][1]), a1 = pack2(e[1][0], e[1][1]); + const u32 a2 = pack2(e[0][2], e[0][3]), a3 = pack2(e[1][2], e[1][3]); +#pragma unroll + for (int i = 0; i < 4; ++i) { + HQ[m][i] *= td; + HKs[m][i] *= td; + } + mma16816(HQ[m], a0, a1, a2, a3, bq0, bq1); + mma16816(HKs[m], a0, a1, a2, a3, bk0, bk1); + } + } + + // rhs = beta (v - e^G hw_k); forward substitution (I + A) U = rhs per value column, every lane of a quad + // solving the quad's four columns after an all-gather + float U[2 * MT][8]; + { + float rhs[MT][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) { + rhs[m][0] = bt0 * (bf2f(vr[m][0]) - eG0 * HKs[m][0]); + rhs[m][1] = bt1 * (bf2f(vr[m][1]) - eG1 * HKs[m][1]); + rhs[m][2] = bt0 * (bf2f(vr[m][2]) - eG0 * HKs[m][2]); + rhs[m][3] = bt1 * (bf2f(vr[m][3]) - eG1 * HKs[m][3]); + } + float R[2 * MT][8]; +#pragma unroll + for (int qq = 0; qq < 4; ++qq) { + const int src = (lane & ~3) | qq; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int h = 0; h < 2; ++h) { + R[2 * m + h][2 * qq] = __shfl_sync(FULL, rhs[m][2 * h], src); + R[2 * m + h][2 * qq + 1] = __shfl_sync(FULL, rhs[m][2 * h + 1], src); + } + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + float u[2 * MT] = {}; + if (i < TT) { +#pragma unroll + for (int col = 0; col < 2 * MT; ++col) u[col] = R[col][i]; +#pragma unroll + for (int s = 0; s < i; ++s) { + const float a = s_A[i][s]; +#pragma unroll + for (int col = 0; col < 2 * MT; ++col) u[col] -= a * U[col][s]; + } + } +#pragma unroll + for (int col = 0; col < 2 * MT; ++col) U[col][i] = u[col]; + } + } + STAMP(6, U[2 * MT - 1][TT - 1] + U[0][TT - 1]); + // this step's update vectors (high, low parts) into the ring + { + u16* dp = p.ring_d + rid * p.d_slot + (long long)hv * RING * VD; + u16* dlp = p.ring_dlo ? p.ring_dlo + rid * p.dlo_slot + (long long)hv * RING * VD : nullptr; + const bool w0 = ok0 && wp + t0 < RING, w1 = ok1 && wp + t1 < RING; + const int ph0 = (base + wp + t0) & (RING - 1), ph1 = (base + wp + t1) & (RING - 1); +#pragma unroll + for (int col = 0; col < 2 * MT; ++col) { + const int v = vb + 16 * (col >> 1) + 8 * (col & 1) + r; + if (w0) { + const float d = pick8(U[col], t0); + const u16 hi = f2bf(d); + dp[ph0 * VD + v] = hi; + if (dlp) dlp[ph0 * VD + v] = f2bf(d - bf2f(hi)); + } + if (w1) { + const float d = pick8(U[col], t1); + const u16 hi = f2bf(d); + dp[ph1 * VD + v] = hi; + if (dlp) dlp[ph1 * VD + v] = f2bf(d - bf2f(hi)); + } + } + } + // output O[v, t] = e^{G_t} hw_q + sum_{s <= t} U[v, s] F[s][t] + float O[MT][4]; + { + float f0[8], f1[8]; +#pragma unroll + for (int s = 0; s < 8; ++s) { + f0[s] = s_F[s][t0]; + f1[s] = s_F[s][t1]; + } +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int h = 0; h < 2; ++h) { + float a0 = 0.f, a1 = 0.f; +#pragma unroll + for (int s = 0; s < TT; ++s) { + a0 += U[2 * m + h][s] * f0[s]; + a1 += U[2 * m + h][s] * f1[s]; + } + O[m][2 * h] = eG0 * HQ[m][2 * h] + a0; + O[m][2 * h + 1] = eG1 * HQ[m][2 * h + 1] + a1; + } + } + + // gated RMSNorm over the head's 128 columns of the BF16-rounded replay output, then static FP8 + float xb[MT][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < 4; ++i) xb[m][i] = bfround(O[m][i]); +#if STORE_O +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int t = (i & 1) ? t1 : t0; + const int v = vb + 16 * m + 8 * (i >> 1) + r; + if (t < TT) p.o_out[((row0 + t) * HV + hv) * VD + v] = f2bf(O[m][i]); + } +#endif + float p0 = 0.f, p1 = 0.f; +#pragma unroll + for (int m = 0; m < MT; ++m) { + p0 += xb[m][0] * xb[m][0] + xb[m][2] * xb[m][2]; + p1 += xb[m][1] * xb[m][1] + xb[m][3] * xb[m][3]; + } +#pragma unroll + for (int msk = 4; msk < 32; msk <<= 1) { + p0 += __shfl_xor_sync(FULL, p0, msk); + p1 += __shfl_xor_sync(FULL, p1, msk); + } + if (r == 0) { + s_norm[warp][t0] = p0; + s_norm[warp][t1] = p1; + } + __syncthreads(); + STAMP(7, 0.f); + float ss0 = 0.f, ss1 = 0.f; +#pragma unroll + for (int w = 0; w < NW; ++w) { + ss0 += s_norm[w][t0]; + ss1 += s_norm[w][t1]; + } + const float rstd0 = rsqrtf(ss0 / (float)VD + p.eps); + const float rstd1 = rsqrtf(ss1 / (float)VD + p.eps); +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int h = i >> 1; + const bool ok = (i & 1) ? ok1 : ok0; + if (ok) { + const int t = (i & 1) ? t1 : t0; + const int v = vb + 16 * m + 8 * h + r; + float yv = xb[m][i] * ((i & 1) ? rstd1 : rstd0) * nw[m][h]; + yv *= zg[m][i]; + p.y[(row0 + t) * p.y_stride + hv * VD + v] = f2e4m3(bfround(yv) * rscale); + } + } +#if PROF + if (lane == 0) { + const long long end = clock64(); + long long* pw = p.prof + ((long long)cta * NW + warp) * 9; +#pragma unroll + for (int i = 0; i < 8; ++i) pw[i] = st_[i]; + pw[8] = end; + } +#endif +} diff --git a/test/kernels/test_qwen3_5_gdn.py b/test/kernels/test_qwen3_5_gdn.py index ff3cf7d7..af274064 100644 --- a/test/kernels/test_qwen3_5_gdn.py +++ b/test/kernels/test_qwen3_5_gdn.py @@ -113,6 +113,68 @@ def test_conv_update_is_sglangs(views, block_n, bs): assert got.transpose(1, 2).is_contiguous() +@pytest.mark.parametrize("bs", REQUESTS) +@pytest.mark.parametrize("tile", [(128, 4), (64, 2)]) +@pytest.mark.parametrize("window_mode", [0, 1]) +def test_hoisted_conv_is_sglangs(window_mode, tile, bs): + """Slot gdn_conv: x rows loaded first (and with window_mode 1 one window row a channel) writes the + output, the conv states and the deduplicated windows bit for bit as SGLang's kernel.""" + from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_update + + from sgl_nvidia_mega.qwen3_5.gdn.conv import conv_update + + g = torch.Generator(device="cuda").manual_seed(window_mode + tile[0] + 10 * bs) + rows = _projection(bs, True, g) + weight = torch.randn(D, WIDTH, device="cuda", generator=g).to(torch.bfloat16) + state = torch.randn(SLOTS, D, WIDTH - 1, device="cuda", generator=g).to( + torch.bfloat16 + ) + shared = torch.randn(SLOTS, D, DRAFT + WIDTH - 2, device="cuda", generator=g).to( + torch.bfloat16 + ) + slots = torch.tensor(STATE_SLOTS[:bs], dtype=torch.int32, device="cuda") + if bs > 2: + slots[1] = -1 # a padded request + inter = torch.arange(bs, dtype=torch.int32, device="cuda") + x = rows.view(bs, DRAFT, D).transpose(1, 2) + + def windows(t): + st = t.stride() + return t.as_strided((SLOTS, DRAFT, D, WIDTH - 1), (st[0], 1, st[1], 1)) + + want_state, want_shared = state.clone(), shared.clone() + want = causal_conv1d_update( + x, + want_state, + weight, + None, + "silu", + conv_state_indices=slots, + intermediate_conv_window=windows(want_shared), + intermediate_state_indices=inter, + ) + got = conv_update( + x, + state, + weight, + None, + "silu", + slots, + windows(shared), + inter, + block_n=tile[0], + token_major_out=True, + hoist=True, + num_warps=tile[1], + window_mode=window_mode, + ) + live = (slots >= 0)[:, None, None] + assert torch.equal( + torch.where(live, got, 0), torch.where(live, want.contiguous(), 0) + ) + assert torch.equal(state, want_state) and torch.equal(shared, want_shared) + + @pytest.mark.parametrize("bs", REQUESTS) @pytest.mark.parametrize("views", [False, True]) def test_norm_fp8_is_sglangs_norm_then_static_quant(views, bs): @@ -259,6 +321,164 @@ def rand(*shape, dtype=torch.bfloat16): assert all(torch.equal(x, y) for x, y in zip(got_ring, want_ring)) +# -------------------------------------------------------------------- slot gdn_core + + +@pytest.mark.parametrize("bs", REQUESTS) +@pytest.mark.parametrize("history", [False, True]) +@pytest.mark.parametrize("pad", [False, True]) +def test_core_is_gdn_verify_numerically(pad, history, bs): + """Slot gdn_core's one launch against slot gdn_verify's verify (32-wide value tiles, + 2 warps) and FP8 gated norm on the same inputs: the replay output and the FP8 bytes + up to rounding, this step's ring entries (update and key high + low parts, log-decay) + up to FP32 summation order; a padded request writes zeros and leaves its ring alone. + History: an empty ring (BF16 checkpoints fold every commit) or 1-9 committed steps.""" + from sgl_nvidia_mega.qwen3_5.gdn import replayssm + from sgl_nvidia_mega.qwen3_5.gdn.norm_fp8 import rms_norm_gated_fp8 + from sgl_nvidia_mega.qwen3_5.sm103.gdn.core import core + + if pad and bs == 1: + pytest.skip("a padded single request leaves nothing to compare") + g = torch.Generator(device="cuda").manual_seed(100 * bs + 10 * history + pad) + ring_len, slots, tokens = 16, bs + 3, bs * DRAFT + + def rand(*shape, scale=1.0): + return (torch.randn(*shape, device="cuda", generator=g) * scale).to( + torch.bfloat16 + ) + + qkvz = rand(tokens, D + HV * V) # the convolved q / k / v rows and the gate z + qkv, z = qkvz[:, :D], qkvz[:, D:].view(tokens, HV, V) + ba = rand(tokens, 2 * HV) + b, a = ba[:, :HV], ba[:, HV:] + A_log = torch.rand(HV, device="cuda", generator=g) + dt_bias = torch.rand(HV, device="cuda", generator=g) + state = rand(slots, HV, V, K, scale=0.05) + ring = [ + rand(slots, HV, ring_len, V, scale=0.05), + rand(slots, H, ring_len, K, scale=0.09), + -torch.rand(slots, HV, ring_len, device="cuda", generator=g) * 0.5, + rand(slots, HV, ring_len, V, scale=1e-4), + rand(slots, H, ring_len, K, scale=1e-4), + ] + states = torch.randperm(slots, device="cuda", generator=g)[:bs].to(torch.int32) + replay = torch.randperm(slots, device="cuda", generator=g)[:bs] + if pad: + states[-1] = -1 + low = 1 if history else 0 + write_pos = torch.randint( + low, + ring_len - DRAFT + 1 if history else 1, + (slots,), + device="cuda", + generator=g, + ).to(torch.int32) + cache_base = torch.randint(0, ring_len, (slots,), device="cuda", generator=g) + cache_base = cache_base.to(torch.int32) + is_flush = torch.zeros(slots, dtype=torch.int8, device="cuda") + qsl = torch.arange(0, tokens + 1, DRAFT, dtype=torch.int32, device="cuda") + weight = rand(V) + scale = torch.tensor([0.02], device="cuda") + want_ring, got_ring = [t.clone() for t in ring], [t.clone() for t in ring] + o_want = torch.zeros(tokens, HV, V, dtype=torch.bfloat16, device="cuda") + replayssm.verify( + qkv[:, : H * K].view(tokens, H, K), + qkv[:, H * K : 2 * H * K].view(tokens, H, K), + qkv[:, 2 * H * K :].view(tokens, HV, V), + a, + b, + A_log, + dt_bias, + o_want, + state, + *want_ring, + None, + qsl, + states, + replay, + write_pos, + cache_base, + is_flush, + max_cache_len=ring_len, + max_spec_len=DRAFT, + scale=K**-0.5, + null_block_id=-1, + block_v=32, + num_warps=2, + ) + y_want = rms_norm_gated_fp8(o_want.view(-1, V), weight, 1e-6, z, scale) + o_got = torch.zeros_like(o_want) + y_got = torch.empty(tokens, HV * V, dtype=torch.float8_e4m3fn, device="cuda") + core( + qkv, + z, + a, + b, + A_log, + dt_bias, + state, + *got_ring, + states, + replay, + write_pos, + cache_base, + is_flush, + weight, + scale, + 1e-6, + y_got, + rows_per_request=DRAFT, + num_k_heads=H, + o_out=o_got, + ) + # the register cap of the widest grids (three CTAs an SM) leaves every byte as it was + capped = [t.clone() for t in ring] + y_cap = torch.empty_like(y_got) + core( + qkv, + z, + a, + b, + A_log, + dt_bias, + state, + *capped, + states, + replay, + write_pos, + cache_base, + is_flush, + weight, + scale, + 1e-6, + y_cap, + rows_per_request=DRAFT, + num_k_heads=H, + min_blocks=3, + ) + assert torch.equal(y_cap.view(torch.uint8), y_got.view(torch.uint8)) + assert all(torch.equal(c, g) for c, g in zip(capped, got_ring)) + live = torch.ones(tokens, dtype=torch.bool, device="cuda") + if pad: + live[-DRAFT:] = False + assert not y_got[~live].float().any() + ow, og = o_want[live].float(), o_got[live].float() + assert (og - ow).norm() <= 2e-3 * ow.norm() + same = ( + y_got.view(torch.uint8)[live] == y_want.view(torch.uint8).view(tokens, -1)[live] + ) + assert same.float().mean() >= 0.99 + yw, yg = y_want.view(tokens, -1)[live].float(), y_got[live].float() + assert (yg - yw).norm() <= 2e-2 * yw.norm() + # the ring: updates and keys as high + low, the log-decays as written + for hi, lo, tol in ((0, 3, 5e-3), (1, 4, 1e-5)): + w = want_ring[hi].float() + want_ring[lo].float() + o = got_ring[hi].float() + got_ring[lo].float() + changed = w - (ring[hi].float() + ring[lo].float()) + assert (o - w).norm() <= tol * max(changed.norm().item(), 1e-30) + assert torch.allclose(got_ring[2], want_ring[2], rtol=0, atol=1e-5) + + # ------------------------------------------------------------------- slot gdn_front FRONT_SLOTS = 5 diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 9c5e61f0..4ba61e49 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -214,6 +214,8 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): assert _on(plan.slots(7, every)) == sorted(plan.Slots.__struct_fields__) # two requests: also MKA's wide entries and their push into the hop assert _on(plan.slots(14, every)) == [ + "gdn_conv", + "gdn_core", "gdn_pdl", "gdn_verify", "hop", @@ -221,9 +223,22 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): "mka_push", "route", ] - # Wider requests take MKA's bitwise front path. - for rows in (28, 56, 112, 224, 448, 896): - assert _on(plan.slots(rows, every)) == ["gdn_pdl", "gdn_verify", "mka"] + # The one-launch GDN core and MKA front serve through 224 rows. + for rows in (28, 56, 112, 224): + assert _on(plan.slots(rows, every)) == [ + "gdn_conv", + "gdn_core", + "gdn_pdl", + "gdn_verify", + "mka", + ] + for rows in (448, 896): + assert _on(plan.slots(rows, every)) == [ + "gdn_conv", + "gdn_pdl", + "gdn_verify", + "mka", + ] # SGLang offers every multiple of 7 up to MAX_ROWS; every request count is a candidate offered = tuple(range(7, plan.MAX_ROWS + 1, 7)) assert plan.candidates(offered, 7) == offered From 378968624f1c4ad74ca9ec15ad81d003d32c63f6 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 15:16:40 -0700 Subject: [PATCH 05/62] feat: add shape-gated low-M projection kernels --- sgl_nvidia_mega/qwen3_5/attention/__init__.py | 24 +- sgl_nvidia_mega/qwen3_5/attention/mka.py | 24 +- sgl_nvidia_mega/qwen3_5/boundary/hop.py | 24 +- sgl_nvidia_mega/qwen3_5/gdn/__init__.py | 5 + sgl_nvidia_mega/qwen3_5/gdn/front.py | 2 +- sgl_nvidia_mega/qwen3_5/gdn/verify.py | 30 +- sgl_nvidia_mega/qwen3_5/lowm.py | 572 ++++++++++++++++++ sgl_nvidia_mega/qwen3_5/moe/__init__.py | 12 +- sgl_nvidia_mega/qwen3_5/plan.py | 45 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 23 +- .../qwen3_5/sm103/lowm/__init__.py | 300 +++++++++ .../qwen3_5/sm103/lowm/src/lowm.cu | 401 ++++++++++++ sgl_nvidia_mega/qwen3_5/spec.py | 16 + test/kernels/test_qwen3_5_hop.py | 46 +- test/kernels/test_qwen3_5_lowm.py | 204 +++++++ test/unit/test_qwen3_5_spec.py | 19 +- test/unit/test_qwen3_5_steps.py | 7 +- 17 files changed, 1726 insertions(+), 28 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/lowm.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/lowm/__init__.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/lowm/src/lowm.cu create mode 100644 test/kernels/test_qwen3_5_lowm.py diff --git a/sgl_nvidia_mega/qwen3_5/attention/__init__.py b/sgl_nvidia_mega/qwen3_5/attention/__init__.py index 288d37f9..0abe5d7a 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/attention/__init__.py @@ -122,9 +122,17 @@ def step(layer, x, st): from sgl_nvidia_mega.qwen3_5.attention import mka return mka.step(layer, x, st) - q, k, v, gate = layer._prepare_qkv_gate( - positions=st.positions, hidden_states=x, forward_batch=st.fb - ) + from sgl_nvidia_mega.qwen3_5 import lowm + + x8 = lowm.take_fp8(st, x) if st.slots.hop_fp8 else None + if st.slots.lowm_in or x8 is not None: + q, k, v, gate = lowm.prepare_qkv_gate( + layer, st.positions, x, x8, low_m=st.slots.lowm_in + ) + else: + q, k, v, gate = layer._prepare_qkv_gate( + positions=st.positions, hidden_states=x, forward_batch=st.fb + ) if st.slots.skv2: from sgl_nvidia_mega.qwen3_5.attention import split_kv @@ -133,5 +141,15 @@ def step(layer, x, st): o = layer.attn(q, k, v, st.fb) if layer.attn_output_gate: o = fused_sigmoid_mul(o, gate, inplace=True) + if st.slots.lowm_oproj: + from sgl_nvidia_mega.qwen3_5 import lowm + from sgl_nvidia_mega.qwen3_5.boundary import hop as hop_step + + hop = None + if st.slots.lowm_push and hop_step.takes_moe_input( + st.bound["boundary"].hop, layer, st, o.shape[0] + ): + hop = st.bound["boundary"].hop + return lowm.row_parallel(layer.o_proj, o, o.dtype, "o_proj", hop) out, _ = layer.o_proj(o) return out diff --git a/sgl_nvidia_mega/qwen3_5/attention/mka.py b/sgl_nvidia_mega/qwen3_5/attention/mka.py index cdf61a8d..4c0bc94c 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/mka.py +++ b/sgl_nvidia_mega/qwen3_5/attention/mka.py @@ -226,6 +226,7 @@ def _front_step(layer, x, st): from sglang.srt.layers.quantization.fp8_utils import flashinfer_bmm_fp8 from sglang.srt.runtime_context import get_parallel + from sgl_nvidia_mega.qwen3_5 import lowm from sgl_nvidia_mega.qwen3_5.attention import split_kv from sgl_nvidia_mega.qwen3_5.sm103.attention.gate_quant import gate_quant from sgl_nvidia_mega.qwen3_5.sm103.draft.front import ( @@ -235,7 +236,19 @@ def _front_step(layer, x, st): m = x.shape[0] attn = layer.attn width = layer.num_heads * layer.head_dim - qkv, _ = layer.qkv_proj(x) + # slot hop_fp8: qkv_proj reads the hop's E4M3 copy of x (no quantize kernel) + x8 = lowm.take_fp8(st, x) if st.slots.hop_fp8 else None + if x8 is not None: + lowm.engaged("hop_fp8", "qkv_proj bmm", m) + qkv = flashinfer_bmm_fp8( + x8, + layer.qkv_proj.weight, + layer.qkv_proj.input_scale, + layer.qkv_proj.weight_scale, + x.dtype, + ) + else: + qkv, _ = layer.qkv_proj(x) q_gate, k, v = qkv.split( [ 2 * width, @@ -269,6 +282,15 @@ def _front_step(layer, x, st): ) o = attn(q.reshape(m, width), None, None, st.fb).reshape(m, width) oq = gate_quant(o.contiguous(), gate.reshape(m, width), layer.o_proj.input_scale) + if st.slots.lowm_oproj: + # slot lowm_out: o_proj on the low-M launch, pushing into the hop with lowm_push + hop = None + if st.slots.lowm_push: + from sgl_nvidia_mega.qwen3_5.boundary import hop as boundary_hop + + if boundary_hop.takes_moe_input(st.bound["boundary"].hop, layer, st, m): + hop = st.bound["boundary"].hop + return lowm.row_parallel(layer.o_proj, oq, torch.bfloat16, "o_proj", hop) # the same allocation as RowParallelLinear's output, which the boundary all-reduces with use_symmetric_memory( get_parallel().tp_group, disabled=not is_allocation_symmetric() diff --git a/sgl_nvidia_mega/qwen3_5/boundary/hop.py b/sgl_nvidia_mega/qwen3_5/boundary/hop.py index 18a445d2..9b1c3458 100644 --- a/sgl_nvidia_mega/qwen3_5/boundary/hop.py +++ b/sgl_nvidia_mega/qwen3_5/boundary/hop.py @@ -311,11 +311,16 @@ def layer_input(hop, layer, hidden, st): and plan.fusions._finalize_eligible(plan, st.fb, pending.owed.m) ): return None + from sgl_nvidia_mega.qwen3_5 import lowm + owed, residual = stream.input(hidden) - norm, residual = _finalize( - hop, "layer input", owed, residual, fused_norm_gamma(plan.norm) + scale = lowm.layer_input_fp8_scale(layer, st, owed.m) + norm, residual, *fp8 = _finalize( + hop, "layer input", owed, residual, fused_norm_gamma(plan.norm), scale ) stream.write(residual) + if fp8: + lowm.stash_fp8(st, norm, fp8[0]) return norm @@ -360,11 +365,19 @@ def moe_input(hop, layer, hidden, st): and takes_moe_input(hop, layer, st, rows) ): return None + from sgl_nvidia_mega.qwen3_5 import lowm + _engage("MoE input", rows) - norm, residual = hop.all_reduce_residual_rms_norm( - local_contribution=part, residual=residual, gamma=fused_norm_gamma(plan.norm) + scale = lowm.moe_input_fp8_scale(layer, st) + norm, residual, *fp8 = hop.all_reduce_residual_rms_norm( + local_contribution=part, + residual=residual, + gamma=fused_norm_gamma(plan.norm), + **({} if scale is None else {"fp8_scale": scale}), ) stream.write(residual) + if fp8: + lowm.stash_fp8(st, norm, fp8[0]) return norm @@ -412,7 +425,7 @@ def _takes(hop, owed) -> bool: ) -def _finalize(hop, point, owed, residual, gamma): +def _finalize(hop, point, owed, residual, gamma, fp8_scale=None): _engage(point, owed.m) return hop.moe_finalize_all_reduce_rms_norm( routed_output=owed.routed_output, @@ -421,6 +434,7 @@ def _finalize(hop, point, owed, residual, gamma): gated_shared_output=owed.gated_shared_output, residual=residual, gamma=gamma, + **({} if fp8_scale is None else {"fp8_scale": fp8_scale}), ) diff --git a/sgl_nvidia_mega/qwen3_5/gdn/__init__.py b/sgl_nvidia_mega/qwen3_5/gdn/__init__.py index 1a653697..d04c767a 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/__init__.py @@ -162,6 +162,11 @@ def step(layer, x, st): ): opts = msgspec.structs.replace(opts, push=True) verify.engaged("gdn_push", x.shape[0]) + # slot lowm_push: out_proj's low-M launch pushes its rows into the hop instead + elif st.slots.lowm_push and hop.takes_moe_input( + st.bound["boundary"].hop, layer, st, x.shape[0] + ): + opts = msgspec.structs.replace(opts, lowm_push=True) if st.slots.gdn_front: from sgl_nvidia_mega.qwen3_5.gdn import front diff --git a/sgl_nvidia_mega/qwen3_5/gdn/front.py b/sgl_nvidia_mega/qwen3_5/gdn/front.py index f36e686f..58c80184 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/front.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/front.py @@ -304,4 +304,4 @@ def ring(t): input_scale=gdn.out_proj.input_scale, weight_scale=gdn.out_proj.weight_scale, ) - return _out_proj(gdn.out_proj, x_fp8, x.dtype) + return _out_proj(gdn.out_proj, x_fp8, x.dtype, st, opts.lowm_push) diff --git a/sgl_nvidia_mega/qwen3_5/gdn/verify.py b/sgl_nvidia_mega/qwen3_5/gdn/verify.py index d28e098a..178f258e 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/verify.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/verify.py @@ -171,6 +171,9 @@ class Options(msgspec.Struct, frozen=True): conv_rows: bool = ( False # the hoisted conv writing each channel's window row once (slot gdn_conv) ) + lowm_push: bool = ( + False # out_proj's low-M launch pushes into the hop (slot lowm_push) + ) STOCK = Options(views=False, block_n=256, block_v=64, num_warps=1, norm_fp8=False) @@ -282,7 +285,15 @@ def step(gdn, x: torch.Tensor, st, opts: Options = SLOT) -> torch.Tensor: logger.info( "NVIDIA megakernel: Qwen3.5 slot gdn_verify runs at %d rows", x.shape[0] ) - qkvz, ba = gdn._forward_input_proj(x) + from sgl_nvidia_mega.qwen3_5 import lowm + + x8 = lowm.take_fp8(st, x) if st.slots.hop_fp8 else None + if st.slots.lowm_in: + qkvz, ba = lowm.in_proj(gdn, x, x8) + elif x8 is not None: + qkvz, ba = lowm.in_proj_fp8(gdn, x, x8) + else: + qkvz, ba = gdn._forward_input_proj(x) split = ( qwen3_5_gdn_prefill_projection_views if opts.views @@ -320,7 +331,7 @@ def step(gdn, x: torch.Tensor, st, opts: Options = SLOT) -> torch.Tensor: input_scale=gdn.out_proj.input_scale, weight_scale=gdn.out_proj.weight_scale, ) - return _out_proj(gdn.out_proj, x, core.dtype) + return _out_proj(gdn.out_proj, x, core.dtype, st, opts.lowm_push) z_shape = z.shape core = gdn.norm(core, z if opts.views else z.reshape(-1, z.shape[-1])) core = core.reshape(z_shape) @@ -328,10 +339,19 @@ def step(gdn, x: torch.Tensor, st, opts: Options = SLOT) -> torch.Tensor: return gdn.out_proj(core)[0] -def _out_proj(out_proj, fp8: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: +def _out_proj( + out_proj, fp8: torch.Tensor, dtype: torch.dtype, st=None, push: bool = False +) -> torch.Tensor: """``RowParallelLinear.forward`` of the FP8 ``out_proj`` on its pre-quantized input: ModelOpt's FlashInfer BMM with the static input scale as the activation scale (what - its own quantization returns), the TP-partial output in SGLang's symmetric pool.""" + its own quantization returns), the TP-partial output in SGLang's symmetric pool; with + slot lowm_out (``Slots.lowm_outproj``) the low-M launch instead (numeric), pushing into the hop with + ``push`` (slot lowm_push).""" + if st is not None and st.slots.lowm_outproj: + from sgl_nvidia_mega.qwen3_5 import lowm + + hop = st.bound["boundary"].hop if push else None + return lowm.row_parallel(out_proj, fp8, dtype, "out_proj", hop) from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) @@ -425,7 +445,7 @@ def ring(t): input_scale=gdn.out_proj.input_scale, weight_scale=gdn.out_proj.weight_scale, ) - return _out_proj(gdn.out_proj, x, qkv.dtype) + return _out_proj(gdn.out_proj, x, qkv.dtype, st, opts.lowm_push) def _core(attn, mixed_qkv, a, b, st, opts): diff --git a/sgl_nvidia_mega/qwen3_5/lowm.py b/sgl_nvidia_mega/qwen3_5/lowm.py new file mode 100644 index 00000000..9151e006 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/lowm.py @@ -0,0 +1,572 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Adapted from SGLang (Apache-2.0, SGLang Team): +# python/sglang/srt/models/qwen3_5.py (Qwen3_5GatedDeltaNet._forward_input_proj, +# Qwen3_5AttentionDecoderLayer.forward_prepare_cuda_fused) and python/sglang/srt/models/qwen2_moe.py +# (Qwen2MoeMLP.forward, Qwen2MoeSparseMoeBlock._gate_shared_output_out_of_place). +"""Slots ``lowm_in``, ``lowm_out`` and ``lowm_se``: a covered verify step's static per-tensor FP8 (W8A8) projections on one +low-M warp-MMA launch each (``sm103/lowm``), the weights loaded before the programmatic dependency wait. + +Numeric: the same E4M3 operands as SGLang's static quantize + FlashInfer/cuBLAS FP8 GEMM (BF16 activations are +quantized in registers with the module's static input scale, the quantize kernel's arithmetic: the same FP8 +values), FP32 sums in another order; BF16 projections stay BF16 x BF16 with FP32 sums. + +- ``lowm_in`` (``Slots.lowm_in``): GDN in_proj_qkvz with in_proj_ba's BF16 rows in the same launch (SGLang + runs in_proj_ba on its alternate stream), attention qkv_proj. +- ``lowm_out`` (``Slots.lowm_outproj`` / ``lowm_oproj``): GDN out_proj on gdn_verify's FP8 norm, attention + o_proj on the gated BF16 output; both write SGLang's symmetric pool, as the all-reduce after them reads it. +- ``hop_fp8`` (``Slots.hop_fp8``, with ``hop``): the hop's all-reduce + norm also writes the next projection's + E4M3 input with that projection's static scale (SGLang's quantization of the norm, byte for byte): the layer + input for in_proj_qkvz / qkv_proj, the MoE input for the shared expert's gate_up; the projection then reads + it (the low-M launch, or FlashInfer's BMM without the quantize kernel). +- ``lowm_se`` (``Slots.lowm_se``): the shared expert of the dual-stream deferred-finalize order in two launches + on the alternate stream: gate_up with the SwiGLU and down_proj's static FP8 in the epilogue (the shared gate's + logits as a BF16 segment of the same launch), then down_proj with its output times the gate's sigmoid. +""" + +from __future__ import annotations + +import logging + +import msgspec +import torch + +logger = logging.getLogger(__name__) + +# The SGLang names this module reads (test/unit/test_contract.py pins them) +SGLANG_READS = { + "srt/models/qwen3_5.py": ( + "Qwen3_5LinearDecoderLayer", + "Qwen3_5AttentionDecoderLayer", + "linear_attn", + "in_proj_qkvz", + "in_proj_ba", + "out_proj", + "qkv_proj", + "o_proj", + "attn_output_gate", + "q_size", + "kv_size", + "q_norm", + "k_norm", + "rotary_emb", + "num_heads", + "num_kv_heads", + "head_dim", + "alt_stream", + ), + "srt/models/qwen2_moe.py": ( + "mlp", + "shared_expert", + "shared_expert_gate", + "gate_up_proj", + "down_proj", + "act_fn", + ), + "srt/layers/quantization/modelopt_quant.py": ( + "ModelOptFp8LinearMethod", + "use_marlin", + "use_sm120_fp8", + "use_flashinfer_bmm", + '"input_scale"', + '"weight_scale"', + ), + "srt/layers/activation.py": ("SiluAndMul",), + "kernels/ops/attention/fused_qk_rmsnorm_rope_gate.py": ( + "fused_qk_gemma_rmsnorm_rope_gate", + ), + "srt/distributed/device_communicators/pynccl_allocator.py": ( + "use_symmetric_memory", + ), + "srt/layers/dp_attention.py": ("is_allocation_symmetric",), + "srt/runtime_context.py": ("get_parallel",), + "srt/layers/quantization/fp8_utils.py": ("flashinfer_bmm_fp8",), +} + +# The SGLang functions whose body this module mirrors, by the digest of their code +# (test/unit/test_contract.py recomputes them on the installed SGLang) +SGLANG_MIRRORS = { + "srt/models/qwen3_5.py:Qwen3_5AttentionDecoderLayer.forward_prepare_cuda_fused": "d10e23f5651990d0", +} + + +# Per (N8, N16, K, rows, E4M3 input): the tactic (k warps, FP8 tiles, BF16 tiles, CTAs a SM, cluster K split) +# measured fastest on GB300 at TP8 (results/trt-1022/l6/data/micro x1-x4: CUDA-graph chains with and without a +# non-PDL kernel between launches); other shapes take sm103/lowm.default_tactic +def _t(*v): + from sgl_nvidia_mega.qwen3_5.sm103.lowm import Tactic + + return Tactic(*v) + + +TACTICS: dict = { + (2560, 16, 4096, 14, False): _t(16, 4, 2, 1, 1), # in_proj_qkvz + in_proj_ba + (2560, 0, 4096, 14, False): _t(8, 4, 1, 1, 1), # qkv_proj + (2560, 16, 4096, 14, True): _t( + 16, 4, 2, 1, 1 + ), # in_proj_qkvz (E4M3 from the hop) + in_proj_ba + (2560, 0, 4096, 14, True): _t(16, 4, 1, 1, 1), # qkv_proj (E4M3 from the hop) + (4096, 0, 1024, 14, True): _t(16, 2, 1, 1, 1), # out_proj + (4096, 0, 1024, 28, True): _t(8, 2, 1, 1, 1), + (4096, 0, 1024, 14, False): _t(8, 4, 1, 1, 1), # o_proj + (4096, 0, 1024, 28, False): _t(4, 2, 1, 1, 1), + (256, 8, 4096, 7, False): _t( + 8, 2, 1, 1, 4 + ), # shared expert gate_up + SwiGLU, shared gate + (256, 8, 4096, 14, False): _t(16, 2, 1, 1, 4), + (256, 8, 4096, 28, False): _t(16, 2, 1, 1, 4), + (256, 8, 4096, 56, False): _t(8, 2, 1, 1, 4), + (256, 8, 4096, 7, True): _t( + 8, 2, 1, 1, 4 + ), # shared expert gate_up on the hop's E4M3 + (256, 8, 4096, 14, True): _t(8, 2, 1, 1, 4), + (256, 8, 4096, 28, True): _t(8, 2, 1, 1, 4), + (4096, 0, 128, 7, True): _t( + 2, 1, 1, 1, 1 + ), # shared expert down_proj x sigmoid(gate) + (4096, 0, 128, 14, True): _t(2, 1, 1, 1, 1), + (4096, 0, 128, 28, True): _t(2, 1, 1, 1, 1), + (4096, 0, 128, 56, True): _t(2, 1, 1, 1, 1), +} + +_ENGAGED: set = set() +HIDDEN = 4096 # the hop's row width + + +class Bound(msgspec.Struct, frozen=True): + """The low-M slots' build result (``st.bound["lowm"]``).""" + + lowm_in: bool = False # every GDN / attention FP8 input projection takes the launch + lowm_out: bool = ( + False # every GDN / attention FP8 output projection takes the launch + ) + se: bool = False # every shared expert takes the two launches + push: bool = False # the output projections' launches also push into the hop + hop_fp8: bool = False # the hop also writes the next projection's E4M3 input + # lowm_se: id(mlp) -> (gate_up rows paired for the SwiGLU, the shared gate zero-padded to 8 rows) + shared: dict = msgspec.field(default_factory=dict) + reason: str = "" + + +def engaged(slot: str, what: str, rows: int) -> None: + """Log once per (slot, projection, rows) that a low-M launch ran, as engagement evidence.""" + if (slot, what, rows) not in _ENGAGED: + _ENGAGED.add((slot, what, rows)) + logger.info( + "NVIDIA megakernel: Qwen3.5 slot %s runs %s at %d rows", slot, what, rows + ) + + +def _w8(linear) -> torch.Tensor: + """ModelOpt's FP8 weight [K, N] (column-major) as the row-major [N, K] it views.""" + return linear.weight.t() + + +def fp8_unsupported(linear) -> str: + """Why ``linear`` is not a static per-tensor FP8 linear on FlashInfer's BMM ("" when it is).""" + from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp8LinearMethod + + method = getattr(linear, "quant_method", None) + if ( + not isinstance(method, ModelOptFp8LinearMethod) + or method.use_marlin + or method.use_sm120_fp8 + or not linear.use_flashinfer_bmm + ): + return "not a ModelOpt FP8 linear on FlashInfer's BMM" + w = _w8(linear) + if ( + w.dtype != torch.float8_e4m3fn + or not w.is_contiguous() + or w.shape[1] % 64 + or w.shape[0] % 8 + ): + return f"weight {w.dtype} {tuple(w.shape)}" + if linear.weight_scale.numel() != 1 or linear.input_scale.numel() != 1: + return "not per-tensor scales" + if ( + linear.weight_scale.dtype != torch.float32 + or linear.input_scale.dtype != torch.float32 + ): + return "scales are not FP32" + if linear.bias is not None: + return "a bias" + return "" + + +def _bf16_unsupported(linear) -> str: + w = linear.weight + if getattr(linear, "quant_method", None) is not None and type( + linear.quant_method + ).__name__ != ("UnquantizedLinearMethod"): + return f"{type(linear.quant_method).__name__}" + if ( + w.dtype != torch.bfloat16 + or not w.is_contiguous() + or w.shape[1] % 64 + or w.shape[0] % 8 + ): + return f"weight {w.dtype} {tuple(w.shape)}" + if getattr(linear, "bias", None) is not None: + return "a bias" + return "" + + +def build(model, ctx) -> Bound: + """Slots ``lowm`` / ``lowm_se`` for every layer, or why a layer cannot take them; compiles every launch + of the covered widths within the slots' rows.""" + from sglang.srt.models.qwen3_5 import ( + Qwen3_5AttentionDecoderLayer, + Qwen3_5LinearDecoderLayer, + ) + + from sgl_nvidia_mega.qwen3_5 import plan + from sgl_nvidia_mega.qwen3_5.runtime import SLOTS + from sgl_nvidia_mega.qwen3_5.sm103 import lowm as sm + + want_in, want_out, want_se = ( + "lowm_in" in SLOTS, + "lowm_out" in SLOTS, + "lowm_se" in SLOTS, + ) + want_push = want_out and "lowm_push" in SLOTS and "hop" in SLOTS + want_fp8 = "hop_fp8" in SLOTS and "hop" in SLOTS + want = want_in or want_out + if not (want or want_se or want_fp8): + return Bound(reason="stock (SGLANG_NVIDIA_MEGAKERNEL_SLOTS)") + m = model.model + dev = torch.cuda.current_device() + layers = [(i, m.layers[i]) for i in range(m.start_layer, m.end_layer)] + launches = set() # (m, k, n8, n16, act, epi8, out2_f32) + if want: + for i, layer in layers: + if isinstance(layer, Qwen3_5LinearDecoderLayer): + gdn = layer.linear_attn + for lin in (gdn.in_proj_qkvz, gdn.out_proj): + reason = fp8_unsupported(lin) + if reason: + return Bound(reason=f"lowm: layer {i}: {reason}") + reason = _bf16_unsupported(gdn.in_proj_ba) + if reason: + return Bound(reason=f"lowm: layer {i} in_proj_ba: {reason}") + n8, k = _w8(gdn.in_proj_qkvz).shape + n16 = gdn.in_proj_ba.weight.shape[0] + for w in ctx.widths: + if want_in and w <= plan.LOWM_IN_MAX_ROWS: + launches.add((w, k, n8, n16, 1, 0, False)) + if want_fp8 and w <= plan.HOP_MAX_ROWS: + launches.add((w, k, n8, n16, 0, 0, False)) + if want_out and w <= plan.LOWM_OUTPROJ_MAX_ROWS: + n, ko = _w8(gdn.out_proj).shape + launches.add((w, ko, n, 0, 0, 0, False)) + elif isinstance(layer, Qwen3_5AttentionDecoderLayer): + attn = layer + if not attn.attn_output_gate: + return Bound(reason=f"lowm: layer {i}: no output gate") + for lin in (attn.qkv_proj, attn.o_proj): + reason = fp8_unsupported(lin) + if reason: + return Bound(reason=f"lowm: layer {i}: {reason}") + for w in ctx.widths: + if want_in and w <= plan.LOWM_IN_MAX_ROWS: + n, k = _w8(attn.qkv_proj).shape + launches.add((w, k, n, 0, 1, 0, False)) + if want_fp8 and w <= plan.HOP_MAX_ROWS: + launches.add((w, k, n, 0, 0, 0, False)) + if want_out and w <= plan.LOWM_OPROJ_MAX_ROWS: + n, k = _w8(attn.o_proj).shape + launches.add((w, k, n, 0, 1, 0, False)) + # MKA's front path hands o_proj its gate's static FP8 output + launches.add((w, k, n, 0, 0, 0, False)) + shared = {} + if want_se: + for i, layer in layers: + mlp = getattr(layer, "mlp", None) + se = getattr(mlp, "shared_expert", None) + gate = getattr(mlp, "shared_expert_gate", None) + if se is None or gate is None: + return Bound(reason=f"lowm_se: layer {i}: no gated shared expert") + for lin in (se.gate_up_proj, se.down_proj): + reason = fp8_unsupported(lin) + if reason: + return Bound(reason=f"lowm_se: layer {i}: {reason}") + if type(se.act_fn).__name__ != "SiluAndMul": + return Bound(reason=f"lowm_se: layer {i}: {type(se.act_fn).__name__}") + gw = gate.weight + if ( + gw.dtype != torch.bfloat16 + or gw.shape[0] != 1 + or getattr(gate, "bias", None) is not None + ): + return Bound( + reason=f"lowm_se: layer {i}: shared gate {gw.dtype} {tuple(gw.shape)}" + ) + gu = _w8(se.gate_up_proj) + if gu.shape[0] % 16: + return Bound(reason=f"lowm_se: layer {i}: gate_up rows {gu.shape[0]}") + gate8 = torch.zeros((8, gw.shape[1]), dtype=gw.dtype, device=gw.device) + gate8[:1].copy_(gw) + shared[id(mlp)] = (sm.swiglu_pairs(gu), gate8) + n, k = gu.shape + nd, kd = _w8(se.down_proj).shape + for w in ctx.widths: + if w <= plan.LOWM_SE_MAX_ROWS: + launches.add((w, k, n, 8, 1, 1, True)) + if want_fp8 and w <= plan.HOP_MAX_ROWS: + launches.add((w, k, n, 8, 0, 1, True)) + launches.add((w, kd, nd, 0, 0, 2, False)) + for w, k, n8, n16, act, epi8, f32 in sorted(launches): + tactic = TACTICS.get((n8, n16, k, w, act == 0)) + sm.prepare(dev, w, k, n8, n16, act, epi8, f32, tactic=tactic) + if ( + want_push + and n8 == HIDDEN + and not n16 + and epi8 == 0 + and w <= plan.HOP_MAX_ROWS + ): + sm.prepare(dev, w, k, n8, n16, act, epi8, f32, tactic=tactic, push=True) + logger.info( + "NVIDIA megakernel: Qwen3.5 low-M launches built: %s", + sorted((w, n8, n16, k) for w, k, n8, n16, *_ in launches), + ) + return Bound( + lowm_in=want_in, + lowm_out=want_out, + se=want_se, + push=want_push, + hop_fp8=want_fp8, + shared=shared, + ) + + +def _gemm(rows: int, **kwargs): + from sgl_nvidia_mega.qwen3_5.sm103 import lowm as sm + + fp8_in = kwargs.get("x8") is not None and kwargs.get("w8") is not None + xa = kwargs["x8"] if fp8_in else kwargs["x16"] + n8 = 0 if kwargs.get("w8") is None else kwargs["w8"].shape[0] + n16 = 0 if kwargs.get("w16") is None else kwargs["w16"].shape[0] + return sm.gemm(tactic=TACTICS.get((n8, n16, xa.shape[1], rows, fp8_in)), **kwargs) + + +def in_proj(gdn, x: torch.Tensor, x8=None): + """``Qwen3_5GatedDeltaNet._forward_input_proj`` in one launch: (qkvz, ba) of the normed rows ``x`` (with + ``x8``, in_proj_qkvz reads the hop's E4M3 rows).""" + rows = x.shape[0] + qkvz_lin, ba_lin = gdn.in_proj_qkvz, gdn.in_proj_ba + w8 = _w8(qkvz_lin) + qkvz = torch.empty((rows, w8.shape[0]), dtype=x.dtype, device=x.device) + ba = torch.empty((rows, ba_lin.weight.shape[0]), dtype=x.dtype, device=x.device) + engaged( + "lowm_in", "in_proj_qkvz+ba" + ("(e4m3 in)" if x8 is not None else ""), rows + ) + _gemm( + rows, + x8=x8, + x16=x, + w8=w8, + w_scale=qkvz_lin.weight_scale, + in_scale=qkvz_lin.input_scale, + w16=ba_lin.weight, + out=qkvz, + out2=ba, + ) + return qkvz, ba + + +def in_proj_fp8(gdn, x: torch.Tensor, x8: torch.Tensor): + """``Qwen3_5GatedDeltaNet._forward_input_proj`` of a captured step with in_proj_qkvz on the hop's E4M3 rows + (FlashInfer's BMM, no quantize kernel) and in_proj_ba on the alternate stream, as SGLang orders them.""" + from sglang.srt.layers.quantization.fp8_utils import flashinfer_bmm_fp8 + + lin = gdn.in_proj_qkvz + engaged("hop_fp8", "in_proj_qkvz bmm", x.shape[0]) + current = torch.cuda.current_stream() + alt = gdn.alt_stream + if alt is not None: + alt.wait_stream(current) + qkvz = flashinfer_bmm_fp8( + x8, lin.weight, lin.input_scale, lin.weight_scale, x.dtype + ) + if alt is None: + return qkvz, gdn.in_proj_ba(x)[0] + with torch.cuda.stream(alt): + ba, _ = gdn.in_proj_ba(x) + current.wait_stream(alt) + return qkvz, ba + + +def stash_fp8(st, x: torch.Tensor, x8: torch.Tensor) -> None: + """Slot hop_fp8: the hop's E4M3 copy ``x8`` of the norm rows ``x`` for their next projection.""" + st.scratch["lowm.fp8"] = (x.data_ptr(), x.shape[0], x8) + + +def take_fp8(st, x: torch.Tensor): + """The hop's E4M3 copy of ``x`` (slot hop_fp8), or None.""" + held = st.scratch.pop("lowm.fp8", None) + if held is None or held[0] != x.data_ptr() or held[1] != x.shape[0]: + return None + return held[2] + + +def layer_input_fp8_scale(layer, st, rows: int): + """The static input scale of ``layer``'s mixer projection that reads the layer input in this step (slot + hop_fp8), or None where none would read an E4M3 copy (gdn_front's and MKA's fused entries quantize + in their own launch).""" + from sglang.srt.models.qwen3_5 import Qwen3_5LinearDecoderLayer + + from sgl_nvidia_mega.qwen3_5.attention import mka + + if not st.slots.hop_fp8: + return None + if isinstance(layer, Qwen3_5LinearDecoderLayer): + if not st.slots.gdn_verify or st.slots.gdn_front: + return None + return layer.linear_attn.in_proj_qkvz.input_scale + if st.slots.mka and rows <= mka.FUSED_MAX_ROWS: + return None + return layer.qkv_proj.input_scale + + +def moe_input_fp8_scale(layer, st): + """The shared expert's gate_up input scale where its low-M launch reads the MoE input (slot hop_fp8).""" + if not (st.slots.hop_fp8 and st.slots.lowm_se) or st.slots.route_shared: + return None + return layer.mlp.shared_expert.gate_up_proj.input_scale + + +def _symmetric_empty( + rows: int, n: int, like: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: + """A TP-partial output in SGLang's symmetric pool (what RowParallelLinear allocates under it).""" + from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, + ) + from sglang.srt.layers.dp_attention import is_allocation_symmetric + from sglang.srt.runtime_context import get_parallel + + with use_symmetric_memory( + get_parallel().tp_group, disabled=not is_allocation_symmetric() + ): + return torch.empty((rows, n), dtype=dtype, device=like.device) + + +def row_parallel( + linear, x: torch.Tensor, dtype: torch.dtype, what: str, hop=None +) -> torch.Tensor: + """``RowParallelLinear.forward`` of a static FP8 ``linear`` (no reduce, no bias) on BF16 ``x`` (quantized + in registers) or its E4M3 input: the TP-partial output [rows, N] in the symmetric pool. With ``hop`` + (slot lowm_push, where the MoE input's all-reduce runs the hop) the launch also pushes the rows into the + hop's mailbox and the all-reduce only consumes them.""" + rows = x.shape[0] + w8 = _w8(linear) + out = _symmetric_empty(rows, w8.shape[0], x, dtype) + engaged("lowm_push" if hop is not None else "lowm_out", what, rows) + act = {"x8": x} if x.dtype == torch.float8_e4m3fn else {"x16": x} + _gemm( + rows, + w8=w8, + w_scale=linear.weight_scale, + in_scale=linear.input_scale, + out=out, + push_site=None if hop is None else hop.site, + **act, + ) + if hop is not None: + hop.mark_pushed(out) + return out + + +def prepare_qkv_gate( + layer, positions: torch.Tensor, x: torch.Tensor, x8=None, low_m: bool = True +): + """``Qwen3_5AttentionDecoderLayer.forward_prepare_cuda_fused`` with qkv_proj on the low-M launch (slot + lowm_in; with ``x8`` on the hop's E4M3 rows), or FlashInfer's BMM on ``x8`` (slot hop_fp8).""" + from sglang.kernels.ops.attention.fused_qk_rmsnorm_rope_gate import ( + fused_qk_gemma_rmsnorm_rope_gate, + ) + from sglang.srt.layers.quantization.fp8_utils import flashinfer_bmm_fp8 + + rows = x.shape[0] + lin = layer.qkv_proj + if low_m: + w8 = _w8(lin) + qkv = torch.empty((rows, w8.shape[0]), dtype=x.dtype, device=x.device) + engaged("lowm_in", "qkv_proj" + ("(e4m3 in)" if x8 is not None else ""), rows) + act = {"x8": x8} if x8 is not None else {"x16": x} + _gemm( + rows, + w8=w8, + w_scale=lin.weight_scale, + in_scale=lin.input_scale, + out=qkv, + **act, + ) + else: + engaged("hop_fp8", "qkv_proj bmm", rows) + qkv = flashinfer_bmm_fp8( + x8, lin.weight, lin.input_scale, lin.weight_scale, x.dtype + ) + if layer.attn_output_gate: + q_gate, k, v = qkv.split( + [layer.q_size * 2, layer.kv_size, layer.kv_size], dim=-1 + ) + else: + q_gate, k, v = qkv.split([layer.q_size, layer.kv_size, layer.kv_size], dim=-1) + q_out, k_out, gate_out = fused_qk_gemma_rmsnorm_rope_gate( + q_gate, + k, + layer.q_norm.weight.data, + layer.k_norm.weight.data, + layer.rotary_emb.cos_sin_cache, + positions, + layer.q_norm.variance_epsilon, + layer.num_heads, + layer.num_kv_heads, + layer.head_dim, + layer.rotary_emb.rotary_dim, + has_gate=layer.attn_output_gate, + mrope_axis_map=(layer.rotary_emb.axis_map if positions.dim() == 2 else None), + ) + q = q_out.view(rows, -1) + k = k_out.view(rows, -1) + gate = gate_out.view(rows, -1) if gate_out is not None else None + return q, k, v, gate + + +def shared_expert(mlp, x: torch.Tensor, bound: Bound, x8=None) -> torch.Tensor: + """The gated shared expert (``Qwen2MoeMLP.forward`` then ``_gate_shared_output_out_of_place``) in two + launches: [rows, hidden] BF16 (with ``x8``, gate_up reads the hop's E4M3 rows).""" + rows = x.shape[0] + se = mlp.shared_expert + gu, dn = se.gate_up_proj, se.down_proj + pairs, gate8 = bound.shared[id(mlp)] + inter = pairs.shape[0] // 2 + h8 = torch.empty((rows, inter), dtype=torch.float8_e4m3fn, device=x.device) + logits = torch.empty((rows, 8), dtype=torch.float32, device=x.device) + engaged("lowm_se", "shared_expert" + ("(e4m3 in)" if x8 is not None else ""), rows) + _gemm( + rows, + x8=x8, + x16=x, + w8=pairs, + w_scale=gu.weight_scale, + in_scale=gu.input_scale, + w16=gate8, + out=h8, + out2=logits, + swiglu_scale=dn.input_scale, + ) + w8 = _w8(dn) + out = torch.empty((rows, w8.shape[0]), dtype=x.dtype, device=x.device) + _gemm( + rows, + x8=h8, + w8=w8, + w_scale=dn.weight_scale, + in_scale=dn.input_scale, + out=out, + gate=logits[:, 0], + ) + return out diff --git a/sgl_nvidia_mega/qwen3_5/moe/__init__.py b/sgl_nvidia_mega/qwen3_5/moe/__init__.py index 8e2ec90d..d57ea7bf 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/moe/__init__.py @@ -157,9 +157,15 @@ def _dual_stream(mlp, x, use_fused_gate: bool, defer: bool, st): # routed FC2 stays on the current stream, the finalize's PDL dependency routed = _router_experts(mlp, x, True, st) with torch.cuda.stream(mlp.alt_stream): - shared = mlp._forward_shared_experts(x, apply_gate=False) - if shared is not None: - shared = mlp._gate_shared_output_out_of_place(x, shared) + if st.slots.lowm_se: + from sgl_nvidia_mega.qwen3_5 import lowm + + x8 = lowm.take_fp8(st, x) if st.slots.hop_fp8 else None + shared = lowm.shared_expert(mlp, x, st.bound["lowm"], x8) + else: + shared = mlp._forward_shared_experts(x, apply_gate=False) + if shared is not None: + shared = mlp._gate_shared_output_out_of_place(x, shared) current.wait_stream(mlp.alt_stream) return routed, shared shared = ( diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index 932dd70c..67a5d5ff 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -21,7 +21,11 @@ nvidia/qwen35-agentx-gb300, ``mka`` the fused attention prologue and epilogue around it (up to 8 rows; with four query heads a rank, its wide entries up to 14 rows and past them its front path, bit for bit SGLang's, at every width), and ``mka_push`` (with -``hop``) has the epilogue push its rows into the hop. +``hop``) has the epilogue push its rows into the hop. ``lowm_in`` / ``lowm_out`` (numeric) run the static FP8 input / output projections on a +low-M warp-MMA launch with the weights loaded before the dependency wait, up to each one's +measured rows (``LOWM_*_MAX_ROWS``), and with ``lowm_push`` (with ``hop``) push their rows into the hop; +``lowm_se`` (numeric) the shared expert in two such launches. ``hop_fp8`` (with ``hop``) has the hop +also write the next projection's static FP8 input, which it then reads without a quantize kernel. """ from __future__ import annotations @@ -53,6 +57,14 @@ MKA_MAX_ROWS = 8 # with four query heads a rank: the wide entries, then the front path (attention/mka.py) MKA_WIDE_MAX_ROWS = MAX_ROWS +# slots lowm / lowm_se: the rows up to which each low-M launch measured faster than SGLang's quantize + +# cuBLAS path on GB300 (results/trt-1022/l6): the input projections (in_proj_qkvz + in_proj_ba, qkv_proj), +# the output projections (out_proj, o_proj) and the shared expert. The wider 28/ +# 56-row launches regressed the target step, so they stay on SGLang's stock path. +LOWM_IN_MAX_ROWS = 14 +LOWM_OUTPROJ_MAX_ROWS = 14 +LOWM_OPROJ_MAX_ROWS = 14 +LOWM_SE_MAX_ROWS = 14 class Slots(msgspec.Struct, frozen=True): @@ -78,6 +90,15 @@ class Slots(msgspec.Struct, frozen=True): skv2: bool = False # the full attention's core through SKV2 (attention/split_kv.py) mka: bool = False # the fused attention prologue and epilogue (attention/mka.py) mka_push: bool = False # the epilogue pushes its rows into the hop + # the input projections on the low-M launch (lowm.py; numeric) + lowm_in: bool = False + lowm_outproj: bool = False # GDN out_proj on the low-M launch (lowm.py; numeric) + lowm_oproj: bool = False # attention o_proj on the low-M launch (lowm.py; numeric) + lowm_se: bool = False # the shared expert in two low-M launches (lowm.py; numeric) + # the low-M output projections push their rows into the hop (lowm.py; numeric) + lowm_push: bool = False + # the hop also writes the next projection's static FP8 input (lowm.py; bitwise) + hop_fp8: bool = False class Facts(msgspec.Struct, frozen=True): @@ -101,6 +122,11 @@ class Facts(msgspec.Struct, frozen=True): mka: bool = False # the fused prologue and epilogue for every full-attention layer mka_wide: bool = False # with its wide entries (four query heads a rank) mka_push: bool = False # with the epilogue's push into the hop + lowm_in: bool = False # the low-M input projections built for every layer + lowm_out: bool = False # the low-M output projections built for every layer + lowm_se: bool = False # the low-M shared expert built for every MoE layer + lowm_push: bool = False # with the output projections' push into the hop + hop_fp8: bool = False # the hop's FP8 copy of the norm for the next projection def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> Slots: @@ -109,6 +135,7 @@ def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> route = facts.route and rows <= ROUTE_MAX_ROWS hop = facts.hop and rows <= HOP_MAX_ROWS mka = facts.mka and rows <= (MKA_WIDE_MAX_ROWS if facts.mka_wide else MKA_MAX_ROWS) + lowm_se = facts.lowm_se and rows <= LOWM_SE_MAX_ROWS return Slots( route=route, mkm=route @@ -127,11 +154,25 @@ def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> and facts.gdn_pdl and facts.gdn_front and rows == GDN_FRONT_ROWS, - route_shared=route and facts.route_shared and rows <= ROUTE_SHARED_ROWS, + # At one request lowm_se replaces both route_shared's shared stage and + # SGLang's shared-expert tail. + route_shared=route + and facts.route_shared + and rows <= ROUTE_SHARED_ROWS + and not lowm_se, # SKV2 attends one sequence skv2=facts.skv2 and rows <= tokens_per_request, mka=mka, mka_push=mka and facts.mka_push and hop, + lowm_in=facts.lowm_in and rows <= LOWM_IN_MAX_ROWS, + lowm_outproj=facts.lowm_out and rows <= LOWM_OUTPROJ_MAX_ROWS, + lowm_oproj=facts.lowm_out and rows <= LOWM_OPROJ_MAX_ROWS, + lowm_se=lowm_se, + lowm_push=facts.lowm_out + and facts.lowm_push + and hop + and rows <= max(LOWM_OUTPROJ_MAX_ROWS, LOWM_OPROJ_MAX_ROWS), + hop_fp8=hop and facts.hop_fp8 and rows <= HOP_MAX_ROWS, ) diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index edab5f39..b328575a 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -33,6 +33,11 @@ "skv2", "mka", "mka_push", + "lowm_in", + "lowm_out", + "lowm_se", + "lowm_push", + "hop_fp8", ) # Steps of more than one request (TP8's decode graphs of batch sizes 2 to 128), covered only # when requested in the same variable: covering them costs the one-request step about @@ -165,6 +170,13 @@ def slot_reasons() -> list[str]: f"{SLOTS_ENV}: slot {name!r} takes slot {base!r}'s long-KV graph " "variant's place from its length up" ) + if "hop_fp8" in SLOTS and "hop" not in SLOTS: + reasons.append(f"{SLOTS_ENV}: slot 'hop_fp8' is an output of slot 'hop'") + if "lowm_push" in SLOTS and not {"lowm_out", "hop"} <= SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'lowm_push' pushes slot 'lowm_out''s rows into slot " + "'hop''s mailbox" + ) if "mka_push" in SLOTS and not {"mka", "hop"} <= SLOTS: reasons.append( f"{SLOTS_ENV}: slot 'mka_push' pushes slot 'mka''s epilogue into slot " @@ -199,7 +211,7 @@ def build(model, ctx): import torch from sglang.srt.runtime_context import get_parallel - from sgl_nvidia_mega.qwen3_5 import attention, boundary, gdn, moe + from sgl_nvidia_mega.qwen3_5 import attention, boundary, gdn, lowm, moe from sgl_nvidia_mega.qwen3_5 import plan as slot_plan from sgl_nvidia_mega.qwen3_5.runtime import parity from sgl_nvidia_mega.qwen3_5.runtime.executor import Qwen35Verify @@ -231,6 +243,7 @@ def build(model, ctx): "gdn": gdn.build(model, pctx), "boundary": boundary.build(model, pctx), "attention": attention.build(model, pctx), + "lowm": lowm.build(model, pctx), } facts = slot_plan.Facts( route=bound["moe"].route, @@ -248,6 +261,11 @@ def build(model, ctx): mka=bound["attention"].mka, mka_wide=bound["attention"].mka_wide, mka_push=bound["attention"].mka_push, + lowm_in=bound["lowm"].lowm_in, + lowm_out=bound["lowm"].lowm_out, + lowm_se=bound["lowm"].se, + lowm_push=bound["lowm"].push, + hop_fp8=bound["lowm"].hop_fp8, ) covered = slot_plan.covered(pctx.widths, facts, pctx.tokens_per_request) slots = slot_plan.plan(covered, facts, pctx.tokens_per_request) @@ -279,7 +297,8 @@ def build(model, ctx): f"NVIDIA megakernel: requested slots {missing} built no width (MoE: " f"{bound['moe'].reason or 'built'}; GDN: {bound['gdn'].reason or 'built'}; " f"boundary: {bound['boundary'].reason or 'built'}; " - f"attention: {bound['attention'].reason or 'built'})" + f"attention: {bound['attention'].reason or 'built'}; " + f"low-M: {bound['lowm'].reason or 'built'})" ) checker = ( parity.Parity(covered, torch.cuda.current_device(), model) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/lowm/__init__.py b/sgl_nvidia_mega/qwen3_5/sm103/lowm/__init__.py new file mode 100644 index 00000000..a280f49c --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/lowm/__init__.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Low-M projections (``src/lowm.cu``): ``C = A @ W^T`` for up to 64 rows on warp MMA, the weights loaded +into registers before the programmatic dependency wait and the activations after it. + +A launch has an FP8 segment (E4M3 weights with a per-tensor scale; E4M3 activations, or BF16 ones quantized +in registers with the static input scale exactly as SGLang's static quantize kernel) and a BF16 segment +(BF16 weights and activations), each a whole number of CTAs. The FP8 segment's epilogue is a BF16 output +(``acc * s_in * s_w``), a SwiGLU on (8 gate rows, 8 up rows) pairs written as E4M3 with a static scale, or a +BF16 output times ``sigmoid(gate)`` per row. Numeric against SGLang's quantize + cuBLAS path: the same E4M3 +operands, FP32 sums in another order. +""" + +from __future__ import annotations + +import ctypes +from functools import cache +from pathlib import Path +from typing import NamedTuple, Optional + +import torch + +from sgl_nvidia_mega.common import jit + +SOURCE = Path(__file__).parent / "src" / "lowm.cu" +MAX_ROWS = 64 +CHUNK = 64 + + +class Params(ctypes.Structure): + """Params of src/lowm.cu.""" + + _fields_ = [ + ("x8", ctypes.c_void_p), + ("ld8", ctypes.c_int64), + ("x16", ctypes.c_void_p), + ("ld16", ctypes.c_int64), + ("w8", ctypes.c_void_p), + ("s_w", ctypes.c_void_p), + ("s_in", ctypes.c_void_p), + ("s_out", ctypes.c_void_p), + ("w16", ctypes.c_void_p), + ("out", ctypes.c_void_p), + ("ldo", ctypes.c_int64), + ("out2", ctypes.c_void_p), + ("ldo2", ctypes.c_int64), + ("gate", ctypes.c_void_p), + ("ldg", ctypes.c_int64), + ("hop_mc", ctypes.c_void_p), + ("hop_uc", ctypes.c_void_p), + ("hop_ctrl", ctypes.c_void_p), + ("hop_rank", ctypes.c_int), + ("hop_world", ctypes.c_int), + ("hop_m_cap", ctypes.c_int), + ] + + +class Tactic(NamedTuple): + """``kw`` warps split K (one 64-column chunk at a time); ``nt8`` / ``nt16`` groups of 8 weight rows a + CTA of the FP8 / BF16 segment; ``min_blocks`` CTAs a SM for the launch bounds; ``ksplit`` CTAs of a + cluster split K further and sum their partials through distributed shared memory in rank order.""" + + kw: int + nt8: int + nt16: int = 1 + min_blocks: int = 1 + ksplit: int = 1 + + +# the compiled launches' register and local-memory use (checks and tuning) +KERNEL_INFO: dict = {} + + +@cache +def _kernel( + device_index: int, + m: int, + k: int, + tactic: Tactic, + ctas8: int, + ctas16: int, + act: int, + epi8: int, + out2_f32: bool, + trigger: bool, + push: bool = False, +) -> jit.Kernel: + major, minor = torch.cuda.get_device_capability(device_index) + defines = dict( + M_ROWS=m, + K_DIM=k, + KW=tactic.kw, + NT8=tactic.nt8, + NT16=tactic.nt16, + CTAS8=ctas8, + CTAS16=ctas16, + ACT=act, + EPI8=epi8, + OUT2_F32=int(out2_f32), + TRIGGER=int(trigger), + MIN_BLOCKS=tactic.min_blocks, + KSPLIT=tactic.ksplit, + PUSH=int(push), + ) + source = "".join(f"#define {key} {value}\n" for key, value in defines.items()) + headers = None + if push: + from sgl_nvidia_mega.qwen3_5.sm103 import hop + + headers = hop.headers() + cubin = jit.compile_source( + source + SOURCE.read_text(), + "q35_lowm", + f"sm_{major}{minor}a", + ("-std=c++17", "-default-device"), + headers, + ) + nc = max(tactic.nt8 if ctas8 else 0, tactic.nt16 if ctas16 else 0) * 8 + shared = tactic.kw * ((m + 15) // 16) * 16 * nc * 4 + if tactic.ksplit > 1: + shared += m * nc * 4 + kernel = jit.load(cubin, "q35_lowm", (32 * tactic.kw, 1, 1), shared) + cu = jit.driver() + attr = cu.CUfunction_attribute + regs = jit.check( + cu.cuFuncGetAttribute(attr.CU_FUNC_ATTRIBUTE_NUM_REGS, kernel.function), "regs" + )[0] + local = jit.check( + cu.cuFuncGetAttribute(attr.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, kernel.function), + "lmem", + )[0] + KERNEL_INFO[(m, k, tactic, ctas8, ctas16, act, epi8, push)] = ( + int(regs), + int(local), + shared, + ) + return kernel + + +def default_tactic(m: int, n8: int, n16: int, k: int, swiglu: bool = False) -> Tactic: + """A tactic for ``m`` rows of an ``[n8 | n16, k]`` launch: each lane prefetches at most 256 bytes of + weights, warps take at least one chunk, and the grid fits two CTAs a SM where it can.""" + chunks = k // CHUNK + best = None + for nt in (1, 2, 4, 8): + if swiglu and nt % 2: + continue + if (n8 and n8 % (8 * nt)) or (n16 and n16 % (8 * min(nt, n16 // 8))): + continue + nt16 = min(nt, n16 // 8) if n16 else 1 + ctas = (n8 // (8 * nt) if n8 else 0) + (n16 // (8 * nt16) if n16 else 0) + for kw in (1, 2, 4, 8, 16): + if chunks % kw: + continue + per_lane = (chunks // kw) * max( + nt * 16 if n8 else 0, nt16 * 32 if n16 else 0 + ) + if per_lane > 256: + continue + waves = -(-ctas // (2 * 148)) + key = (waves, chunks // kw > 8, -min(ctas, 2 * 148), chunks // kw, -kw) + if best is None or key < best[0]: + best = (key, Tactic(kw, nt, nt16, 1)) + if best is None: + raise ValueError(f"no low-M tactic for N={n8}+{n16}, K={k}") + return best[1] + + +def check_launch( + m: int, n8: int, n16: int, k: int, tactic: Tactic, swiglu: bool = False +) -> str: + """Why ``tactic`` cannot run this launch ("" when it can).""" + if not 0 < m <= MAX_ROWS: + return f"{m} rows" + if k % CHUNK or (k // CHUNK) % (tactic.kw * tactic.ksplit): + return f"K={k} over {tactic.kw} warps x {tactic.ksplit} CTAs" + if tactic.ksplit not in (1, 2, 4, 8): + return f"a cluster of {tactic.ksplit}" + if n8 and n8 % (8 * tactic.nt8): + return f"N8={n8} over {tactic.nt8} tiles" + if swiglu and tactic.nt8 % 2: + return "a SwiGLU segment takes an even number of tiles a CTA" + if n16 and n16 % (8 * tactic.nt16): + return f"N16={n16} over {tactic.nt16} tiles" + return "" + + +def prepare( + device_index: int, + m: int, + k: int, + n8: int, + n16: int, + act: int, + epi8: int = 0, + out2_f32: bool = False, + tactic: Optional[Tactic] = None, + trigger: bool = True, + push: bool = False, +) -> Tactic: + """Compile a launch's kernel ahead of its first call (no compile inside CUDA-graph capture); + returns its tactic.""" + if tactic is None: + tactic = default_tactic(m, n8, n16, k, epi8 == 1) + reason = check_launch(m, n8, n16, k, tactic, epi8 == 1) + if reason: + raise ValueError(f"low-M launch: {reason}") + ctas8 = n8 // (8 * tactic.nt8) * tactic.ksplit if n8 else 0 + ctas16 = n16 // (8 * tactic.nt16) * tactic.ksplit if n16 else 0 + _kernel( + device_index, m, k, tactic, ctas8, ctas16, act, epi8, out2_f32, trigger, push + ) + return tactic + + +def gemm( + *, + x8: Optional[torch.Tensor] = None, + x16: Optional[torch.Tensor] = None, + w8: Optional[torch.Tensor] = None, + w_scale: Optional[torch.Tensor] = None, + in_scale: Optional[torch.Tensor] = None, + w16: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, + out2: Optional[torch.Tensor] = None, + swiglu_scale: Optional[torch.Tensor] = None, + gate: Optional[torch.Tensor] = None, + tactic: Optional[Tactic] = None, + trigger: bool = True, + push_site=None, +): + """``[x] @ [w8; w16]^T`` into ``out`` (FP8 segment) and ``out2`` (BF16 segment); returns (out, out2). + + x8 [M, K] E4M3 or x16 [M, K] BF16 (unit column stride, 16-byte aligned rows); w8 [N8, K] E4M3 row-major + with the per-tensor ``w_scale`` and the static ``in_scale``; w16 [N16, K] BF16 row-major (takes x16). + out: BF16 [M, N8], or E4M3 [M, N8 / 2] with ``swiglu_scale`` (w8 rows paired 8 gate / 8 up, see + ``swiglu_pairs``); ``gate`` [M] FP32 logits multiply the BF16 output by their sigmoid. out2: BF16 or + FP32 [M, N16]. ``push_site`` (a sm103/hop ``Site``; BF16 output of hidden width only): the output rows + are also pushed into the decode hop's mailbox, whose all-reduce must come next.""" + xa = x16 if x16 is not None else x8 + m, k = xa.shape + n8 = 0 if w8 is None else w8.shape[0] + n16 = 0 if w16 is None else w16.shape[0] + act = 0 if (n8 and x8 is not None) else 1 + epi8 = 1 if swiglu_scale is not None else (2 if gate is not None else 0) + if tactic is None: + tactic = default_tactic(m, n8, n16, k, epi8 == 1) + reason = check_launch(m, n8, n16, k, tactic, epi8 == 1) + if reason: + raise ValueError(f"low-M launch: {reason}") + ctas8 = n8 // (8 * tactic.nt8) * tactic.ksplit if n8 else 0 + ctas16 = n16 // (8 * tactic.nt16) * tactic.ksplit if n16 else 0 + out2_f32 = out2 is not None and out2.dtype == torch.float32 + push = push_site is not None + if push and (epi8 != 0 or n16 or out is None or out.stride(0) % 8): + raise ValueError("the hop push takes a plain BF16 FP8-segment output") + kernel = _kernel( + xa.device.index, m, k, tactic, ctas8, ctas16, act, epi8, out2_f32, trigger, push + ) + p = Params( + x8=0 if x8 is None else x8.data_ptr(), + ld8=0 if x8 is None else x8.stride(0), + x16=0 if x16 is None else x16.data_ptr(), + ld16=0 if x16 is None else x16.stride(0), + w8=0 if w8 is None else w8.data_ptr(), + s_w=0 if w_scale is None else w_scale.data_ptr(), + s_in=0 if in_scale is None else in_scale.data_ptr(), + s_out=0 if swiglu_scale is None else swiglu_scale.data_ptr(), + w16=0 if w16 is None else w16.data_ptr(), + out=0 if out is None else out.data_ptr(), + ldo=0 if out is None else out.stride(0), + out2=0 if out2 is None else out2.data_ptr(), + ldo2=0 if out2 is None else out2.stride(0), + gate=0 if gate is None else gate.data_ptr(), + ldg=0 if gate is None else gate.stride(0), + ) + if push: + p.hop_mc, p.hop_uc, p.hop_ctrl = push_site.mc, push_site.uc, push_site.ctrl + p.hop_rank, p.hop_world, p.hop_m_cap = ( + push_site.rank, + push_site.world, + push_site.m_cap, + ) + jit.launch( + kernel, + (ctas8 + ctas16, 1, 1), + [p], + torch.cuda.current_stream(xa.device).cuda_stream, + cluster=(tactic.ksplit, 1, 1) if tactic.ksplit > 1 else None, + ) + return out, out2 + + +def swiglu_pairs(gate_up: torch.Tensor) -> torch.Tensor: + """Rows ``[gate (I); up (I)]`` -> ``[gate 0..7; up 0..7; gate 8..15; up 8..15; ...]``.""" + rows, k = gate_up.shape + inter = rows // 2 + gate = gate_up[:inter].reshape(inter // 8, 8, k) + up = gate_up[inter:].reshape(inter // 8, 8, k) + return torch.stack((gate, up), dim=1).reshape(rows, k).contiguous() diff --git a/sgl_nvidia_mega/qwen3_5/sm103/lowm/src/lowm.cu b/sgl_nvidia_mega/qwen3_5/sm103/lowm/src/lowm.cu new file mode 100644 index 00000000..081bdcd3 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/lowm/src/lowm.cu @@ -0,0 +1,401 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Low-M projections of a Qwen3.5 verify step on warp MMA with register-prefetched weights: C = A @ W^T for +// M_ROWS (<= 64) rows. The work layout follows TensorRT-LLM's lowm_mma_gemm (Qwen3.5 decode, Apache-2.0): +// a CTA owns NT groups of 8 weight rows and splits K over its KW warps; each lane holds 16 contiguous K +// columns of a 64-wide chunk for its weight row and its activation rows, which are the MMA fragments once +// the reduction order inside the chunk is permuted the same way for both operands; the warps reduce through +// shared memory in warp order. +// +// Two weight segments, each a whole number of CTAs: blocks [0, CTAS8) take the FP8 segment (E4M3 weights, +// per-tensor weight scale), blocks [CTAS8, CTAS8 + CTAS16) the BF16 segment (BF16 activations). +// ACT 0: the FP8 segment reads E4M3 activations quantized by the producer; ACT 1: BF16 activations quantized +// in registers with the static input scale, with the arithmetic of SGLang's static quantize kernel +// (1 / s by div.full, multiply, clamp to +-448, cvt.rn.satfinite: the same E4M3 values). +// EPI8 0: BF16 out = acc * (s_in * s_w); 1: SwiGLU on (8 gate rows, 8 up rows) tile pairs of BF16-rounded +// outputs, written as E4M3 with the static scale s_out; 2: BF16 out times sigmoid(gate[row]) (the shared +// expert's gate, FP32 logits). +// PDL: weights and scales load before griddepcontrol.wait, activations after it; TRIGGER 1 releases the +// dependent launch once the weight loads are issued (a dependent waits for this grid before reading). + +#include +#ifndef PUSH +#define PUSH 0 +#endif +#if PUSH +#include "q35_hop.cuh" +#endif + +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef long long int64_t; + +#ifndef M_ROWS +#error "M_ROWS" +#endif +#ifndef KSPLIT +#define KSPLIT 1 +#endif +#define MT ((M_ROWS + 15) / 16) +#define CHUNKS (K_DIM / 64) +#define CPW (CHUNKS / (KW * KSPLIT)) +#define NTMAX (NT8 > NT16 ? NT8 : NT16) +#define NC (NTMAX * 8) + +struct Params { + const uint8_t* x8; // [M, K] E4M3 (ACT 0) + int64_t ld8; + const __nv_bfloat16* x16; // [M, K] BF16 (ACT 1, BF16 segment) + int64_t ld16; + const uint8_t* w8; // [N8, K] E4M3, row-major + const float* s_w; + const float* s_in; + const float* s_out; // EPI8 1 + const __nv_bfloat16* w16; // [N16, K] BF16, row-major + void* out; // FP8 segment: BF16 [M, N8] (EPI8 0, 2) or E4M3 [M, N8 / 2] (EPI8 1) + int64_t ldo; + void* out2; // BF16 segment: BF16 or FP32 [M, N16] + int64_t ldo2; + const float* gate; // EPI8 2: FP32 logits, row stride ldg + int64_t ldg; + // PUSH: the decode hop's mailbox site (q35hop::Site); the BF16 output rows are also pushed into it + void* hop_mc; + void* hop_uc; + void* hop_ctrl; + int hop_rank, hop_world, hop_m_cap; +}; + +__device__ __forceinline__ void pdl_wait() { asm volatile("griddepcontrol.wait;" ::: "memory"); } +__device__ __forceinline__ void pdl_launch() { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); } + +#ifndef WLD +#define WLD 2 +#endif +#ifndef ALD +#define ALD 1 +#endif +#ifndef QCLAMP +#define QCLAMP 0 +#endif +__device__ __forceinline__ uint4 ld_weight(const void* p) { + uint4 v; +#if WLD == 0 + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p)); +#elif WLD == 1 + asm volatile("ld.global.cs.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p)); +#else + asm volatile("ld.global.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p)); +#endif + return v; +} +// activations: not volatile, so the compiler may schedule them among the MMAs; their pointers come from +// launder() after the dependency wait, which keeps every such load behind it +__device__ __forceinline__ uint4 ld_act(const void* p) { + uint4 v; +#if ALD == 0 + asm("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p)); +#else + asm("ld.global.ca.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p)); +#endif + return v; +} +template +__device__ __forceinline__ T* launder(T* p) { + asm volatile("" : "+l"(p)); + return p; +} + +__device__ __forceinline__ void mma_e4m3(float (&c)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, + uint32_t b0, uint32_t b1) { + asm( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} +__device__ __forceinline__ void mma_bf16(float (&c)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, + uint32_t b0, uint32_t b1) { + asm( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// SGLang's static quantize (Triton): y * (1.0 / s) (div.full.f32), clamp to [-448, 448], E4M3 rn satfinite +__device__ __forceinline__ float inverse(float s) { + float r; + asm("div.full.f32 %0, %1, %2;" : "=f"(r) : "f"(1.0f), "f"(s)); + return r; +} +__device__ __forceinline__ uint16_t e4m3x2(float lo, float hi) { + uint16_t r; + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(r) : "f"(hi), "f"(lo)); + return r; +} +// the clamp of SGLang's kernel; satfinite alone gives the same E4M3 for every non-NaN input +#if QCLAMP +__device__ __forceinline__ float qclamp(float x) { return fminf(fmaxf(x, -448.f), 448.f); } +#else +__device__ __forceinline__ float qclamp(float x) { return x; } +#endif +// two bf16x2 words (element 0 in the low half) -> four E4M3, element 0 in the low byte +__device__ __forceinline__ uint32_t quant4(uint32_t w01, uint32_t w23, float inv) { + const float f0 = qclamp(__uint_as_float(w01 << 16) * inv); + const float f1 = qclamp(__uint_as_float(w01 & 0xffff0000u) * inv); + const float f2 = qclamp(__uint_as_float(w23 << 16) * inv); + const float f3 = qclamp(__uint_as_float(w23 & 0xffff0000u) * inv); + return uint32_t(e4m3x2(f0, f1)) | (uint32_t(e4m3x2(f2, f3)) << 16); +} + +// 16 activation columns of a row as E4M3 (zero beyond M_ROWS) +// rows past M_ROWS read the last row: no branch around the loads; their sums are never stored +__device__ __forceinline__ uint4 act8(const uint8_t* x8, int64_t ld8, const __nv_bfloat16* x16, int64_t ld16, + int row, int col, float inv) { + row = row < M_ROWS ? row : M_ROWS - 1; +#if ACT == 0 + return ld_act(x8 + row * ld8 + col); +#else + const __nv_bfloat16* src = x16 + row * ld16 + col; + const uint4 lo = ld_act(src), hi = ld_act(src + 8); + return make_uint4(quant4(lo.x, lo.y, inv), quant4(lo.z, lo.w, inv), quant4(hi.x, hi.y, inv), + quant4(hi.z, hi.w, inv)); +#endif +} + +__device__ __forceinline__ float bf16r(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float SUM(const float* red, int row, int col) { + float s = 0.f; +#pragma unroll + for (int w = 0; w < KW; ++w) s += red[(w * MT * 16 + row) * NC + col]; + return s; +} + +#if KSPLIT > 1 +__device__ __forceinline__ void cluster_sync() { + asm volatile("barrier.cluster.arrive.release.aligned;\n\tbarrier.cluster.wait.acquire.aligned;" ::: "memory"); +} +__device__ __forceinline__ unsigned cluster_rank() { + unsigned r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} +// the cluster's sum of an output: the KSPLIT CTAs' partial sums in rank order (distributed shared memory) +__device__ __forceinline__ float TOTAL(const float* part, int row, int col) { + const unsigned a = static_cast(__cvta_generic_to_shared(part + row * NC + col)); + float s = 0.f; +#pragma unroll + for (int r = 0; r < KSPLIT; ++r) { + unsigned ra; + float v; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(ra) : "r"(a), "r"(r)); + asm volatile("ld.shared::cluster.f32 %0, [%1];" : "=f"(v) : "r"(ra) : "memory"); + s += v; + } + return s; +} +#define T_SUM(row, col) TOTAL(part, row, col) +#else +#define T_SUM(row, col) SUM(red, row, col) +#endif + +extern "C" __global__ void __launch_bounds__(KW * 32, MIN_BLOCKS) q35_lowm(const Params p) { + extern __shared__ float red[]; // [KW][MT * 16][NC], then (KSPLIT > 1) the CTA's sums [M_ROWS][NC] + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, q = lane & 3; +#if KSPLIT > 1 + const int rank = cluster_rank(); + float* part = red + KW * MT * 16 * NC; +#else + const int rank = 0; +#endif + const int group = blockIdx.x / KSPLIT; // the CTA's weight-row group; rank: its K slice + const int k0 = (rank * KW + warp) * CPW * 64 + 16 * q; + float acc[MT][NTMAX][4]; +#pragma unroll + for (int t = 0; t < MT; ++t) +#pragma unroll + for (int n = 0; n < NTMAX; ++n) +#pragma unroll + for (int i = 0; i < 4; ++i) acc[t][n][i] = 0.f; + const bool fp8 = blockIdx.x < CTAS8; + int ntiles; + if (fp8) { +#if CTAS8 > 0 + ntiles = NT8; + const int n0 = group * NT8 * 8; + uint4 w[CPW][NT8]; +#pragma unroll + for (int c = 0; c < CPW; ++c) +#pragma unroll + for (int n = 0; n < NT8; ++n) w[c][n] = ld_weight(p.w8 + (int64_t)(n0 + 8 * n + g) * K_DIM + k0 + 64 * c); +#if ACT == 1 + const float s_in = *p.s_in; +#endif +#if TRIGGER + pdl_launch(); +#endif + pdl_wait(); + const uint8_t* x8 = launder(p.x8); + const __nv_bfloat16* x16 = launder(p.x16); +#if ACT == 1 + const float inv = inverse(s_in); +#else + const float inv = 0.f; +#endif +#pragma unroll + for (int c = 0; c < CPW; ++c) { + uint4 a[MT][2]; +#pragma unroll + for (int t = 0; t < MT; ++t) { + a[t][0] = act8(x8, p.ld8, x16, p.ld16, 16 * t + g, k0 + 64 * c, inv); + a[t][1] = (16 * t + 8 < M_ROWS) ? act8(x8, p.ld8, x16, p.ld16, 16 * t + g + 8, k0 + 64 * c, inv) + : make_uint4(0, 0, 0, 0); + } +#pragma unroll + for (int t = 0; t < MT; ++t) +#pragma unroll + for (int n = 0; n < NT8; ++n) { + mma_e4m3(acc[t][n], a[t][0].x, a[t][1].x, a[t][0].y, a[t][1].y, w[c][n].x, w[c][n].y); + mma_e4m3(acc[t][n], a[t][0].z, a[t][1].z, a[t][0].w, a[t][1].w, w[c][n].z, w[c][n].w); + } + } +#endif + } else { +#if CTAS16 > 0 + ntiles = NT16; + const int n0 = (group - CTAS8 / KSPLIT) * NT16 * 8; + uint4 w[CPW][NT16][2]; +#pragma unroll + for (int c = 0; c < CPW; ++c) +#pragma unroll + for (int n = 0; n < NT16; ++n) { + const __nv_bfloat16* src = p.w16 + (int64_t)(n0 + 8 * n + g) * K_DIM + k0 + 64 * c; + w[c][n][0] = ld_weight(src); + w[c][n][1] = ld_weight(src + 8); + } +#if TRIGGER + pdl_launch(); +#endif + pdl_wait(); + const __nv_bfloat16* x16 = launder(p.x16); +#pragma unroll + for (int c = 0; c < CPW; ++c) { + uint4 a[MT][2][2]; +#pragma unroll + for (int t = 0; t < MT; ++t) +#pragma unroll + for (int h = 0; h < 2; ++h) { + if (16 * t + 8 * h < M_ROWS) { + int row = 16 * t + g + 8 * h; + row = row < M_ROWS ? row : M_ROWS - 1; + const __nv_bfloat16* src = x16 + row * p.ld16 + k0 + 64 * c; + a[t][h][0] = ld_act(src); + a[t][h][1] = ld_act(src + 8); + } else { + a[t][h][0] = a[t][h][1] = make_uint4(0, 0, 0, 0); + } + } +#pragma unroll + for (int t = 0; t < MT; ++t) +#pragma unroll + for (int n = 0; n < NT16; ++n) { + // step j takes words 2j / 2j + 1 of the lane's 16 BF16 columns (logical k 2q+i / 2q+8+i) + mma_bf16(acc[t][n], a[t][0][0].x, a[t][1][0].x, a[t][0][0].y, a[t][1][0].y, w[c][n][0].x, w[c][n][0].y); + mma_bf16(acc[t][n], a[t][0][0].z, a[t][1][0].z, a[t][0][0].w, a[t][1][0].w, w[c][n][0].z, w[c][n][0].w); + mma_bf16(acc[t][n], a[t][0][1].x, a[t][1][1].x, a[t][0][1].y, a[t][1][1].y, w[c][n][1].x, w[c][n][1].y); + mma_bf16(acc[t][n], a[t][0][1].z, a[t][1][1].z, a[t][0][1].w, a[t][1][1].w, w[c][n][1].z, w[c][n][1].w); + } + } +#endif + } + // reduce the KW warps' partial sums in warp order +#pragma unroll + for (int t = 0; t < MT; ++t) +#pragma unroll + for (int n = 0; n < NTMAX; ++n) { + if (n < ntiles) { + float* dst = red + ((warp * MT * 16) + 16 * t + g) * NC + 8 * n + 2 * q; + dst[0] = acc[t][n][0]; + dst[1] = acc[t][n][1]; + dst[8 * NC] = acc[t][n][2]; + dst[8 * NC + 1] = acc[t][n][3]; + } + } + __syncthreads(); +#if KSPLIT > 1 + for (int i = threadIdx.x; i < M_ROWS * ntiles * 8; i += KW * 32) { + const int row = i / (ntiles * 8), col = i % (ntiles * 8); + part[row * NC + col] = SUM(red, row, col); + } + cluster_sync(); +#define OUT_START (rank * KW * 32 + threadIdx.x) +#define OUT_STEP (KW * 32 * KSPLIT) +#else +#define OUT_START (threadIdx.x) +#define OUT_STEP (KW * 32) +#endif + if (fp8) { +#if CTAS8 > 0 + const float alpha = *p.s_in * *p.s_w; +#if EPI8 == 1 + // SwiGLU: tile 2j = gate rows, tile 2j + 1 = the matching up rows; NT8 * 4 outputs a CTA + const float inv_out = inverse(*p.s_out); + uint8_t* out = reinterpret_cast(p.out); + for (int i = OUT_START; i < M_ROWS * NT8 * 4; i += OUT_STEP) { + const int row = i / (NT8 * 4), col = i % (NT8 * 4), pair = col >> 3, c8 = col & 7; + const float gt = bf16r(T_SUM(row, 16 * pair + c8) * alpha); + const float up = bf16r(T_SUM(row, 16 * pair + 8 + c8) * alpha); + const float h = bf16r(gt / (1.f + __expf(-gt)) * up); + const uint16_t v = e4m3x2(qclamp(h * inv_out), 0.f); + out[row * p.ldo + group * NT8 * 4 + col] = uint8_t(v & 0xff); + } +#elif PUSH + // the decode hop's push (its all-reduce then only consumes): one 8-column fragment a thread, pushed + // before it is stored + const q35hop::Site site{p.hop_mc, p.hop_uc, static_cast(p.hop_ctrl), p.hop_rank, p.hop_world, + p.hop_m_cap}; + const uint32_t stage = q35hop::stage_peek(site); + for (int i = OUT_START; i < M_ROWS * NT8; i += OUT_STEP) { + const int row = i / NT8, f = i % NT8; + float v[8]; +#pragma unroll + for (int c = 0; c < 8; ++c) v[c] = T_SUM(row, 8 * f + c) * alpha; + const uint4 frag = q35hop::pack8(v); + q35hop::push_frag(site, stage, row, group * NT8 + f, frag); + *reinterpret_cast(reinterpret_cast<__nv_bfloat16*>(p.out) + row * p.ldo + group * NT8 * 8 + 8 * f) = frag; + } +#else + __nv_bfloat16* out = reinterpret_cast<__nv_bfloat16*>(p.out); + for (int i = OUT_START; i < M_ROWS * NT8 * 8; i += OUT_STEP) { + const int row = i / (NT8 * 8), col = i % (NT8 * 8); + float v = T_SUM(row, col) * alpha; +#if EPI8 == 2 + const float gl = p.gate[row * p.ldg]; + v = bf16r(v) * (1.f / (1.f + __expf(-gl))); +#endif + out[row * p.ldo + group * NT8 * 8 + col] = __float2bfloat16_rn(v); + } +#endif +#endif + } else { +#if CTAS16 > 0 + for (int i = OUT_START; i < M_ROWS * NT16 * 8; i += OUT_STEP) { + const int row = i / (NT16 * 8), col = i % (NT16 * 8); + const float v = T_SUM(row, col); + const int col0 = (group - CTAS8 / KSPLIT) * NT16 * 8; +#if OUT2_F32 + reinterpret_cast(p.out2)[row * p.ldo2 + col0 + col] = v; +#else + reinterpret_cast<__nv_bfloat16*>(p.out2)[row * p.ldo2 + col0 + col] = __float2bfloat16_rn(v); +#endif + } +#endif + } +#if KSPLIT > 1 + cluster_sync(); // the peers' partial sums stay until every rank has read them +#endif +} diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index 8d79db35..4484c9ea 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -137,6 +137,22 @@ "mka.o_proj.rel_l2_vs_stock": 1.05, "mka.o_proj.max_abs_vs_stock": 2.0, }, + # the low-M launches against fp64 on the same E4M3 operands (test/kernels/test_qwen3_5_lowm.py) + "lowm_in": { + "lowm_in.in_proj_qkvz.rel_max_vs_stock": 1.05, + "lowm_in.in_proj_ba.rel_max_vs_stock": 1.05, + "lowm_in.qkv_proj.rel_max_vs_stock": 1.05, + }, + "lowm_out": { + "lowm_out.out_proj.rel_max_vs_stock": 1.05, + "lowm_out.o_proj.rel_max_vs_stock": 1.05, + }, + "lowm_se": {"lowm_se.shared_expert.rel_l2_vs_stock": 2e-2}, + # the pushed partial is lowm_out's; the all-reduce that consumes it is the fusion's (test_qwen3_5_hop.py) + "lowm_push": { + "lowm_push.out_proj.rel_max_vs_stock": 1.05, + "lowm_push.o_proj.rel_max_vs_stock": 1.05, + }, } diff --git a/test/kernels/test_qwen3_5_hop.py b/test/kernels/test_qwen3_5_hop.py index 39993e8d..dde595ea 100644 --- a/test/kernels/test_qwen3_5_hop.py +++ b/test/kernels/test_qwen3_5_hop.py @@ -13,7 +13,8 @@ And the GatedDeltaNet out_proj push (slot gdn_push, numeric): its partial against an fp64 reference on the same FP8 operands, at most 1.05 x the mean relative L2 and 2 x the max abs error of FlashInfer's BMM (SGLang's op), and its all-reduce, which only consumes, -bit for bit the fusion's on the same partial. +bit for bit the fusion's on the same partial. And the low-M launch's push (slot lowm_push): its partial is +the launch's without the push, and the all-reduce that only consumes it is the fusion's. Run with ``python test/kernels/test_qwen3_5_hop.py --num-gpu 4``. """ @@ -241,5 +242,48 @@ def test_out_proj_push_holds_to_fp64_and_its_all_reduce_only_consumes(rank, pair assert push[:, 1].max() <= 2 * stock[:, 1].max() +def test_lowm_push_is_its_partial_and_its_all_reduce_only_consumes(rank, pair): + """Slot lowm_push: the low-M out_proj / o_proj launch (sm103/lowm) with the hop push stores the + same partial as without it, and the all-reduce that only consumes it is the fusion's, bit for bit.""" + from sgl_nvidia_mega.qwen3_5.sm103 import lowm + + fusion, hop = pair + device = torch.device(f"cuda:{torch.cuda.current_device()}") + k = 1024 + with torch.inference_mode(): + for seed in range(2): + gen = torch.Generator(device=device).manual_seed(2000 * seed + rank) + w = _fp8(HIDDEN, k, device, gen, scale=0.5) + sw = torch.tensor([0.01], device=device) + sx = torch.tensor([0.02], device=device) + for m in (9, 14): + for bf16_in in (False, True): + x8 = _fp8(m, k, device, gen) + act = {"x16": x8.to(torch.bfloat16)} if bf16_in else {"x8": x8} + plain = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=device) + lowm.gemm(w8=w, w_scale=sw, in_scale=sx, out=plain, **act) + partial = torch.empty_like(plain) + lowm.gemm( + w8=w, + w_scale=sw, + in_scale=sx, + out=partial, + push_site=hop.site, + **act, + ) + hop.mark_pushed(partial) + x2 = _operands(m, device, rank, seed=seed) + got = hop.all_reduce_residual_rms_norm( + local_contribution=partial, + residual=x2["residual"], + gamma=x2["gamma"], + ) + assert torch.equal( + partial.view(torch.int16), plain.view(torch.int16) + ) + _same(got, _run(fusion, "all_reduce", dict(x2, partial=partial))) + torch.cuda.synchronize() + + if __name__ == "__main__": multigpu_pytest_main(__name__, __file__, num_gpus=(4,)) diff --git a/test/kernels/test_qwen3_5_lowm.py b/test/kernels/test_qwen3_5_lowm.py new file mode 100644 index 00000000..4db2aba7 --- /dev/null +++ b/test/kernels/test_qwen3_5_lowm.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""The low-M projection launches (slots lowm_in / lowm_out / lowm_se, sm103/lowm) on one SM100-family GPU +(GB300), numeric against SGLang's static quantize + FlashInfer FP8 BMM, eager and replayed from a CUDA graph: +- the in-register quantization is SGLang's static quantize bit for bit: a BF16-input launch equals, byte for + byte, the same launch on SGLang's E4M3 quantization of the rows; +- each TP8 projection (in_proj_qkvz with in_proj_ba's BF16 rows, qkv_proj, out_proj, o_proj) at the tactics + the slots use and at the default one, against fp64 on the same E4M3 operands: at most 1.05 x the max + relative error of FlashInfer's BMM plus one BF16 ulp of the output scale; +- the cluster K split equals the same tactic without it up to the FP32 sum order; +- the shared expert's two launches against SGLang's chain (quantize, gate_up BMM, SiluAndMul, quantize, + down BMM, fused_gate_sigmoid_mul).""" + +import pytest +import torch + +from sgl_nvidia_mega.qwen3_5.sm103 import lowm + +FP8 = torch.float8_e4m3fn + + +def _gpu_ready() -> bool: + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: + return False + try: + import flashinfer # noqa: F401 + from cuda.bindings import nvrtc # noqa: F401 + from sglang.kernels.ops.quantization.fp8_kernel import ( # noqa: F401 + static_quant_fp8, + ) + except ImportError: + return False + return True + + +pytestmark = pytest.mark.skipif(not _gpu_ready(), reason="needs an SM100-family GPU") + + +def _weight(g, n, k, std=0.05): + w = torch.randn(n, k, device="cuda", generator=g) * std + s = (w.abs().amax() / 448.0).float().reshape(1) + return (w / s).to(FP8).contiguous(), s + + +def _rows(g, m, k): + x = (torch.randn(m, k, device="cuda", generator=g) * 2).to(torch.bfloat16) + s = (x.float().abs().amax() / 448.0 * 1.1).reshape(1).float() + return x, s + + +def _bmm(x8, w8, s, sw): + from sglang.srt.layers.quantization.fp8_utils import flashinfer_bmm_fp8 + + return flashinfer_bmm_fp8(x8, w8.t(), s, sw, torch.bfloat16) + + +def _rel(a, ref): + return float((a.float() - ref).abs().max() / ref.abs().max()) + + +def _tactics(): + from sgl_nvidia_mega.qwen3_5.lowm import TACTICS + + return TACTICS + + +@pytest.mark.parametrize( + "n, n16, k, rows", + [ + (2560, 16, 4096, 14), + (2560, 0, 4096, 14), + (4096, 0, 1024, 14), + (4096, 0, 1024, 28), + (2560, 0, 4096, 28), + (4096, 0, 1024, 56), + ], +) +def test_projection_is_bmm_numerically(n, n16, k, rows): + from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 + + g = torch.Generator(device="cuda").manual_seed(rows * 7 + n16) + w8, sw = _weight(g, n, k) + w16 = ( + (torch.randn(n16, k, device="cuda", generator=g) * 0.02).to(torch.bfloat16) + if n16 + else None + ) + x, s = _rows(g, rows, k) + x8, _ = static_quant_fp8(x, s) + exact = (x8.double() @ w8.double().t()) * float(s) * float(sw) + ref = _bmm(x8, w8, s, sw) + bound = 1.05 * _rel(ref, exact) + 2.0**-8 + used = _tactics().get((n, n16, k, rows, False)) + for tactic in ([used] if used is not None else []) + [ + lowm.default_tactic(rows, n, n16, k), + None, + ]: + out_q = torch.empty(rows, n, dtype=torch.bfloat16, device="cuda") + out_8 = torch.empty_like(out_q) + out2 = ( + torch.empty(rows, n16, dtype=torch.bfloat16, device="cuda") if n16 else None + ) + lowm.gemm( + x16=x, + w8=w8, + w_scale=sw, + in_scale=s, + w16=w16, + out=out_q, + out2=out2, + tactic=tactic, + ) + torch.cuda.synchronize() + assert _rel(out_q, exact) <= bound, (tactic, _rel(out_q, exact), bound) + if tactic is not None: + # the same FP8 CTAs on SGLang's E4M3 rows: the quantization is SGLang's, byte for byte + lowm.gemm(x8=x8, w8=w8, w_scale=sw, in_scale=s, out=out_8, tactic=tactic) + torch.cuda.synchronize() + assert torch.equal(out_q.view(torch.int16), out_8.view(torch.int16)), tactic + if n16: + e2 = x.double() @ w16.double().t() + r2 = torch.nn.functional.linear(x, w16) + assert _rel(out2, e2) <= 1.05 * _rel(r2, e2) + 2.0**-8 + + +@pytest.mark.parametrize("rows", [14, 28]) +def test_cluster_split_is_the_same_sums(rows): + g = torch.Generator(device="cuda").manual_seed(rows) + w8, sw = _weight(g, 4096, 1024) + x, s = _rows(g, rows, 1024) + a = torch.empty(rows, 4096, dtype=torch.bfloat16, device="cuda") + b = torch.empty_like(a) + lowm.gemm( + x16=x, w8=w8, w_scale=sw, in_scale=s, out=a, tactic=lowm.Tactic(4, 4, 1, 1, 1) + ) + lowm.gemm( + x16=x, w8=w8, w_scale=sw, in_scale=s, out=b, tactic=lowm.Tactic(4, 4, 1, 1, 2) + ) + torch.cuda.synchronize() + d = (a.float() - b.float()).abs() + assert float(d.max()) <= float(a.float().abs().max()) * 2.0**-7 + + +@pytest.mark.parametrize("rows", [14, 28, 56]) +def test_shared_expert_is_sglang_numerically(rows): + from sgl_kernel import silu_and_mul + from sglang.kernels.ops.elementwise.elementwise import fused_gate_sigmoid_mul + from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 + + g = torch.Generator(device="cuda").manual_seed(100 + rows) + gu, sgu = _weight(g, 256, 4096) + dn, sdn = _weight(g, 4096, 128) + gate = (torch.randn(1, 4096, device="cuda", generator=g) * 0.02).to(torch.bfloat16) + x, s = _rows(g, rows, 4096) + s_h = torch.tensor([0.02], device="cuda") + # SGLang: quantize, gate_up BMM, SiluAndMul, quantize, down BMM, gate + q, _ = static_quant_fp8(x, s) + h = torch.empty(rows, 128, dtype=torch.bfloat16, device="cuda") + silu_and_mul(_bmm(q, gu, s, sgu), h) + q2, _ = static_quant_fp8(h, s_h) + ref = fused_gate_sigmoid_mul(x, gate.squeeze(0), _bmm(q2, dn, s_h, sdn)) + gate8 = torch.zeros(8, 4096, dtype=torch.bfloat16, device="cuda") + gate8[:1].copy_(gate) + t = _tactics() + h8 = torch.empty(rows, 128, dtype=FP8, device="cuda") + logits = torch.empty(rows, 8, dtype=torch.float32, device="cuda") + out = torch.empty(rows, 4096, dtype=torch.bfloat16, device="cuda") + + def run(): + lowm.gemm( + x16=x, + w8=lowm.swiglu_pairs(gu), + w_scale=sgu, + in_scale=s, + w16=gate8, + out=h8, + out2=logits, + swiglu_scale=s_h, + tactic=t.get((256, 8, 4096, rows, False)), + ) + lowm.gemm( + x8=h8, + w8=dn, + w_scale=sdn, + in_scale=s_h, + out=out, + gate=logits[:, 0], + tactic=t.get((4096, 0, 128, rows, True)), + ) + + run() + torch.cuda.synchronize() + # the hidden FP8 rows differ from SGLang's only where the FP32 sum order moves a BF16 rounding + assert float((h8.float() != q2.float()).float().mean()) < 0.02 + rel = float((out.float() - ref.float()).norm() / ref.float().norm()) + assert rel < 2e-2, rel + first = out.clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, first) diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 4ba61e49..6456a1a4 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -198,6 +198,11 @@ def test_the_slot_plan_covers_one_request_of_7_rows(): shared = dict(route=True, route_shared=True) assert plan.plan((7,), plan.Facts(**shared)) == {7: plan.Slots(**shared)} assert not plan.slots(7, plan.Facts(route_shared=True)).route_shared + # lowm_se replaces route_shared's shared stage where both could run. + assert _on(plan.slots(7, plan.Facts(**shared, lowm_se=True))) == [ + "lowm_se", + "route", + ] with pytest.raises(ValueError): plan.slots(plan.MAX_ROWS + 1, plan.Facts()) @@ -211,19 +216,27 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): every = plan.Facts( **{**dict.fromkeys(plan.Facts.__struct_fields__, True), "mkm_wide": False} ) - assert _on(plan.slots(7, every)) == sorted(plan.Slots.__struct_fields__) - # two requests: also MKA's wide entries and their push into the hop + assert _on(plan.slots(7, every)) == sorted( + set(plan.Slots.__struct_fields__) - {"route_shared"} + ) + # Two requests also take every measured low-M launch and the hop's FP8 copy. assert _on(plan.slots(14, every)) == [ "gdn_conv", "gdn_core", "gdn_pdl", "gdn_verify", "hop", + "hop_fp8", + "lowm_in", + "lowm_oproj", + "lowm_outproj", + "lowm_push", + "lowm_se", "mka", "mka_push", "route", ] - # The one-launch GDN core and MKA front serve through 224 rows. + # The 28/56-row low-M paths stay stock; GDN core and MKA serve through 224. for rows in (28, 56, 112, 224): assert _on(plan.slots(rows, every)) == [ "gdn_conv", diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 17794b2e..7bf31cb7 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -765,6 +765,7 @@ def _gdn_verify_setup(log, *, ring_len=16, temporal_dtype=torch.bfloat16, draft= fb=SimpleNamespace(req_pool_indices=k["req"]), binding=SimpleNamespace(tokens_per_request=draft), view=SimpleNamespace(gdn=meta, mamba=pools), + slots=plan.Slots(), scratch={}, ) return verify, gdn, st @@ -1092,7 +1093,7 @@ def norm(x, z, weight, scale, quant, rstd, eps): patch.object( verify, "_out_proj", - lambda linear, fp8, dtype: ( + lambda linear, fp8, dtype, *slot: ( log.append(("out_proj", linear, fp8, dtype)) or "partial" ), ), @@ -2436,7 +2437,9 @@ def test_the_hop_slot_runs_where_sglang_runs_its_cutedsl_kernel(cutedsl_gates, p for run in ("sglang", "plugin"): stream = ResidualStream(residual) fb = _verify_batch(stream) - st = SimpleNamespace(fb=fb, model=SimpleNamespace(norm=final)) + st = SimpleNamespace( + fb=fb, model=SimpleNamespace(norm=final), slots=plan.Slots() + ) if point == "MoE input": handle = layer.attn_boundary.finish(part, fb) x = ( From 332edbc52a40005aae7caf7412c52e9d14b2de98 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 15:22:24 -0700 Subject: [PATCH 06/62] feat: add bitwise wide MoE engines --- sgl_nvidia_mega/qwen3_5/moe/__init__.py | 83 +- sgl_nvidia_mega/qwen3_5/moe/mkm.py | 5 +- sgl_nvidia_mega/qwen3_5/moe/route.py | 189 ++- sgl_nvidia_mega/qwen3_5/plan.py | 53 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 92 +- sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py | 232 ++- sgl_nvidia_mega/qwen3_5/sm103/moe/router.py | 124 ++ .../qwen3_5/sm103/moe/src/mkm_moe_m128.cu | 1412 ++++++++++++++++ .../qwen3_5/sm103/moe/src/mkm_moe_m32.cu | 1346 +++++++++++++++ .../qwen3_5/sm103/moe/src/mkm_moe_m64.cu | 1448 +++++++++++++++++ .../qwen3_5/sm103/moe/src/mkm_moe_wide.cu | 1331 +++++++++++++++ .../qwen3_5/sm103/moe/src/route_topk.cu | 259 +++ .../qwen3_5/sm103/moe/src/router_front.cu | 12 +- sgl_nvidia_mega/qwen3_5/spec.py | 2 + test/kernels/test_qwen3_5_moe.py | 209 ++- test/unit/test_qwen3_5_runtime.py | 31 +- test/unit/test_qwen3_5_spec.py | 89 +- test/unit/test_qwen3_5_steps.py | 2 +- 18 files changed, 6852 insertions(+), 67 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m128.cu create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m32.cu create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m64.cu create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_wide.cu create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_topk.cu diff --git a/sgl_nvidia_mega/qwen3_5/moe/__init__.py b/sgl_nvidia_mega/qwen3_5/moe/__init__.py index d57ea7bf..07dcf050 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/moe/__init__.py @@ -12,7 +12,12 @@ alternate stream), otherwise the eager order. With ``defer`` (the boundary takes the finalize) it returns SGLang's ``MoeDeferredFinalize`` handoff, otherwise the summed rows. Slot ``route`` (``route.py``) runs the deferred routed experts on the router -front, and slot ``mkm`` (``mkm.py``) runs them in one launch on its outputs. In the +front, and slot ``mkm`` (``mkm.py``) runs them in one launch on its outputs; past route's +rows, slot ``route_topk`` (``route.py``) runs them on a materialized routing, or at 17 to +32 rows slot ``mkm_m32`` in one launch of mkm's 32-token engine, and at 33 to 64 slot +``mkm_m64`` on route_topk's routing in one launch of its 64-token engine (65 to 128: +``mkm_m128``, its 128-token engine; 129 to 256 / 257 to 512: ``mkm_m256`` / ``mkm_m512``, its +wide engines). In the dual-stream order of up to 8 rows, slot ``route_shared`` computes the shared expert's gate/up in the router front's launch; the alternate stream waits for that launch only and runs the rest of the shared expert (numeric). @@ -69,23 +74,59 @@ def build(model, ctx) -> route.Bound: - """Slots ``route`` and ``mkm`` for every MoE layer; the result is - ``st.bound["moe"]``.""" + """Slots ``route``, ``mkm``, ``route_topk`` and ``mkm_m32`` for every MoE layer; the + result is ``st.bound["moe"]``.""" + import msgspec + from sgl_nvidia_mega.qwen3_5 import plan from sgl_nvidia_mega.qwen3_5.runtime import SLOTS - if "route" not in SLOTS: - return route.Bound(reason="stock (SGLANG_NVIDIA_MEGAKERNEL_SLOTS)") - rows = plan.widest(ctx.widths, plan.ROUTE_MAX_ROWS) - if not rows: - return route.Bound(reason=f"no width of {list(ctx.widths)} within its rows") - return route.build( - model, - rows, - fused="mkm" in SLOTS, - shared="route_shared" in SLOTS, - wide="mkm_wide" in SLOTS, - ) + bound = route.Bound(reason="stock (SGLANG_NVIDIA_MEGAKERNEL_SLOTS)") + if "route" in SLOTS: + rows = plan.widest(ctx.widths, plan.ROUTE_MAX_ROWS) + if not rows: + bound = route.Bound( + reason=f"no width of {list(ctx.widths)} within its rows" + ) + else: + bound = route.build( + model, + rows, + fused="mkm" in SLOTS, + shared="route_shared" in SLOTS, + wide="mkm_wide" in SLOTS, + ) + if "route_topk" in SLOTS: + reason = route.topk_unsupported(model) + if reason: + return msgspec.structs.replace(bound, reason=f"route_topk: {reason}") + from sgl_nvidia_mega.qwen3_5.sm103.moe import router + + router.build_topk(torch.cuda.current_device()) + bound = msgspec.structs.replace(bound, topk=True) + if "route_gemm" in SLOTS: + router.build(torch.cuda.current_device(), router.MAX_TOKENS) + bound = msgspec.structs.replace(bound, gemm=True) + m32 = "mkm_m32" in SLOTS and plan.widest(ctx.widths, plan.MKM_M32_MAX_ROWS) > ( + plan.MKM_WIDE_MAX_ROWS + ) + m64 = "mkm_m64" in SLOTS and plan.widest(ctx.widths, plan.MKM_M64_MAX_ROWS) > ( + plan.MKM_M32_MAX_ROWS + ) + m128 = "mkm_m128" in SLOTS and plan.widest( + ctx.widths, plan.MKM_M128_MAX_ROWS + ) > (plan.MKM_M64_MAX_ROWS) + m256 = "mkm_m256" in SLOTS and plan.widest( + ctx.widths, plan.MKM_M256_MAX_ROWS + ) > (plan.MKM_M128_MAX_ROWS) + m512 = "mkm_m512" in SLOTS and plan.widest( + ctx.widths, plan.MKM_M512_MAX_ROWS + ) > (plan.MKM_M256_MAX_ROWS) + if m32 or m64 or m128 or m256 or m512: + bound = route.build_m32( + model, bound, m32=m32, m64=m64, m128=m128, m256=m256, m512=m512 + ) + return bound def step(layer, x: torch.Tensor, st, defer: bool): @@ -186,6 +227,18 @@ def _router_experts(mlp, x, defer: bool, st): if defer and st.slots.route: return route.routed(mlp, x, st.bound["moe"] if st.slots.mkm else None) + if defer and st.slots.mkm_m32: + return route.routed_m32(mlp, x, st.bound["moe"], gemm=st.slots.route_gemm) + if defer and st.slots.mkm_m64: + return route.routed_given(mlp, x, st.bound["moe"], "mkm_m64") + if defer and st.slots.mkm_m128: + return route.routed_given(mlp, x, st.bound["moe"], "mkm_m128") + if defer and st.slots.mkm_m256: + return route.routed_given(mlp, x, st.bound["moe"], "mkm_m256") + if defer and st.slots.mkm_m512: + return route.routed_given(mlp, x, st.bound["moe"], "mkm_m512") + if defer and st.slots.route_topk: + return route.routed_topk(mlp, x, gemm=st.slots.route_gemm) router_logits, _ = mlp.gate(x) topk_output = mlp.topk(x, router_logits) if defer: diff --git a/sgl_nvidia_mega/qwen3_5/moe/mkm.py b/sgl_nvidia_mega/qwen3_5/moe/mkm.py index 56cfb2b8..e3a1fa0a 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/mkm.py +++ b/sgl_nvidia_mega/qwen3_5/moe/mkm.py @@ -161,10 +161,10 @@ def routing_unsupported(mlp) -> str: return "" -def routed(mlp, maps: tuple, identity: torch.Tensor, logits, x_fp4, x_sf): +def routed(mlp, maps: tuple, identity: torch.Tensor, logits, x_fp4, x_sf, routing=None): """The deferred-finalize output of ``mlp``'s routed experts from the router front's logits and quantized input (``maps``: the layer's weight TMA maps; ``identity``: - 0, 1, ... int32).""" + 0, 1, ... int32; ``routing``: route_topk's ids and weights, for the 64-token engine).""" from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( FlashInferTrtllmDeferredFinalizeOutput, ) @@ -183,6 +183,7 @@ def routed(mlp, maps: tuple, identity: torch.Tensor, logits, x_fp4, x_sf): experts.g1_alphas, experts.g1_scale_c, experts.g2_alphas, + routing=routing, ) deferred = FlashInferTrtllmDeferredFinalizeOutput( gemm2_out=gemm2_out, diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index f7d169d6..bb430398 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -16,6 +16,23 @@ nvidia/qwen35-agentx-gb300), so the shared expert on the alternate stream starts at its activation, after that launch only; the gate/up sums in another order than the FP8 linear's GEMM. + +Slot ``route_topk`` serves the steps slot ``route`` does not (17 to 896 rows): behind +SGLang's own router GEMM, one launch (``sm103/moe/router.route_topk_quant``) writes each +token's top-10 experts and renormalized weights as the trtllm-gen routing computes them +and the NVFP4 input as the runner's quantizer does; the runner then serves the +materialized routing (FlashInfer's routed MoE), which skips its own top-k kernel. The +logits are SGLang's, so the step's output is the stock one's. + +Slot ``mkm_m32`` takes route_topk's place at 17 to 32 rows: slot mkm's 32-token engine runs +the top-10, FC1, SwiGLU, the NVFP4 requantization and FC2 in one launch on the router +GEMM's logits and the NVFP4 input (with route_gemm the router front's tiles give both; +otherwise route_topk's launch quantizes the input), bit for bit the runner's rows on +those logits, so the step's output is route_topk's. Slots ``mkm_m64`` / ``mkm_m128`` / +``mkm_m256`` / ``mkm_m512`` take the runner's place at 33 to 64 / 65 to 128 / 129 to 256 / +257 to 512 rows: after route_topk's launch (the routing and the NVFP4 input) the 64 / 128 / +256 / 512-token engine runs FC1, SwiGLU, the requantization and FC2 on that routing in one +launch, bit for bit the runner's rows. """ from __future__ import annotations @@ -47,7 +64,11 @@ "_silu_fp4_quant_fused", ), "srt/models/qwen3_5.py": ("layers", "start_layer", "end_layer", "mlp"), - "srt/layers/moe/topk.py": ("TopKOutputChecker", "format_is_bypassed"), + "srt/layers/moe/topk.py": ( + "TopKOutputChecker", + "format_is_bypassed", + "StandardTopKOutput", + ), "srt/layers/moe/fused_moe_triton/layer.py": ( "dispatcher", "run_moe_core", @@ -99,6 +120,17 @@ class Bound(msgspec.Struct, frozen=True): # [160] int32 0, 1, ...: slot mkm's rows are in expert order identity: object = None shared: bool = False # slot route_shared: the shared gate/up in the router launch + topk: bool = False # slot route_topk: the top-10 + NVFP4 input launch built + gemm: bool = ( + False # slot route_gemm: the router front's 16-row tiles stand in for the GEMM + ) + m32: dict = msgspec.field( + default_factory=dict + ) # slot mkm_m32: layer id -> weight TMA maps + m64: bool = False # slot mkm_m64: the 64-token engine built (on m32's maps) + m128: bool = False # slot mkm_m128: the 128-token engine built (on m32's maps) + m256: bool = False # slot mkm_m256: the 256-token engine built (on m32's maps) + m512: bool = False # slot mkm_m512: the 512-token engine built (on m32's maps) reason: str = "" # why a requested slot is not built @@ -161,6 +193,59 @@ def build( ) +def build_m32( + model, + bound: Bound, + m32: bool = True, + m64: bool = False, + m128: bool = False, + m256: bool = False, + m512: bool = False, +) -> Bound: + """Slots mkm_m32 / mkm_m64 / mkm_m128 / mkm_m256 / mkm_m512 on ``bound`` (route_topk + built): the 32 / 64 / 128 / 256 / 512-token engines and every MoE layer's weight maps + (slot mkm's when built), or why they cannot serve.""" + from sgl_nvidia_mega.qwen3_5.sm103.moe import mkm as mkm_kernel + + m = model.model + layers = [m.layers[i].mlp for i in range(m.start_layer, m.end_layer)] + device = torch.cuda.current_device() + intermediate = layers[0].experts.w2_weight.shape[-1] * 2 + reason = mkm_kernel.m32_unsupported(device, intermediate) + if not reason and (m256 or m512): + reason = mkm_kernel.wide_unsupported(device, intermediate) + if reason: + return msgspec.structs.replace(bound, reason=f"mkm_m32: {reason}") + # each requested engine (an engine between two requested ones is built and unused) + wanted = (m32, m64, m128, m256, m512) + bounds = mkm_kernel.ENGINE_ROWS[2:] + lo = bounds[wanted.index(True)] + hi = bounds[len(wanted) - 1 - wanted[::-1].index(True)] + mkm_kernel.build( + device, + intermediate, + hi, + min_rows=mkm_kernel.ENGINE_ROWS[bounds.index(lo) + 1] + 1, + ) + return msgspec.structs.replace( + bound, + m64=m64, + m128=m128, + m256=m256, + m512=m512, + m32={ + mlp.layer_id: bound.mkm.get(mlp.layer_id) + or mkm_kernel.weight_maps(mlp.experts.w13_weight, mlp.experts.w2_weight) + for mlp in layers + }, + identity=torch.arange( + max(hi, mkm_kernel.M128_MAX_TOKENS) * mkm_kernel.TOP_K, + dtype=torch.int32, + device="cuda", + ), + ) + + def unsupported(mlp) -> str: """Why the router front cannot stand in for ``mlp``'s gate and the runner's input quantizer ("" when it can).""" @@ -186,6 +271,21 @@ def unsupported(mlp) -> str: return "" +def topk_unsupported(model) -> str: + """Why slot route_topk cannot serve every MoE layer of ``model`` ("" when it can): + the router front's gate and input quantizer, and the trtllm-gen weights with a + softmax top-10 renormalized over 512 local experts (mkm.unsupported).""" + from sgl_nvidia_mega.qwen3_5.moe import mkm + + m = model.model + for i in range(m.start_layer, m.end_layer): + mlp = m.layers[i].mlp + reason = unsupported(mlp) or mkm.unsupported(mlp) + if reason: + return f"layer {i}: {reason}" + return "" + + def shared_unsupported(mlp) -> str: """Why the router front cannot compute ``mlp``'s shared gate/up ("" when it can): a static per-tensor FP8 projection of at most 512 rows, then SiLU-and-mul.""" @@ -289,3 +389,90 @@ def _routed_experts(mlp, x, fused, logits, x_fp4, x_sf): with flashinfer_trtllm_deferred_finalize_context(): combine_input = experts.run_moe_core(dispatch_output=dispatch_output) return experts.dispatcher.combine(combine_input=combine_input) + + +def routed_topk(mlp, x: torch.Tensor, gemm: bool = False): + """``FusedMoE.forward_deferred_finalize`` of ``mlp``'s routed experts on ``x`` + [rows, hidden] with SGLang's router GEMM, then slot route_topk's launch (the top-10 + and the NVFP4 input), and the runner on the materialized routing. With ``gemm`` + (slot route_gemm, numeric) the router front's GEMM and quantizer in 16-row tiles + stand in for SGLang's router GEMM, and route_topk's launch only routes.""" + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + flashinfer_trtllm_deferred_finalize_context, + ) + from sglang.srt.layers.moe.topk import StandardTopKOutput + + from sgl_nvidia_mega.qwen3_5.sm103.moe import router + + if ("route_topk", x.shape[0]) not in _ENGAGED: + _ENGAGED.add(("route_topk", x.shape[0])) + logger.info( + "NVIDIA megakernel: Qwen3.5 slot route_topk runs at %d rows", x.shape[0] + ) + experts = mlp.experts + scale = experts.w13_input_scale_quant.reshape(1) + if gemm: + logits, x_fp4, x_sf = router.router_gemm_quant_tiled(x, mlp.gate.weight, scale) + ids, weights, _, _ = router.route_topk_quant(logits, x, scale, quantize=False) + else: + logits, _ = mlp.gate(x) + ids, weights, x_fp4, x_sf = router.route_topk_quant(logits, x, scale) + topk_output = StandardTopKOutput( + topk_weights=weights, topk_ids=ids, router_logits=logits + ) + dispatch_output = experts.dispatcher.dispatch( + hidden_states=x, topk_output=topk_output + )._replace(hidden_states=x_fp4, hidden_states_scale=x_sf.view(torch.float8_e4m3fn)) + with flashinfer_trtllm_deferred_finalize_context(): + combine_input = experts.run_moe_core(dispatch_output=dispatch_output) + return experts.dispatcher.combine(combine_input=combine_input) + + +def routed_m32(mlp, x: torch.Tensor, bound: Bound, gemm: bool = False): + """``FusedMoE.forward_deferred_finalize`` of ``mlp``'s routed experts on ``x`` + [17-32 rows, hidden] in one launch of the 32-token engine (slot mkm_m32), on SGLang's + router GEMM and route_topk's quantized input, or with ``gemm`` (slot route_gemm, + numeric) on the router front's tiles (logits and quantized input in one launch).""" + from sgl_nvidia_mega.qwen3_5.moe import mkm + from sgl_nvidia_mega.qwen3_5.sm103.moe import router + + if ("mkm_m32", x.shape[0]) not in _ENGAGED: + _ENGAGED.add(("mkm_m32", x.shape[0])) + logger.info( + "NVIDIA megakernel: Qwen3.5 slot mkm_m32 runs at %d rows", x.shape[0] + ) + scale = mlp.experts.w13_input_scale_quant.reshape(1) + if gemm: + logits, x_fp4, x_sf = router.router_gemm_quant_tiled(x, mlp.gate.weight, scale) + else: + logits, _ = mlp.gate(x) + _, _, x_fp4, x_sf = router.route_topk_quant(logits, x, scale) + return mkm.routed(mlp, bound.m32[mlp.layer_id], bound.identity, logits, x_fp4, x_sf) + + +def routed_given(mlp, x: torch.Tensor, bound: Bound, slot: str): + """``FusedMoE.forward_deferred_finalize`` of ``mlp``'s routed experts on ``x`` + [33-512 rows, hidden]: SGLang's router GEMM, route_topk's launch (the routing and the + NVFP4 input, as slot route_topk runs them), then the 64, 128, 256 or 512-token engine on + that routing in place of the runner (slots mkm_m64, mkm_m128, mkm_m256, mkm_m512).""" + from sgl_nvidia_mega.qwen3_5.moe import mkm + from sgl_nvidia_mega.qwen3_5.sm103.moe import router + + if (slot, x.shape[0]) not in _ENGAGED: + _ENGAGED.add((slot, x.shape[0])) + logger.info( + "NVIDIA megakernel: Qwen3.5 slot %s runs at %d rows", slot, x.shape[0] + ) + logits, _ = mlp.gate(x) + ids, weights, x_fp4, x_sf = router.route_topk_quant( + logits, x, mlp.experts.w13_input_scale_quant.reshape(1) + ) + return mkm.routed( + mlp, + bound.m32[mlp.layer_id], + bound.identity, + logits, + x_fp4, + x_sf, + routing=(ids, weights), + ) diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index 67a5d5ff..cea84f8a 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -7,8 +7,11 @@ target verify of 1 to 128 requests (``REQUESTS``: every decode graph batch size SGLang captures up to 128); a width of more than one request is covered only where a slot runs. ``route`` (the router GEMM and the routed experts' NVFP4 input in one launch) runs up to -16 rows, ``mkm`` (on its outputs, the top-10 and both expert GEMMs in one launch) up to -7, and with ``mkm_wide`` (its 16-token engine) up to 16. ``gdn_verify`` (the +16 rows; ``route_topk`` materializes the top-10 and NVFP4 input behind SGLang's router +GEMM at wider rows. ``mkm`` (on route's outputs, the top-10 and both expert GEMMs in one +launch) runs up to 7 rows, and with ``mkm_wide`` (its 16-token engine) up to 16. +``mkm_m32`` runs the same routed-expert arithmetic on route_topk's 17–32 rows. +``gdn_verify`` (the GatedDeltaNet verify chain with its launch choices changed) covers every width, with ``gdn_core`` (numeric: its ReplaySSM verify, gated norm and FP8 output in one launch) up to 224 rows, and ``hop`` (the boundary's fused all-reduce points on the decode hop) the @@ -47,9 +50,18 @@ # hop the rows up to which FlashInfer's MNNVL CuTe DSL fusion runs the one-shot # all-reduce the hop equals bit for bit (its finalize runs one-shot up to 23 rows) ROUTE_MAX_ROWS = 16 +ROUTE_TOPK_MAX_ROWS = MAX_ROWS +# The router front's 16-row tiles replace SGLang's router GEMM through two +# waves; wider route_topk rows keep the stock GEMM. +ROUTE_GEMM_MAX_ROWS = 32 GDN_CORE_MAX_ROWS = 224 # and gdn_core's launch MKM_MAX_ROWS = 7 MKM_WIDE_MAX_ROWS = 16 +MKM_M32_MAX_ROWS = 32 +MKM_M64_MAX_ROWS = 64 +MKM_M128_MAX_ROWS = 128 +MKM_M256_MAX_ROWS = 256 +MKM_M512_MAX_ROWS = 512 HOP_MAX_ROWS = 15 GDN_PUSH_MAX_ROWS = 8 # the rows slot route_shared's shared stage fits shared memory for @@ -71,9 +83,19 @@ class Slots(msgspec.Struct, frozen=True): """The kernels that fill one width's step.""" route: bool = False # router GEMM + NVFP4 input quantization (sm103/moe/router.py) + # behind SGLang's router GEMM: top-10 + NVFP4 input in one launch (moe/route.py) + route_topk: bool = False + route_gemm: bool = False # with the router front's tiles in place of that GEMM mkm: bool = ( False # top-10 + FC1 + SwiGLU + FC2 on route's outputs (sm103/moe/mkm.py) ) + # in route_topk's place: the 32-token engine on the router GEMM's logits + mkm_m32: bool = False + # in route_topk's runner: the 64-token engine on its materialized routing + mkm_m64: bool = False + mkm_m128: bool = False # likewise, its 128-token engine + mkm_m256: bool = False # likewise, its 256-token engine + mkm_m512: bool = False # likewise, its 512-token engine gdn_verify: bool = False # the GDN verify chain's launch choices (gdn/verify.py) gdn_pdl: bool = False # its ReplaySSM verify launched dependent on the conv # its ReplaySSM verify, gated norm and FP8 output in one launch (sm103/gdn/core.py) @@ -106,8 +128,17 @@ class Facts(msgspec.Struct, frozen=True): function runs its stock code.""" route: bool = False # the router front built for every MoE layer + route_topk: bool = ( + False # the top-10 + NVFP4 input launch built for every MoE layer + ) + route_gemm: bool = False # and the router front's tiles mkm: bool = False # the fused routed experts built for every MoE layer mkm_wide: bool = False # with their 16-token engine + mkm_m32: bool = False # the 32-token engine built for every MoE layer + mkm_m64: bool = False # the 64-token engine built for every MoE layer + mkm_m128: bool = False # the 128-token engine built for every MoE layer + mkm_m256: bool = False # the 256-token engine built for every MoE layer + mkm_m512: bool = False # the 512-token engine built for every MoE layer gdn_verify: bool = False # the GDN verify slot built for every GDN layer gdn_pdl: bool = False # with its dependent verify launch gdn_core: bool = False # with the one-launch verify core @@ -133,14 +164,32 @@ def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> if not 1 <= rows <= MAX_ROWS: raise ValueError(f"{rows} rows are outside the covered widths 1-{MAX_ROWS}") route = facts.route and rows <= ROUTE_MAX_ROWS + route_topk = facts.route_topk and not route and rows <= ROUTE_TOPK_MAX_ROWS hop = facts.hop and rows <= HOP_MAX_ROWS mka = facts.mka and rows <= (MKA_WIDE_MAX_ROWS if facts.mka_wide else MKA_MAX_ROWS) lowm_se = facts.lowm_se and rows <= LOWM_SE_MAX_ROWS return Slots( route=route, + route_topk=route_topk, + route_gemm=route_topk and facts.route_gemm and rows <= ROUTE_GEMM_MAX_ROWS, mkm=route and facts.mkm and rows <= (MKM_WIDE_MAX_ROWS if facts.mkm_wide else MKM_MAX_ROWS), + mkm_m32=route_topk + and facts.mkm_m32 + and MKM_WIDE_MAX_ROWS < rows <= MKM_M32_MAX_ROWS, + mkm_m64=route_topk + and facts.mkm_m64 + and MKM_M32_MAX_ROWS < rows <= MKM_M64_MAX_ROWS, + mkm_m128=route_topk + and facts.mkm_m128 + and MKM_M64_MAX_ROWS < rows <= MKM_M128_MAX_ROWS, + mkm_m256=route_topk + and facts.mkm_m256 + and MKM_M128_MAX_ROWS < rows <= MKM_M256_MAX_ROWS, + mkm_m512=route_topk + and facts.mkm_m512 + and MKM_M256_MAX_ROWS < rows <= MKM_M512_MAX_ROWS, gdn_verify=facts.gdn_verify, gdn_pdl=facts.gdn_verify and facts.gdn_pdl, gdn_core=facts.gdn_verify and facts.gdn_core and rows <= GDN_CORE_MAX_ROWS, diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index b328575a..fd17216c 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -19,8 +19,15 @@ SHIPPED_SLOTS: tuple[str, ...] = () OPT_IN_SLOTS: tuple[str, ...] = ( "route", + "route_topk", + "route_gemm", "mkm", "mkm_wide", + "mkm_m32", + "mkm_m64", + "mkm_m128", + "mkm_m256", + "mkm_m512", "gdn_verify", "gdn_pdl", "gdn_core", @@ -67,7 +74,18 @@ SLOTS_ENV = "SGLANG_NVIDIA_MEGAKERNEL_SLOTS" # These slots run all 512 experts themselves. Under expert parallelism, keep the # route-only and boundary slots but let SGLang's EP path run the routed experts. -EP_STOCK_SLOTS = frozenset({"mkm", "mkm_wide", "draft_moe"}) +EP_STOCK_SLOTS = frozenset( + { + "mkm", + "mkm_wide", + "mkm_m32", + "mkm_m64", + "mkm_m128", + "mkm_m256", + "mkm_m512", + "draft_moe", + } +) def _requested_slots() -> frozenset: @@ -123,10 +141,47 @@ def slot_reasons() -> list[str]: f"{SLOTS_ENV}: slot 'mkm_wide' serves the steps of more than one request " "('multi_request')" ) + if "mkm_m32" in SLOTS and "route_topk" not in SLOTS: + reasons.append(f"{SLOTS_ENV}: slot 'mkm_m32' takes slot 'route_topk''s rows") + if "mkm_m32" in SLOTS and "multi_request" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'mkm_m32' serves the steps of more than two requests " + "('multi_request')" + ) + if "mkm_m64" in SLOTS and "route_topk" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'mkm_m64' runs on slot 'route_topk''s routing" + ) + if "mkm_m64" in SLOTS and "multi_request" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'mkm_m64' serves the steps of more than four requests " + "('multi_request')" + ) + if "mkm_m128" in SLOTS and "route_topk" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'mkm_m128' runs on slot 'route_topk''s routing" + ) + if "mkm_m128" in SLOTS and "multi_request" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'mkm_m128' serves the steps of more than nine requests " + "('multi_request')" + ) + for name, more in (("mkm_m256", "18"), ("mkm_m512", "36")): + if name in SLOTS and "route_topk" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot '{name}' runs on slot 'route_topk''s routing" + ) + if name in SLOTS and "multi_request" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot '{name}' serves the steps of more than {more} requests " + "('multi_request')" + ) if "route_shared" in SLOTS and "route" not in SLOTS: reasons.append( f"{SLOTS_ENV}: slot 'route_shared' runs in slot 'route''s launch" ) + if "route_gemm" in SLOTS and "route_topk" not in SLOTS: + reasons.append(f"{SLOTS_ENV}: slot 'route_gemm' runs inside slot 'route_topk'") if "gdn_push" in SLOTS and not {"gdn_verify", "hop"} <= SLOTS: reasons.append( f"{SLOTS_ENV}: slot 'gdn_push' pushes slot 'gdn_verify''s out_proj into " @@ -247,8 +302,15 @@ def build(model, ctx): } facts = slot_plan.Facts( route=bound["moe"].route, + route_topk=bound["moe"].topk, + route_gemm=bound["moe"].gemm, mkm=bool(bound["moe"].mkm), mkm_wide=bound["moe"].mkm_wide, + mkm_m32=bool(bound["moe"].m32) and "mkm_m32" in SLOTS, + mkm_m64=bool(bound["moe"].m64), + mkm_m128=bool(bound["moe"].m128), + mkm_m256=bool(bound["moe"].m256), + mkm_m512=bool(bound["moe"].m512), gdn_verify=bound["gdn"].gdn_verify, gdn_pdl=bound["gdn"].gdn_pdl, gdn_core=bound["gdn"].gdn_core, @@ -291,6 +353,34 @@ def build(model, ctx): <= slot_plan.MKM_MAX_ROWS else set() ) + if ( + slot_plan.widest(pctx.widths, slot_plan.MKM_M32_MAX_ROWS) + <= slot_plan.MKM_WIDE_MAX_ROWS + ): + vacuous.add("mkm_m32") + if ( + slot_plan.widest(pctx.widths, slot_plan.MKM_M64_MAX_ROWS) + <= slot_plan.MKM_M32_MAX_ROWS + ): + vacuous.add("mkm_m64") + if ( + slot_plan.widest(pctx.widths, slot_plan.MKM_M128_MAX_ROWS) + <= slot_plan.MKM_M64_MAX_ROWS + ): + vacuous.add("mkm_m128") + if ( + slot_plan.widest(pctx.widths, slot_plan.MKM_M256_MAX_ROWS) + <= slot_plan.MKM_M128_MAX_ROWS + ): + vacuous.add("mkm_m256") + if ( + slot_plan.widest(pctx.widths, slot_plan.MKM_M512_MAX_ROWS) + <= slot_plan.MKM_M256_MAX_ROWS + ): + vacuous.add("mkm_m512") + # route_topk has nothing to serve where slot route covers every width + if "route" in SLOTS and max(pctx.widths) <= slot_plan.ROUTE_MAX_ROWS: + vacuous.update({"route_topk", "route_gemm"}) missing = sorted(name for name in target_slots - vacuous if not built[name]) if missing: raise RuntimeError( diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py b/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py index 576a4e69..5aa31d1e 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py @@ -7,7 +7,10 @@ runner's deferred finalize takes them. One engine per intermediate size a rank holds: 128 at TP8 (``src/mkm_moe.cu``), 256 at TP4 (``src/mkm_moe_tp4.cu``), a 2-CTA cluster per selectable expert. At TP8, 8 to 16 rows run the 16-token engine (``src/mkm_moe_m16.cu``: -a CTA per (token, top-k) slot; a slot whose expert an earlier slot holds exits). +a CTA per (token, top-k) slot; a slot whose expert an earlier slot holds exits), and 17 to +32 rows the 32-token engine (``src/mkm_moe_m32.cu``: the 16-token engine's work per +distinct expert on at most one CTA per SM, each CTA running its share of the experts in +turn and gathering each one's token rows). """ from __future__ import annotations @@ -35,6 +38,7 @@ class Engine(NamedTuple): kernel: str smem_bytes: int # sizeof(Smem) + 1024, checked when the source compiles w2_swizzle: str # the W2 TMA map's swizzle: its box is a row of I / 2 bytes + defines: tuple = () # extra compile options # per intermediate size a rank holds (TP8, TP4) @@ -54,6 +58,64 @@ class Engine(NamedTuple): "mkm_moe_m16.cu", "q35_mkm_moe_m16", 221184, "CU_TENSOR_MAP_SWIZZLE_64B" ), } +# the 32-token engines, per intermediate size a rank holds (TP8) +M32_MAX_TOKENS = 32 +M32_ENGINES = { + 128: Engine( + "mkm_moe_m32.cu", "q35_mkm_moe_m32", 224256, "CU_TENSOR_MAP_SWIZZLE_64B" + ), +} +# SMs the 32-token engine leaves to the alternate stream's shared expert, whose +# GEMMs do not fit beside an engine CTA on an SM. The wider engines keep every SM. +# each CTA of the 32-token engine runs at most 3 distinct experts (src/mkm_moe_m32.cu +# kMaxItems): it needs 32 * 10 / 3 SMs; the 64-token engine's 8 items fit as well +M32_MIN_SMS = 107 +# the 64-token engines (the 32-token engine's work in (expert, row group) units) +M64_MAX_TOKENS = 64 +M64_ENGINES = { + 128: Engine( + "mkm_moe_m64.cu", "q35_mkm_moe_m64", 228352, "CU_TENSOR_MAP_SWIZZLE_64B" + ), +} +# the 128-token engines (the 64-token engine's units, on a given routing only) +M128_MAX_TOKENS = 128 +M128_ENGINES = { + 128: Engine( + "mkm_moe_m128.cu", "q35_mkm_moe_m128", 231424, "CU_TENSOR_MAP_SWIZZLE_64B" + ), +} +# the wide engines (src/mkm_moe_wide.cu, the 128-token engine's units on a static work split: CTA c runs experts +# c, c + G, ...), on a given routing only, built per token bound +M256_MAX_TOKENS = 256 +M512_MAX_TOKENS = 512 +M256_ENGINES = { + 128: Engine( + "mkm_moe_wide.cu", + "q35_mkm_moe_wide", + 226304, + "CU_TENSOR_MAP_SWIZZLE_64B", + ("-DQ35_MKM_MAX_TOK=256",), + ), +} +M512_ENGINES = { + 128: Engine( + "mkm_moe_wide.cu", + "q35_mkm_moe_wide", + 231424, + "CU_TENSOR_MAP_SWIZZLE_64B", + ("-DQ35_MKM_MAX_TOK=512",), + ), +} +# a wide engine's CTA runs at most 4 own experts (512 / 128) and 3 further row groups (its kMaxItems) +WIDE_MIN_SMS = 128 +WIDE_ROWS = (M256_MAX_TOKENS, M512_MAX_TOKENS) +RESERVED_SMS = { + M32_MAX_TOKENS: 32, + M64_MAX_TOKENS: 0, + M128_MAX_TOKENS: 0, + M256_MAX_TOKENS: 0, + M512_MAX_TOKENS: 0, +} class Params(ctypes.Structure): @@ -77,11 +139,45 @@ class Params(ctypes.Structure): ] + [("num_tokens", ctypes.c_int)] +class ParamsM32(Params): + """qwen_mkm_moe_m32::Params / qwen_mkm_moe_m64::Params: each item gathers its token rows + from x_fp4; the 64-token engine also takes a given routing (ids_in, null: its own).""" + + _fields_ = [("x_fp4", ctypes.c_void_p), ("ids_in", ctypes.c_void_p)] + + +# the rows each engine serves: (7-row, 16-token, 32-token, 64-token, 128-token, 256-token, 512-token) engines up to +# these +ENGINE_ROWS = ( + MAX_TOKENS, + M16_MAX_TOKENS, + M32_MAX_TOKENS, + M64_MAX_TOKENS, + M128_MAX_TOKENS, + M256_MAX_TOKENS, + M512_MAX_TOKENS, +) +# the engines that take a given routing (route_topk's ids); those past 64 tokens only that way +ROUTED_ENGINES = (M64_MAX_TOKENS, M128_MAX_TOKENS, M256_MAX_TOKENS, M512_MAX_TOKENS) + + +def engine_rows(m: int) -> int: + """The bound of the engine that serves ``m`` rows.""" + return next(r for r in ENGINE_ROWS if m <= r) + + @cache -def _kernel( - device_index: int, intermediate: int = 128, wide: bool = False -) -> jit.Kernel: - engine = (WIDE_ENGINES if wide else ENGINES)[intermediate] +def _kernel(device_index: int, intermediate: int, rows: int) -> jit.Kernel: + """The engine of ``engine_rows`` bound ``rows`` at ``intermediate`` per rank.""" + engine = { + MAX_TOKENS: ENGINES, + M16_MAX_TOKENS: WIDE_ENGINES, + M32_MAX_TOKENS: M32_ENGINES, + M64_MAX_TOKENS: M64_ENGINES, + M128_MAX_TOKENS: M128_ENGINES, + M256_MAX_TOKENS: M256_ENGINES, + M512_MAX_TOKENS: M512_ENGINES, + }[rows][intermediate] major, minor = torch.cuda.get_device_capability(device_index) cubin = jit.compile_source( (Path(__file__).parent / "src" / engine.source).read_text(), @@ -91,24 +187,74 @@ def _kernel( "-std=c++17", "-default-device", f"-DQ35_MKM_SMEM_BYTES={engine.smem_bytes}", - ), + ) + + engine.defines, ) return jit.load(cubin, engine.kernel, (THREADS, 1, 1), engine.smem_bytes) -def max_tokens(intermediate: int) -> int: - """The widest rows an engine serves for ``intermediate`` per rank.""" +def max_tokens( + intermediate: int, + m32: bool = False, + m64: bool = False, + m128: bool = False, + m256: bool = False, + m512: bool = False, +) -> int: + """The widest rows an engine serves for ``intermediate`` per rank (``m32`` / ``m64`` / + ``m128`` / ``m256`` / ``m512``: with the 32 / 64 / 128 / 256 / 512-token engine).""" + if m512 and intermediate in M512_ENGINES: + return M512_MAX_TOKENS + if m256 and intermediate in M256_ENGINES: + return M256_MAX_TOKENS + if m128 and intermediate in M128_ENGINES: + return M128_MAX_TOKENS + if m64 and intermediate in M64_ENGINES: + return M64_MAX_TOKENS + if m32 and intermediate in M32_ENGINES: + return M32_MAX_TOKENS return M16_MAX_TOKENS if intermediate in WIDE_ENGINES else MAX_TOKENS +def m32_unsupported(device_index: int, intermediate: int) -> str: + """Why the 32-token engine cannot serve ``intermediate`` per rank on this GPU ("" when + it can).""" + if intermediate not in M32_ENGINES: + return f"no 32-token engine at intermediate {intermediate} per rank" + sms = torch.cuda.get_device_properties(device_index).multi_processor_count - max( + RESERVED_SMS.values() + ) + if sms < M32_MIN_SMS: + return f"{sms} SMs (the 32-token engine needs {M32_MIN_SMS})" + return "" + + +def wide_unsupported(device_index: int, intermediate: int) -> str: + """Why the 256 / 512-token engines cannot serve ``intermediate`` per rank on this GPU + ("" when they can).""" + if intermediate not in M256_ENGINES: + return f"no 256 / 512-token engine at intermediate {intermediate} per rank" + sms = torch.cuda.get_device_properties(device_index).multi_processor_count - max( + RESERVED_SMS[r] for r in WIDE_ROWS + ) + if sms < WIDE_MIN_SMS: + return f"{sms} SMs (the 256 / 512-token engines need {WIDE_MIN_SMS})" + return "" + + def build( - device_index: int, intermediate: int = 128, max_rows: int = MAX_TOKENS + device_index: int, + intermediate: int = 128, + max_rows: int = MAX_TOKENS, + min_rows: int = 1, ) -> None: - """Compile the engines that serve ``intermediate`` per rank up to ``max_rows`` rows - (no compile inside a step).""" - _kernel(device_index, intermediate) - if max_rows > MAX_TOKENS: - _kernel(device_index, intermediate, wide=True) + """Compile the engines that serve ``intermediate`` per rank on ``min_rows`` to + ``max_rows`` rows (no compile inside a step).""" + lo = 0 + for hi in ENGINE_ROWS: + if min_rows <= hi and max_rows > lo: + _kernel(device_index, intermediate, hi) + lo = hi def _map( @@ -177,24 +323,34 @@ def mkm_moe( g1_alpha: torch.Tensor, g1_scale_c: torch.Tensor, g2_alpha: torch.Tensor, + routing: tuple | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """(gemm2_out bf16 [M * 10, 4096], expert_weights bf16 [M, 10]) for logits [M, 512] - bf16 (M <= 7, at TP8 M <= 16) and the router front's x_fp4 [M, 2048] / x_sf [M, 256] + bf16 (M <= 7, at TP8 M <= 128) and the router front's x_fp4 [M, 2048] / x_sf [M, 256] (uint8, linear e4m3); w13_sf [512, 2 I, 256] / w2_sf [512, 4096, I / 16] e4m3 viewed - as uint8 (the engine follows I); g1_alpha, g1_scale_c, g2_alpha [512] fp32.""" + as uint8 (the engine follows I); g1_alpha, g1_scale_c, g2_alpha [512] fp32. With + ``routing`` (route_topk's ids int32 / weights bf16 [M, 10]; 33 to 64 rows, and required + at 65 to 512) the engine takes that routing instead of ranking the logits, and returns + those weights.""" m = logits.shape[0] intermediate = w2_sf.shape[2] * 16 - if not 1 <= m <= max_tokens(intermediate): - raise ValueError(f"mkm: {m} rows (1-{max_tokens(intermediate)})") - wide = m > MAX_TOKENS + widest = max_tokens( + intermediate, m32=True, m64=True, m128=True, m256=True, m512=True + ) + if not 1 <= m <= widest: + raise ValueError(f"mkm: {m} rows (1-{widest})") + rows = engine_rows(m) + if routing is not None and rows not in ROUTED_ENGINES: + raise ValueError( + f"mkm: a given routing needs a 64-512-token engine, not {m} rows" + ) + if routing is None and rows > M64_MAX_TOKENS: + raise ValueError(f"mkm: {m} rows need a given routing (route_topk's)") device = logits.device gemm2_out = torch.empty((m * TOP_K, HIDDEN), dtype=torch.bfloat16, device=device) expert_weights = torch.empty((m, TOP_K), dtype=torch.bfloat16, device=device) topk_ids = torch.empty((m, TOP_K), dtype=torch.int32, device=device) - swizzle = jit.driver().CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B - # the x rows past m read as zeros: 8 rows a box, 16 for the 16-token engine - x_map = _map(x_fp4, m, HIDDEN // 2, 128, 16 if wide else 8, swizzle) - p = Params( + fields = dict( logits=logits.data_ptr(), expert_weights=expert_weights.data_ptr(), topk_ids=topk_ids.data_ptr(), @@ -207,10 +363,34 @@ def mkm_moe( out=gemm2_out.data_ptr(), num_tokens=m, ) + kernel = _kernel(device.index, intermediate, rows) + stream = torch.cuda.current_stream().cuda_stream + if rows in (M32_MAX_TOKENS, M64_MAX_TOKENS, M128_MAX_TOKENS) + WIDE_ROWS: + # min(M * 10, SMs) CTAs, one per SM, each running its share of the distinct experts + # (the wide engines: every SM, CTA c running experts c, c + G, ...) + sms = ( + torch.cuda.get_device_properties(device).multi_processor_count + - RESERVED_SMS[rows] + ) + ids_in = 0 if routing is None else routing[0].data_ptr() + if rows in WIDE_ROWS and not ( + routing[0].is_contiguous() and ids_in % 16 == 0 and sms >= WIDE_MIN_SMS + ): + raise ValueError( + "mkm: the wide engines read contiguous 16-B aligned ids on >= 128 SMs" + ) + p = ParamsM32(x_fp4=x_fp4.data_ptr(), ids_in=ids_in, **fields) + grid = sms if rows in WIDE_ROWS else min(m * TOP_K, sms) + jit.launch(kernel, (grid, 1, 1), [*maps, p], stream) + return gemm2_out, (expert_weights if routing is None else routing[1]) + swizzle = jit.driver().CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B + wide = rows == M16_MAX_TOKENS + # the x rows past m read as zeros: 8 rows a box, 16 for the 16-token engine + x_map = _map(x_fp4, m, HIDDEN // 2, 128, 16 if wide else 8, swizzle) jit.launch( - _kernel(device.index, intermediate, wide), + kernel, (M16_CTAS if wide else CTAS, 1, 1), - [*maps, x_map, p], - torch.cuda.current_stream().cuda_stream, + [*maps, x_map, Params(**fields)], + stream, ) return gemm2_out, expert_weights diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py b/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py index c9f035a8..d0eb1477 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py @@ -204,3 +204,127 @@ def router_gemm_quant_shared( kernel = _shared_kernel(x.device.index, max_m_bucket(m)) jit.launch(kernel, (CTAS, 1, 1), [p], torch.cuda.current_stream().cuda_stream) return logits, x_fp4, x_sf, gate_up + + +# The top-10 and NVFP4 input behind SGLang's router GEMM (slot route_topk, src/route_topk.cu) +TOPK_SOURCE = Path(__file__).parent / "src" / "route_topk.cu" +TOP_K = 10 +TOPK_MAX_ROWS = 896 # the widest verify step it is built for (128 requests of 7 rows) + + +class TopkParams(ctypes.Structure): + """q35_route_topk::Params.""" + + _fields_ = [ + (name, ctypes.c_void_p) + for name in ( + "logits", + "x", + "global_scale", + "topk_ids", + "topk_weights", + "x_fp4", + "x_sf", + ) + ] + + +@cache +def _topk_kernel(device_index: int) -> jit.Kernel: + major, minor = torch.cuda.get_device_capability(device_index) + cubin = jit.compile_source( + TOPK_SOURCE.read_text(), + "q35_route_topk", + f"sm_{major}{minor}a", + ("-std=c++17", "-default-device"), + ) + return jit.load(cubin, "q35_route_topk_quant", (THREADS, 1, 1), 0) + + +def build_topk(device_index: int) -> None: + """Compile slot route_topk's launch (no compile inside a step).""" + _topk_kernel(device_index) + + +def route_topk_quant( + logits: torch.Tensor, + x: torch.Tensor, + global_scale: torch.Tensor, + quantize: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """logits [M, 512] bf16 (the router GEMM's), x [M, 4096] bf16, global_scale [1] + fp32: (topk_ids int32 [M, 10], topk_weights bf16 [M, 10], x_fp4 uint8 [M, 2048], + x_sf uint8 [M, 256]); the routing equals the trtllm-gen RenormalizeNaive routing's + and the input the runner's NVFP4 quantizer's, bit for bit. Without ``quantize`` only + the routing (x_fp4 / x_sf None).""" + m = x.shape[0] + for t, shape in ((logits, (m, EXPERTS)), (x, (m, HIDDEN))): + if ( + t.dtype != torch.bfloat16 + or tuple(t.shape) != shape + or not t.is_contiguous() + ): + raise ValueError(f"route_topk: {t.dtype} {tuple(t.shape)}, want {shape}") + if not 1 <= m <= TOPK_MAX_ROWS: + raise ValueError(f"route_topk: {m} rows (1-{TOPK_MAX_ROWS})") + ids = torch.empty((m, TOP_K), dtype=torch.int32, device=x.device) + weights = torch.empty((m, TOP_K), dtype=torch.bfloat16, device=x.device) + x_fp4 = x_sf = None + if quantize: + x_fp4 = torch.empty((m, HIDDEN // 2), dtype=torch.uint8, device=x.device) + x_sf = torch.empty((m, HIDDEN // 16), dtype=torch.uint8, device=x.device) + p = TopkParams( + logits=logits.data_ptr(), + x=x.data_ptr(), + global_scale=global_scale.data_ptr(), + topk_ids=ids.data_ptr(), + topk_weights=weights.data_ptr(), + x_fp4=jit.null_or_pointer(x_fp4).value, + x_sf=jit.null_or_pointer(x_sf).value, + ) + jit.launch( + _topk_kernel(x.device.index), + (m, 1, 1), + [p], + torch.cuda.current_stream().cuda_stream, + ) + return ids, weights, x_fp4, x_sf + + +GEMM_MAX_ROWS = 56 # slot route_gemm: the router front in tiles of 16 rows up to here + + +def router_gemm_quant_tiled( + x: torch.Tensor, weight: torch.Tensor, global_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``router_gemm_quant`` past 16 rows (up to GEMM_MAX_ROWS): one launch of the + 16-row front over token tiles (grid y). Each tile's logits sum as the 16-row front's, + in another order than SGLang's cuBLAS GEMM at these rows (numeric).""" + m = x.shape[0] + if x.dtype != torch.bfloat16 or x.shape[1:] != (HIDDEN,) or not x.is_contiguous(): + raise ValueError( + f"router front: x {x.dtype} {tuple(x.stride())} {tuple(x.shape)}" + ) + if not 1 <= m <= GEMM_MAX_ROWS: + raise ValueError(f"tiled router front: {m} rows (1-{GEMM_MAX_ROWS})") + logits = torch.empty((m, EXPERTS), dtype=torch.bfloat16, device=x.device) + x_fp4 = torch.empty((m, HIDDEN // 2), dtype=torch.uint8, device=x.device) + x_sf = torch.empty((m, HIDDEN // 16), dtype=torch.uint8, device=x.device) + p = Params( + x=x.data_ptr(), + weight=weight.data_ptr(), + global_scale=global_scale.data_ptr(), + logits=logits.data_ptr(), + x_fp4=x_fp4.data_ptr(), + x_sf=x_sf.data_ptr(), + num_tokens=m, + mode=MODE_QUANT, + ) + tiles = -(-m // MAX_TOKENS) + jit.launch( + _kernel(x.device.index, MAX_TOKENS), + (CTAS, tiles, 1), + [p], + torch.cuda.current_stream().cuda_stream, + ) + return logits, x_fp4, x_sf diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m128.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m128.cu new file mode 100644 index 00000000..6845a58c --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m128.cu @@ -0,0 +1,1412 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Fused small-M NVFP4 MoE expert phase for Qwen3.5 TP8 decode verify (M <= 128 tokens, top-10, 512 experts, +// hidden 4096, intermediate 128 per rank): the top-10 of the router logits in the trtllm-gen routing's order and +// weights, FC1, SwiGLU, the NVFP4 requantization and FC2 in one launch over the trtllm-gen weight layout. Writes the +// deferred-finalize rows (row t * 10 + k, identity permutation) and the expert weights. +// +// The 64-token engine (mkm_moe_m64.cu) for up to 128 tokens, on a given routing only (route_topk's ids): each item +// carries its rows as a token list. As there, the work is (expert, row group of at most 16 of its +// tokens) units, so an expert that more than 16 tokens chose runs as several units on different CTAs. CTA c speculates +// on slot c's expert and runs its first unit when slot c is that expert's first; every other unit (experts first seen +// past the grid, further row groups) goes to the CTAs whose slot repeats an earlier expert, then round robin to +// CTAs 0, 1, ... Per unit the arithmetic and summation order are the 16-token engine's: rows bit for bit +// trtllm-gen's. +// +// Built with NVRTC: the kernel has C linkage (q35_mkm_moe_m128), and the launcher's dynamic shared memory size +// (Q35_MKM_SMEM_BYTES) is checked against the layout at compile time. +#include +#include +#include + +// NVRTC has no , or +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef int int32_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef unsigned long long uintptr_t; +typedef struct __align__(64) { + uint64_t opaque[16]; +} CUtensorMap; +#ifndef INFINITY +#define INFINITY __int_as_float(0x7f800000) +#endif +// Minimal SM100 (tcgen05 / TMA / mbarrier) PTX helpers for the small-M NVFP4 MoE engine. + +namespace qwen35 { + +__device__ __forceinline__ uint32_t smem_u32(const void* p) { + return static_cast(__cvta_generic_to_shared(p)); +} + +// ---------------------------------------------------------------- mbarrier +__device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(smem_u32(bar)), "r"(count) : "memory"); +} +__device__ __forceinline__ void fence_mbar_init() { + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} +__device__ __forceinline__ void mbar_expect_tx(uint64_t* bar, uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smem_u32(bar)), "r"(bytes) : "memory"); +} +__device__ __forceinline__ void mbar_arrive(uint64_t* bar) { + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(smem_u32(bar)) : "memory"); +} +__device__ __forceinline__ bool mbar_try_wait(uint64_t* bar, uint32_t phase) { + uint32_t ok; + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ok) + : "r"(smem_u32(bar)), "r"(phase) + : "memory"); + return ok != 0; +} +__device__ __forceinline__ void mbar_wait(uint64_t* bar, uint32_t phase) { + while (!mbar_try_wait(bar, phase)) { + } +} + +// ---------------------------------------------------------------- TMA / bulk copies +__device__ __forceinline__ void tma_load_2d(void* dst, const void* tmap, uint64_t* bar, int c0, int c1) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes [%0], [%1, {%3, %4}], [%2];" :: + "r"(smem_u32(dst)), + "l"(tmap), "r"(smem_u32(bar)), "r"(c0), "r"(c1) + : "memory"); +} +__device__ __forceinline__ void bulk_load(void* dst, const void* src, uint32_t bytes, uint64_t* bar) { + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)) + : "memory"); +} +__device__ __forceinline__ void prefetch_tmap(const void* tmap) { + asm volatile("prefetch.tensormap [%0];" ::"l"(tmap) : "memory"); +} +// generic-proxy smem writes -> visible to the async proxy (UMMA / UTCCP operand reads) +__device__ __forceinline__ void fence_proxy_async_smem() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + +// ---------------------------------------------------------------- tcgen05 +__device__ __forceinline__ void tmem_alloc(uint32_t* dst_smem, uint32_t ncols) { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smem_u32(dst_smem)), + "r"(ncols) + : "memory"); +} +__device__ __forceinline__ void tmem_relinquish() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;" ::: "memory"); +} +__device__ __forceinline__ void tmem_dealloc(uint32_t taddr, uint32_t ncols) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(ncols) : "memory"); +} +__device__ __forceinline__ void tc_fence_before() { + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); +} +__device__ __forceinline__ void tc_fence_after() { + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); +} +// arrive on `bar` when all prior tcgen05 async ops issued by this thread complete +__device__ __forceinline__ void tc_commit(uint64_t* bar) { + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" ::"r"(smem_u32(bar)) + : "memory"); +} +// smem -> TMEM copy of one 512-byte scale-factor chunk (32 rows x 16 B), broadcast to the 4 lane quarters +__device__ __forceinline__ void utccp_32x128b_warpx4(uint32_t taddr, uint64_t sdesc) { + asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" ::"r"(taddr), "l"(sdesc) : "memory"); +} +// D(tmem) (+)= A(smem) * B(smem)^T with NVFP4 block-16 scales from TMEM +__device__ __forceinline__ void umma_nvf4(uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, + uint32_t tmem_sfa, uint32_t tmem_sfb, uint32_t accumulate) { + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %4, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%5], [%6], p;\n}\n" ::"r"( + tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(accumulate), "r"(tmem_sfa), "r"(tmem_sfb) + : "memory"); +} +// each warp reads its 32-lane quarter: 32 lanes x 16 consecutive 32-bit columns +__device__ __forceinline__ void tmem_ld_32x32b_x16(uint32_t taddr, uint32_t (&r)[16]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15}, [%16];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(taddr) + : "memory"); +} +// L2 prefetch of a contiguous global range (bytes multiple of 16, 16-byte aligned) +__device__ __forceinline__ void bulk_prefetch_l2(const void* gptr, uint32_t bytes) { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(gptr), "r"(bytes) : "memory"); +} +__device__ __forceinline__ void tmem_ld_wait() { asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); } +// each warp writes its 32-lane quarter: 32 lanes x 16 consecutive 32-bit columns +__device__ __forceinline__ void tmem_st_32x32b_x16(uint32_t taddr, const uint32_t (&r)[16]) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16};" ::"r"( + taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_wait() { asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); } + +// ---------------------------------------------------------------- descriptors +// SM100 shared-memory matrix descriptor (version 1). layout: 0 none, 2 SW128, 4 SW64, 6 SW32. +__device__ __forceinline__ uint64_t make_sdesc(uint32_t saddr, uint32_t lbo, uint32_t sbo, uint32_t layout) { + uint64_t d = 0; + d |= (uint64_t)((saddr >> 4) & 0x3FFF); + d |= (uint64_t)((lbo >> 4) & 0x3FFF) << 16; + d |= (uint64_t)((sbo >> 4) & 0x3FFF) << 32; + d |= (uint64_t)1 << 46; // version + d |= (uint64_t)(layout & 0x7) << 61; + return d; +} +// Block-scaled instruction descriptor: E2M1 x E2M1, UE4M3 scales, both K-major, F32 accumulate. +__host__ __device__ constexpr uint32_t make_idesc_nvf4(uint32_t M, uint32_t N) { + return (1u << 7) // a_format E2M1 + | (1u << 10) // b_format E2M1 + | ((N >> 3) << 17) // n_dim + | (0u << 23) // scale format UE4M3 + | ((M >> 4) << 24); // m_dim +} + +} // namespace qwen35 + +namespace sglang { + +namespace qwen_mkm_moe_m128 { + +using namespace qwen35; + + +constexpr int kH = 4096, kE = 512, kTopK = 10, kMaxTok = 128; +constexpr int kGroups = kMaxTok / 32; // 32-token groups: lane l holds tokens l + 32 h +constexpr int kRows = 16; // A rows (tokens) per item +// per CTA: at most 1 + ceil(units / CTAs) items; units <= 560 at 128 tokens (each expert's ceil(tokens / 16)) over +// >= 107 CTAs (the launcher's bound) +constexpr int kMaxItems = 8; +constexpr int kI = 128; // intermediates per expert per rank (TP8) +constexpr int kN1 = 2 * kI; // W13 rows per expert (gate + up) +constexpr int kStageBytes = 36864; // 32 KB weights + 4 KB scale factors +constexpr int kStages = 4; +constexpr int kW13BoxRows = 128; // rows of one W13 TMA box; the host tensor map must use the same +constexpr int kFc1Stages = 16; // the expert's 256 W13 rows, 4 k-blocks each (full K = 4096) +constexpr int kFc2Stages = 8; // two tiles of 256 W2 rows x 64 B each +constexpr int kTotStages = kFc1Stages + kFc2Stages; +constexpr int kThreads = 256; +// TMEM columns. FC1: accumulator [0,256), x SFA ring [256,320) (16 slots x 4), SFB ring [384,512) (16 slots x 8). +// FC2: accumulator [0,256), h SFA [256,320) (16 atoms x 4), SFB ring [384,512) (16 slots x 8). +constexpr uint32_t kColSfaX = 256, kColSfa = 256, kColSfb = 384; +constexpr int kSfbSlots = 16, kSfaSlots = 16; +constexpr int kARow = 260; // FC1 accumulator row stride in floats: 16 B aligned, rows on distinct banks +#ifndef L2_EVICT_FIRST // stream expert weights with an L2 evict-first policy (they are read once per layer) +#define L2_EVICT_FIRST 1 +#endif +#ifndef TOPK_RANK +#define TOPK_RANK 0 +#endif +#ifndef SF_UTCCP_FC1 +#define SF_UTCCP_FC1 0 +#endif +#ifndef SPEC_STAGES // ring stages issued for the speculated expert before the routing verdict (rest: after it) +#define SPEC_STAGES 4 +#endif +#ifndef STAGE_STAMPS // DIAG: per-stage issue (slots 64 + st) and landing (16 + st, seen by warp 4) stamps +#define STAGE_STAMPS 0 +#endif +#ifndef DIAG_SPEC // DIAG: warp 1's speculative ranking, stamps 16 = logits loaded, 17 = keys, 18 = top-10 (stride 32) +#define DIAG_SPEC 0 +#endif +#ifndef DIAG_ROUTE // DIAG: stamps 13 = warp 0's routing rounds done, 15 = their barrier, 14 = item 0's rows gathered +#define DIAG_ROUTE 0 +#endif +constexpr int kTraceStride = STAGE_STAMPS ? 128 : (DIAG_SPEC || DIAG_ROUTE) ? 32 : 16; +#ifndef SMEM_SPACE // 1: shared-memory accesses through Smem compile to shared-space instructions +#define SMEM_SPACE 1 +#endif +#ifndef PF_SET // L2 prefetch of the stages that wait for a ring slot at: 1 = start, 4 = FC2, 8 = FC2 pass 1 +// (none: from 56 rows every CTA streams and HBM is the bound, so these prefetches only compete with the ring's loads) +#define PF_SET 0 +#endif +#ifndef PF_N1 // stages prefetched at the start trigger +#define PF_N1 4 +#endif +#ifndef PF_N2 // stages prefetched at the FC2 trigger +#define PF_N2 4 +#endif +#ifndef TAILV // FC2 drains: bit 0 = last pass on all 8 warps; bit 1 = TMEM loads double-buffered; bit 2 = coalesced stores +#define TAILV 3 +#endif +#define TAIL8 (TAILV & 1) +#define DRAIN_PIPE ((TAILV >> 1) & 1) +#define DRAIN_SMEM ((TAILV >> 2) & 1) +static_assert(DRAIN_SMEM == 0, "x's shared memory holds the next item's rows while the drains run"); +#ifndef EARLY_SF2 // FC2 stages whose scale factors are staged before the h epilogue (old order: 2) +#define EARLY_SF2 0 +#endif +#ifndef GATHER_BATCH // x row granules each gathering thread loads before storing them +#define GATHER_BATCH 8 +#endif +#ifndef GATHER_ASYNC // 1: the gather copies by cp.async (every granule in flight at once, no register staging) +#define GATHER_ASYNC 1 +#endif +#ifndef SPEC_PF // stages of the speculated expert past the first ring fill prefetched to L2 while the routing runs +#define SPEC_PF 4 +#endif +#ifndef FC2_PF // FC2 tiles (beyond the first kStages) prefetched to L2 ahead of the ring; 0 = off +#define FC2_PF 0 +#endif +#ifndef FC2_PF_LEAD // issue the FC2 prefetch when the producer reaches FC1 stage kFc1Stages - FC2_PF_LEAD +#define FC2_PF_LEAD 4 +#endif + +struct Params { + const uint16_t* logits; // [M, 512] bf16 router logits (router_gemm_quant output) + uint16_t* expert_weights; // [M, 10] bf16 out: DynBlock softmax-over-top-10 weights (the last CTA writes) + int32_t* topk_ids; // [M, 10] out (the last CTA writes) + const uint8_t* x_sf; // [M, 256] linear e4m3 + const uint8_t* w13_sf; // [E, 2 row blocks x 64 k blocks x 512 B] (128x4 interleaved) + const uint8_t* w2_sf; // [E, 32 row blocks x 2 k blocks x 512 B] + const float* g1_alpha; // [E] gate dequant + const float* g1_scale_c; // [E] up dequant * FC2 input quant scale + const float* g2_alpha; // [E] + __nv_bfloat16* out; // [M * 10, 4096] + long long* trace; // optional [grid, 16] %globaltimer stamps + int num_tokens; + const uint8_t* x_fp4; // [M, 2048] NVFP4 activations (each item gathers its rows) + // [M, 10] the routing's expert ids (route_topk's, bit for bit the trtllm-gen routing): when set, the engine takes them + // in place of ranking the logits (and writes neither expert_weights nor topk_ids, which route_topk wrote) + const int32_t* ids_in; +}; + +struct __align__(1024) Smem { + uint8_t ring[kStages][kStageBytes]; + alignas(1024) uint8_t xa[32 * 1024]; // x: per 128-B K chunk c, 2 KB = SW128 atoms of item rows 0..7 | 8..15 + alignas(1024) uint8_t ha[16 * 1024]; // h: 8 copies of (rows 0..7 | rows 8..15) x 128 B (K = 128: 64 B), SW128 + float acc1[16 * kARow]; // FC1 accumulator rows 0..15 (item rows) x 256 columns + uint8_t xsf[16][256]; // x scale bytes of the item's rows, full K + uint32_t hsf[16][2]; // h scale words: row r, k-block kb (4 e4m3 bytes) + uint64_t full[kStages], empty[kStages], sf_full[kStages]; + uint64_t x_full, fc1_done, acc_ready, h_ready, fc2_pass[2], fc2_free, acc_drained, route_done; + uint32_t tmem; + int nitems, spec_ok; + int item_e[kMaxItems]; + uint8_t item_rows[kMaxItems][kRows]; // each item's tokens, ascending + int item_n[kMaxItems]; + int nrows[2], rtok[2][kRows], rkslot[2][kRows]; // rows of the current / next item (by item parity) + int first[kE]; // first slot of each expert (atomicMin) + int cnt[kE]; // tokens of each expert (atomicAdd) + uint32_t fbits[kMaxTok]; // token t: its top-k slots that are their expert's first (bit k) + int ounits[kMaxTok]; // token t: the units its first-occurrence slots hand out (not a CTA's own) + uint8_t uslot[kMaxTok][kTopK]; // per slot: those units (0 unless a first occurrence) + int32_t ids[kMaxTok][kTopK]; +}; + +__device__ __forceinline__ uint64_t gtime() { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)::"memory"); + return t; +} +__device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} +__device__ __forceinline__ void cluster_arrive() { + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); +} +// no release: the phase only orders DSMEM traffic that mbarriers already completed, so pending global stores +// (stamps, routing outputs, output rows) need not drain before the arrive +__device__ __forceinline__ void cluster_arrive_relaxed() { + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} +__device__ __forceinline__ void cluster_wait() { + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); +} +__device__ __forceinline__ uint32_t mapa(uint32_t saddr, uint32_t rank) { + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(r) : "r"(saddr), "r"(rank)); + return r; +} +// smem (this CTA) -> smem (peer CTA in the cluster), completes tx bytes on the peer's mbarrier +__device__ __forceinline__ void bulk_s2s(uint32_t dst_cluster, uint32_t src, uint32_t bytes, uint32_t bar_cluster) { + asm volatile("cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + dst_cluster), + "r"(src), "r"(bytes), "r"(bar_cluster) + : "memory"); +} +// registers -> the peer CTA's shared memory; completes the bytes on the peer's mbarrier +__device__ __forceinline__ void st_async_v4(uint32_t dst, const uint4& v, uint32_t bar) { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"(dst), + "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "r"(bar) + : "memory"); +} +__device__ __forceinline__ void st_async_b32(uint32_t dst, uint32_t v, uint32_t bar) { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(dst), "r"(v), "r"(bar) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x4(uint32_t taddr, const uint4& v) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x4.b32 [%0], {%1, %2, %3, %4};" ::"r"(taddr), "r"(v.x), "r"(v.y), + "r"(v.z), "r"(v.w) + : "memory"); +} +// bar.sync counts arrivals per warp execution: reconverge first, so a warp whose lane 0 just ran a stamp or an +// mbarrier arrive cannot arrive twice and release the barrier early +__device__ __forceinline__ void named_bar(int id, int n) { + __syncwarp(); + asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(n) : "memory"); +} +__device__ __forceinline__ void tmem_ld_x32(uint32_t taddr, uint32_t (&r)[32]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,%19," + "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), "=r"(r[16]), + "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), "=r"(r[24]), + "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr) + : "memory"); +} +__device__ __forceinline__ void tmem_ld_x64(uint32_t taddr, uint32_t (&r)[64]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x64.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,%19," + "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31,%32,%33,%34,%35,%36,%37,%38,%39,%40,%41,%42,%43,%44,%45,%46,%47," + "%48,%49,%50,%51,%52,%53,%54,%55,%56,%57,%58,%59,%60,%61,%62,%63}, [%64];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), "=r"(r[16]), + "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), "=r"(r[24]), + "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]), "=r"(r[32]), + "=r"(r[33]), "=r"(r[34]), "=r"(r[35]), "=r"(r[36]), "=r"(r[37]), "=r"(r[38]), "=r"(r[39]), "=r"(r[40]), + "=r"(r[41]), "=r"(r[42]), "=r"(r[43]), "=r"(r[44]), "=r"(r[45]), "=r"(r[46]), "=r"(r[47]), "=r"(r[48]), + "=r"(r[49]), "=r"(r[50]), "=r"(r[51]), "=r"(r[52]), "=r"(r[53]), "=r"(r[54]), "=r"(r[55]), "=r"(r[56]), + "=r"(r[57]), "=r"(r[58]), "=r"(r[59]), "=r"(r[60]), "=r"(r[61]), "=r"(r[62]), "=r"(r[63]) + : "r"(taddr) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x8(uint32_t taddr, const uint32_t (&r)[8]) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x8.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" ::"r"(taddr), "r"(r[0]), + "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x16(uint32_t taddr, const uint32_t (&r)[16]) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x16.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16};" :: + "r"(taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x32(uint32_t taddr, const uint32_t (&r)[32]) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x32.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18," + "%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31,%32};" ::"r"(taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]), "r"(r[16]), "r"(r[17]), "r"(r[18]), + "r"(r[19]), "r"(r[20]), "r"(r[21]), "r"(r[22]), "r"(r[23]), "r"(r[24]), "r"(r[25]), "r"(r[26]), "r"(r[27]), + "r"(r[28]), "r"(r[29]), "r"(r[30]), "r"(r[31]) + : "memory"); +} +__device__ __forceinline__ float ex2_approx(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float mul_ftz(float a, float b) { + float y; + asm("mul.ftz.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ float rcp_approx(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ uint32_t cvt_e4m3(float a) { + uint32_t out; + asm("{\n .reg .b16 pair;\n .reg .f32 zero;\n mov.f32 zero, 0f00000000;\n" + " cvt.rn.satfinite.e4m3x2.f32 pair, zero, %1;\n cvt.u32.u16 %0, pair;\n}" + : "=r"(out) + : "f"(a)); + return out & 0xFFu; +} +__device__ __forceinline__ float e4m3_to_f32(uint32_t v) { + float out; + asm("{\n .reg .b16 pair;\n .reg .b32 h2;\n .reg .b16 lo, hi;\n cvt.u16.u32 pair, %1;\n" + " cvt.rn.f16x2.e4m3x2 h2, pair;\n mov.b32 {lo, hi}, h2;\n cvt.f32.f16 %0, lo;\n}" + : "=f"(out) + : "r"(v)); + return out; +} +__device__ __forceinline__ uint32_t cvt_e2m1x8(const float* v) { + uint32_t out; + asm("{\n .reg .b8 b0, b1, b2, b3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\n cvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\n cvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" + " mov.b32 %0, {b0, b1, b2, b3};\n}" + : "=r"(out) + : "f"(v[0]), "f"(v[1]), "f"(v[2]), "f"(v[3]), "f"(v[4]), "f"(v[5]), "f"(v[6]), "f"(v[7])); + return out; +} +// ---- router top-k, bit-exact with FlashInfer's single-block routing (routingCustom, RenormalizeNaive on bf16): +// key = bf16 order (cub float twiddle) above 65535 - expert, i.e. the higher score first, the smaller expert on ties. +__device__ __forceinline__ uint32_t topk_key(uint32_t bits, uint32_t expert) { + const uint32_t twiddled = bits ^ ((bits & 0x8000u) ? 0xFFFFu : 0x8000u); + return (twiddled << 16) | (65535u - expert); +} +__device__ __forceinline__ void topk_unkey(uint32_t key, uint32_t& bits, uint32_t& expert) { + expert = 65535u - (key & 0xFFFFu); + const uint32_t twiddled = key >> 16; + bits = twiddled ^ ((twiddled & 0x8000u) ? 0x8000u : 0xFFFFu); +} +__device__ __forceinline__ uint32_t max16(const uint32_t (&k)[16]) { + const uint32_t a = max(max(max(k[0], k[1]), max(k[2], k[3])), max(max(k[4], k[5]), max(k[6], k[7]))); + const uint32_t b = max(max(max(k[8], k[9]), max(k[10], k[11])), max(max(k[12], k[13]), max(k[14], k[15]))); + return max(a, b); +} +__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_sum_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.add.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_topk_of_lanes(uint32_t v, uint32_t lane) { + uint32_t mine = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(v); + v = v == top ? 0u : v; + mine = lane == k ? top : mine; + } + return mine; +} +__device__ __forceinline__ uint32_t count_greater(const uint32_t* s_cand, uint32_t v) { + uint32_t r = 0; +#pragma unroll + for (int j = 0; j < 32; j += 4) { + const uint4 q = *reinterpret_cast(s_cand + j); + r += (q.x > v) + (q.y > v) + (q.z > v) + (q.w > v); + } + return r; +} +// Top-10 keys of 512 (16 per lane), lane k < 10 holding the k-th largest (keys unique and nonzero). Ranks come from +// all-pairs compares over shared memory rather than chains of warp reductions. +__device__ __forceinline__ uint32_t warp_top10(uint32_t (&keys)[16], uint32_t lane, uint32_t* s_cand) { + // bound = the 10th largest lane maximum; every top-10 key is >= bound + const uint32_t m = max16(keys); + s_cand[lane] = m; + __syncwarp(); + const uint32_t r = count_greater(s_cand, m); + const uint32_t bound = __shfl_sync(0xffffffffu, m, __ffs(__ballot_sync(0xffffffffu, r == kTopK - 1)) - 1); + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) count += keys[i] >= bound; + // exclusive prefix of count over the lanes, from the ballots of its bits (count <= 16) + const uint32_t lt = (1u << lane) - 1u; + uint32_t offset = 0, total = 0; +#pragma unroll + for (int b = 0; b < 5; ++b) { + const uint32_t bb = __ballot_sync(0xffffffffu, (count >> b) & 1u); + offset += (uint32_t)__popc(bb & lt) << b; + total += (uint32_t)__popc(bb) << b; + } + if (total <= 32) { + __syncwarp(); // every lane has read the lane maxima +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) { + if (keys[i] >= bound) s_cand[offset++] = keys[i]; + } + if (lane >= total) s_cand[lane] = 0u; + __syncwarp(); + const uint32_t cand = s_cand[lane]; + const uint32_t rk = count_greater(s_cand, cand); + __syncwarp(); + if (lane < total && rk < kTopK) s_cand[rk] = cand; + __syncwarp(); + return lane < kTopK ? s_cand[lane] : 0u; + } + uint32_t mine = 0; +#pragma unroll 1 + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(max16(keys)); +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) keys[i] = keys[i] == top ? 0u : keys[i]; + mine = lane == k ? top : mine; + } + return mine; +} +// Top-10 keys of 512 (16 per lane), lane k holding the k-th largest (keys unique and nonzero). +__device__ __forceinline__ uint32_t warp_top10_redux(uint32_t (&keys)[16], uint32_t lane, uint32_t* s_cand) { + uint32_t v = max16(keys); + uint32_t bound = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + bound = warp_max_u32(v); + v = v == bound ? 0u : v; + } + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) count += keys[i] >= bound; + const uint32_t total = warp_sum_u32(count); + if (total <= 32) { + uint32_t offset = count; +#pragma unroll + for (uint32_t d = 1; d < 32; d <<= 1) { + const uint32_t up = __shfl_up_sync(0xffffffffu, offset, d); + offset += lane >= d ? up : 0u; + } + offset -= count; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) { + if (keys[i] >= bound) s_cand[offset++] = keys[i]; + } + __syncwarp(); + const uint32_t cand = lane < total ? s_cand[lane] : 0u; + __syncwarp(); + return warp_topk_of_lanes(cand, lane); + } + uint32_t mine = 0; +#pragma unroll 1 + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(max16(keys)); +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) keys[i] = keys[i] == top ? 0u : keys[i]; + mine = lane == k ? top : mine; + } + return mine; +} +// The top-10 keys of one logits row (16 bf16 a lane in v0, v1), lane k < 10 holding the k-th largest +__device__ __forceinline__ uint32_t row_top10(const uint4 v0, const uint4 v1, uint32_t lane, uint32_t* s_cand) { + uint32_t keys[16]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t w = i < 4 ? (&v0.x)[i] : (&v1.x)[i - 4]; + keys[2 * i] = topk_key(w & 0xFFFFu, lane * 16 + 2 * i); + keys[2 * i + 1] = topk_key(w >> 16, lane * 16 + 2 * i + 1); + } + return TOPK_RANK ? warp_top10(keys, lane, s_cand) : warp_top10_redux(keys, lane, s_cand); +} +// FlashInfer calcSoftmax (bf16 scores) as built with -use_fast_math: expf -> ex2.approx.ftz of x * log2(e), +// division -> div.full.ftz; reductions are cooperative_groups::reduce like the reference. +__device__ __forceinline__ float fast_expf(float x) { + float y; + asm("{\n .reg .f32 t;\n mul.ftz.f32 t, %1, 0f3FB8AA3B;\n ex2.approx.ftz.f32 %0, t;\n}" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float fast_div(float a, float b) { + float y; + asm("div.full.ftz.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ uint16_t softmax_topk(uint16_t score_bits, uint32_t lane) { + namespace cg = cooperative_groups; + const auto warp = cg::tiled_partition<32>(cg::this_thread_block()); + const float si = __uint_as_float((uint32_t)score_bits << 16); + float max_score = -INFINITY; + if (lane < kTopK) max_score = si >= max_score ? si : max_score; + max_score = cg::reduce(warp, max_score, cg::greater()); + float sum_score = 0.f, new_score = 0.f; + if (lane < kTopK) { + new_score = fast_expf(__fsub_rn(si, max_score)); + sum_score = __fadd_rn(sum_score, new_score); + } + sum_score = cg::reduce(warp, sum_score, cg::plus()); + uint16_t out = score_bits; + if (lane < kTopK) { + const __nv_bfloat16 b = __float2bfloat16_rn(fast_div(new_score, sum_score)); + out = *reinterpret_cast(&b); + } + return out; +} + +__device__ __forceinline__ uint64_t policy_evict_first() { + uint64_t pol; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(pol)); + return pol; +} +__device__ __forceinline__ void tma_load_2d_w(void* dst, const void* tmap, uint64_t* bar, int c0, int c1, uint64_t pol) { + if (L2_EVICT_FIRST) + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%3, %4}], [%2], %5;" ::"r"(smem_u32(dst)), + "l"(tmap), "r"(smem_u32(bar)), "r"(c0), "r"(c1), "l"(pol) + : "memory"); + else + tma_load_2d(dst, tmap, bar, c0, c1); +} +__device__ __forceinline__ void bulk_load_w(void* dst, const void* src, uint32_t bytes, uint64_t* bar, uint64_t pol) { + if (L2_EVICT_FIRST) + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)), "l"(pol) + : "memory"); + else + bulk_load(dst, src, bytes, bar); +} + +__device__ __forceinline__ void stamp(const Params& p, int slot) { + if (p.trace) p.trace[blockIdx.x * kTraceStride + slot] = (long long)gtime(); +} +// stamp after `dep` is available (the branch needs its value; keys are never 0) +__device__ __forceinline__ void stamp_dep(const Params& p, int slot, uint32_t dep) { + if (dep == 0u) __trap(); + stamp(p, slot); +} +__device__ __forceinline__ void tma_prefetch_2d(const void* tmap, int c0, int c1) { + asm volatile("cp.async.bulk.prefetch.tensor.2d.L2.global.tile [%0, {%1, %2}];" ::"l"(tmap), "r"(c0), "r"(c1) : "memory"); +} +// Physical W13 column (accumulator column) holding reordered row q: block-32 shuffle of FlashInfer's +// get_shuffle_matrix_a_row_indices (epilogue_tile_m = 128); reordered row 2j is up_j, 2j + 1 is gate_j. +__device__ __forceinline__ int w13_col(int q) { return (q & ~31) + ((q & 31) % 4) * 8 + (q & 31) / 4; } +__host__ __device__ constexpr int w13_logical(int p) { + return ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) % 2 ? ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) / 2 + : 128 + ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) / 2; +} + + +// FC2 drain of one pass: TMEM lane quarter q (tiles 2 q, 2 q + 1 of the pass), 256 accumulator columns +__device__ __forceinline__ void reg_fence64(uint32_t (&r)[64]) { // keep uses of r after the tcgen05.wait::ld +#pragma unroll + for (int i = 0; i < 64; ++i) asm volatile("" : "+r"(r[i])); +} + +// FC2 accumulator columns [64 C0, 64 C1) of this warp's lane quarter -> bf16 gemm2_out rows (row r of a tile = the +// r-th row of the item, row info buffer rb) +template +__device__ __forceinline__ void fc2_drain16(const Params& p, Smem& s, uint32_t tm, int q, int lane, int pass, float a2, + int rb) { + const int row = q * 32 + lane, tile = pass * 8 + row / 16, r = row % 16; + const bool valid = r < s.nrows[rb]; + __nv_bfloat16* orow = valid ? p.out + (size_t)(s.rtok[rb][r] * kTopK + s.rkslot[rb][r]) * kH + tile * 256 : nullptr; + const uint32_t base = tm + ((uint32_t)(q * 32) << 16); + auto put = [&](const uint32_t (&r)[64], int c) { + if (valid) { +#pragma unroll + for (int b32 = 0; b32 < 2; ++b32) { + uint32_t pk[16]; +#pragma unroll + for (int l = 0; l < 32; l += 2) { + const float v0 = __uint_as_float(r[b32 * 32 + (l % 4) * 8 + l / 4]) * a2; + const float v1 = __uint_as_float(r[b32 * 32 + ((l + 1) % 4) * 8 + (l + 1) / 4]) * a2; + const __nv_bfloat162 b2 = __floats2bfloat162_rn(v0, v1); + pk[l / 2] = *reinterpret_cast(&b2); + } +#pragma unroll + for (int i = 0; i < 4; ++i) + *reinterpret_cast(orow + c * 64 + b32 * 32 + i * 8) = + make_uint4(pk[4 * i], pk[4 * i + 1], pk[4 * i + 2], pk[4 * i + 3]); + } + } + }; + if (DRAIN_PIPE) { // the next chunk's TMEM load is in flight while this chunk converts and stores + uint32_t ra[64], rb2[64]; + tmem_ld_x64(base + C0 * 64, ra); + tmem_ld_wait(); + reg_fence64(ra); +#pragma unroll + for (int c = C0; c < C1; c += 2) { + if (c + 1 < C1) tmem_ld_x64(base + (uint32_t)((c + 1) * 64), rb2); + put(ra, c); + if (c + 1 < C1) { + tmem_ld_wait(); + reg_fence64(rb2); + if (c + 2 < C1) tmem_ld_x64(base + (uint32_t)((c + 2) * 64), ra); + put(rb2, c + 1); + if (c + 2 < C1) { tmem_ld_wait(); reg_fence64(ra); } + } + } + } else { +#pragma unroll 1 + for (int c = C0; c < C1; ++c) { + uint32_t r[64]; + tmem_ld_x64(base + (uint32_t)(c * 64), r); + tmem_ld_wait(); + put(r, c); + } + } +} + +extern "C" __global__ void __launch_bounds__(kThreads, 1) +q35_mkm_moe_m128(const __grid_constant__ CUtensorMap tm13, const __grid_constant__ CUtensorMap tm2, const Params p) { + extern __shared__ __align__(1024) uint8_t raw[]; +#if SMEM_SPACE // pointer arithmetic on the __shared__ array keeps the shared address space (STS/LDS, not generic) + Smem& s = *reinterpret_cast(raw + ((1024u - (smem_u32(raw) & 1023u)) & 1023u)); +#else + Smem& s = *reinterpret_cast((reinterpret_cast(raw) + 1023) & ~uintptr_t(1023)); +#endif + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + const int slot = blockIdx.x; // the (token, top-k) slot this CTA speculates on; gridDim.x <= M * 10 + const int M = p.num_tokens; + for (int i = threadIdx.x; i < kE; i += kThreads) { + s.first[i] = 0x7FFFFFFF; + s.cnt[i] = 0; + } + if (threadIdx.x < kMaxTok) { + s.fbits[threadIdx.x] = 0u; + s.ounits[threadIdx.x] = 0; + } + if (threadIdx.x == 0) { + stamp(p, 0); + for (int i = 0; i < kStages; ++i) { + mbar_init(&s.full[i], 1); + mbar_init(&s.empty[i], 1); + mbar_init(&s.sf_full[i], 4); + } + // every barrier below completes once per item, so item i waits parity i & 1 + mbar_init(&s.x_full, 1); + mbar_init(&s.fc1_done, 1); + mbar_init(&s.acc_ready, 2); // warps 0 and 4 drained their columns into acc1 + mbar_init(&s.h_ready, 4); + mbar_init(&s.fc2_pass[0], 1); + mbar_init(&s.fc2_pass[1], 1); + mbar_init(&s.fc2_free, 4); + mbar_init(&s.acc_drained, 4); // warps 4..7 read the last FC2 pass: the next item's FC1 may overwrite columns 0..255 + mbar_init(&s.route_done, 1); + fence_mbar_init(); + } + if (warp == 2) { tmem_alloc(&s.tmem, 512); tmem_relinquish(); tc_fence_before(); } + if (warp == 1 && lane == 0) { prefetch_tmap(&tm13); prefetch_tmap(&tm2); } + __syncthreads(); + tc_fence_after(); + const uint32_t tm = s.tmem; + const int my_t = slot / kTopK, my_k = slot % kTopK; + // ring stage g (global over the CTA's items) of item-relative stage st of expert e + auto issue_stage = [&](int g, int st, int e, uint64_t pol) { + const int b = g % kStages; + mbar_expect_tx(&s.full[b], kStageBytes); + uint8_t* dst = s.ring[b]; + if (st < kFc1Stages) { + // the expert's rows as two boxes of 128 rows, K bytes [128 st, 128 st + 128) + tma_load_2d_w(dst, &tm13, &s.full[b], st * 128, e * kN1, pol); + tma_load_2d_w(dst + 16384, &tm13, &s.full[b], st * 128, e * kN1 + 128, pol); + // scale atoms of row blocks 0 and 1, k-blocks [4 st, 4 st + 4) + const uint8_t* sf = p.w13_sf + (size_t)e * (kN1 * kH / 16); + bulk_load_w(dst + 32768, sf + (size_t)(st * 4) * 512, 2048, &s.full[b], pol); + bulk_load_w(dst + 32768 + 2048, sf + (size_t)(64 + st * 4) * 512, 2048, &s.full[b], pol); + } else { + // W2 tiles 2 f and 2 f + 1 (256 rows x 64 B each) and their scale atoms (2 row blocks x 2 k-blocks) + const int f = st - kFc1Stages; +#pragma unroll + for (int tl = 0; tl < 2; ++tl) { + const int row0 = (2 * f + tl) * 256; + tma_load_2d_w(dst + tl * 16384, &tm2, &s.full[b], 0, e * kH + row0, pol); + bulk_load_w(dst + 32768 + tl * 2048, p.w2_sf + (size_t)e * (kH * 8) + (size_t)(row0 / 128) * 2 * 512, 2048, + &s.full[b], pol); + } + } + }; + auto prefetch_stage = [&](int st, int e) { // the same bytes as issue_stage, to L2 only + if (st < kFc1Stages) { + tma_prefetch_2d(&tm13, st * 128, e * kN1); + tma_prefetch_2d(&tm13, st * 128, e * kN1 + 128); + const uint8_t* sf = p.w13_sf + (size_t)e * (kN1 * kH / 16); + bulk_prefetch_l2(sf + (size_t)(st * 4) * 512, 2048); + bulk_prefetch_l2(sf + (size_t)(64 + st * 4) * 512, 2048); + } else { + const int row0 = (st - kFc1Stages) * 512; + tma_prefetch_2d(&tm2, 0, e * kH + row0); + tma_prefetch_2d(&tm2, 0, e * kH + row0 + 256); + bulk_prefetch_l2(p.w2_sf + (size_t)e * (kH * 8) + (size_t)(row0 / 128) * 2 * 512, 4096); + } + }; + // Rows of an item (the tokens of mask in increasing order, at most 16) into row buffer rb, by one warp + auto set_rows = [&](int rb, int it) { + const int n = s.item_n[it], ei = s.item_e[it]; + if (lane < n) { + const int t = s.item_rows[it][lane]; + int ks = 0; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + if (s.ids[t][k] == ei) ks = k; + s.rtok[rb][lane] = t; + s.rkslot[rb][lane] = ks; + } + if (lane == 0) s.nrows[rb] = n; + }; + // their x rows and scale bytes into xa / xsf, by threads tid = 0 .. nt - 1 + auto copy_rows = [&](int rb, int tid, int nt) { + const int n1 = s.nrows[rb] * 128, n = s.nrows[rb] * 144; // 16-B granules: x rows, then their scale bytes + if (GATHER_ASYNC) { + for (int idx = tid; idx < n; idx += nt) { + const uint8_t* src; + uint8_t* dst; + if (idx < n1) { + const int r = idx >> 7, c = (idx >> 3) & 15, q = idx & 7; + src = p.x_fp4 + (size_t)s.rtok[rb][r] * (kH / 2) + c * 128 + q * 16; + dst = s.xa + c * 2048 + (r / 8) * 1024 + (r % 8) * 128 + ((q ^ (r % 8)) * 16); + } else { + const int r = (idx - n1) >> 4, q = (idx - n1) & 15; + src = p.x_sf + (size_t)s.rtok[rb][r] * (kH / 16) + q * 16; + dst = &s.xsf[r][q * 16]; + } + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(smem_u32(dst)), "l"(src) : "memory"); + } + asm volatile("cp.async.wait_all;" ::: "memory"); + fence_proxy_async_smem(); + return; + } + // GATHER_BATCH loads a thread ahead of their stores: the global latency is paid once a batch + for (int i0 = tid; i0 < n; i0 += GATHER_BATCH * nt) { + uint4 v[GATHER_BATCH]; +#pragma unroll + for (int j = 0; j < GATHER_BATCH; ++j) { + const int idx = i0 + j * nt; + if (idx < n1) { // row r, K chunk c (128 B), 16-B granule q of the chunk + const int r = idx >> 7, c = (idx >> 3) & 15, q = idx & 7; + v[j] = *reinterpret_cast(p.x_fp4 + (size_t)s.rtok[rb][r] * (kH / 2) + c * 128 + q * 16); + } else if (idx < n) { + const int r = (idx - n1) >> 4, q = (idx - n1) & 15; + v[j] = *reinterpret_cast(p.x_sf + (size_t)s.rtok[rb][r] * (kH / 16) + q * 16); + } + } +#pragma unroll + for (int j = 0; j < GATHER_BATCH; ++j) { + const int idx = i0 + j * nt; + if (idx < n1) { + const int r = idx >> 7, c = (idx >> 3) & 15, q = idx & 7; + *reinterpret_cast(s.xa + c * 2048 + (r / 8) * 1024 + (r % 8) * 128 + ((q ^ (r % 8)) * 16)) = v[j]; + } else if (idx < n) { + const int r = (idx - n1) >> 4, q = (idx - n1) & 15; + *reinterpret_cast(&s.xsf[r][q * 16]) = v[j]; + } + } + } + fence_proxy_async_smem(); // the tensor core reads xa through the async proxy + }; + // the given routing: warps 0, 2..7 load every slot's expert; warp 1 starts this CTA's weight stream at once on its + // own slot's expert (a wrong guess only wastes its first ring fill) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (threadIdx.x == 32) stamp(p, 10); + if (warp == 1) { + // speculate on this CTA's own slot: its expert's weight stream starts before the work split + const uint32_t espec = (uint32_t)p.ids_in[slot]; + if (lane == 0) { + stamp(p, 9); + const uint64_t pol = policy_evict_first(); + for (int st = 0; st < SPEC_STAGES; ++st) { + issue_stage(st, st, (int)espec, pol); + if (STAGE_STAMPS) stamp(p, 64 + st); + } + // HBM is otherwise idle until the routing is done: the next stages go to L2 now (another CTA streams the + // expert when this slot is not its first) + for (int st = SPEC_STAGES; st < SPEC_STAGES + SPEC_PF && st < kTotStages; ++st) prefetch_stage(st, (int)espec); + stamp(p, 2); + } + mbar_wait(&s.route_done, 0); + } else { + // the routing as given: the 224 routing threads load every slot's expert at once, then record and count them + const int tid = (warp == 0 ? 0 : warp - 1) * 32 + lane, S = M * kTopK; + constexpr int kPer = (kMaxTok * kTopK + 223) / 224; + int ex[kPer]; +#pragma unroll + for (int j = 0; j < kPer; ++j) ex[j] = tid + 224 * j < S ? p.ids_in[tid + 224 * j] : 0; +#pragma unroll + for (int j = 0; j < kPer; ++j) { + const int s2 = tid + 224 * j; + if (s2 < S) { + s.ids[s2 / kTopK][s2 % kTopK] = ex[j]; + atomicMin(&s.first[ex[j]], s2); + atomicAdd(&s.cnt[ex[j]], 1); + } + } + } + if (warp != 1) { + const int base = warp == 0 ? 0 : warp - 1; + if (DIAG_ROUTE && threadIdx.x == 0) stamp(p, 13); + named_bar(4, 224); + if (DIAG_ROUTE && threadIdx.x == 0) stamp(p, 15); + // each routing warp marks its tokens' slots that are their expert's first and the units they hand out (s.first and + // s.cnt are final after the barrier): a first occurrence at slot s2 of an expert with n tokens hands out its + // ceil(n / 16) row groups, less the first one when s2 < G (that CTA runs it) + const int G = (int)gridDim.x; + { // every slot at once over the 224 routing threads (per-token sums by shared-memory atomics) + const int tid = base * 32 + lane, S = M * kTopK; +#pragma unroll + for (int j = 0; j < (kMaxTok * kTopK + 223) / 224; ++j) { + const int s2 = tid + 224 * j; + if (s2 < S) { + const int t = s2 / kTopK, k = s2 % kTopK, e2 = s.ids[t][k]; + const bool f = s.first[e2] == s2; + const int u = f ? (s.cnt[e2] + kRows - 1) / kRows - (s2 < G ? 1 : 0) : 0; + s.uslot[t][k] = (uint8_t)u; + if (f) { + atomicOr(&s.fbits[t], 1u << k); + atomicAdd(&s.ounits[t], u); + } + } + } + } + named_bar(4, 224); + if (DIAG_ROUTE && threadIdx.x == 0) stamp(p, 20); + if (warp == 0) { + // Work items: (expert, row group) units. Slot c's expert's first unit belongs to CTA c when c is the expert's first + // slot (its speculative stream already carries it). The other units, in order of their expert's first slot then + // row group, go to the CTAs whose own slot repeats an earlier expert, in order; units beyond those CTAs become + // further items of CTAs 0, 1, ... Every CTA derives the same assignment from the same routing. + // Per-token prefixes over 128 tokens: lane l holds tokens l + 32 h, h < kGroups. + auto scan = [&](int v) -> int { // inclusive prefix over the lanes +#pragma unroll + for (int d = 1; d < 32; d <<= 1) { + const int up = __shfl_up_sync(0xffffffffu, v, d); + v += lane >= d ? up : 0; + } + return v; + }; + uint32_t fb[kGroups]; + int ou[kGroups], fi[kGroups], oi[kGroups]; + int fcarry = 0, ocarry = 0; +#pragma unroll + for (int h = 0; h < kGroups; ++h) { + const int t = lane + 32 * h; + fb[h] = t < M ? s.fbits[t] : 0u; + ou[h] = t < M ? s.ounits[t] : 0; + fi[h] = scan(__popc(fb[h])) + fcarry; + oi[h] = scan(ou[h]) + ocarry; + fcarry = __shfl_sync(0xffffffffu, fi[h], 31); + ocarry = __shfl_sync(0xffffffffu, oi[h], 31); + } + const int Uf = fcarry, O = ocarry; // distinct experts, units handed out + // token t's exclusive prefixes and first bits (t < M; h = t / 32 is uniform over the warp) + auto tok = [&](int t, int& fex, int& oex, uint32_t& ftb) { + const int h = t >> 5, l = t & 31; + int fin = 0, oin = 0, our = 0; + uint32_t f = 0; +#pragma unroll + for (int g = 0; g < kGroups; ++g) { + if (g == h) { + fin = fi[g]; + oin = oi[g]; + our = ou[g]; + f = fb[g]; + } + } + fin = __shfl_sync(0xffffffffu, fin, l); + oin = __shfl_sync(0xffffffffu, oin, l); + our = __shfl_sync(0xffffffffu, our, l); + ftb = __shfl_sync(0xffffffffu, f, l); + oex = oin - our; + fex = fin - __popc(ftb); + }; + auto before = [&](int s2) -> int { // first-occurrence slots before slot s2 (s2 <= M * 10) + const int t = s2 / kTopK; + if (t >= M) return Uf; + int fex, oex; + uint32_t ftb; + tok(t, fex, oex, ftb); + return fex + __popc(ftb & ((1u << (s2 % kTopK)) - 1u)); + }; + // the n-th (0-based) handed-out unit: (expert, row group) + auto nth_unit = [&](int n, int& e, int& grp) { + int t = 0; +#pragma unroll + for (int h = 0; h < kGroups; ++h) t += __popc(__ballot_sync(0xffffffffu, lane + 32 * h < M && oi[h] <= n)); + int fex, oex; + uint32_t ftb; + tok(t, fex, oex, ftb); + int k = 0, left = n - oex; + while (left >= (int)s.uslot[t][k]) left -= s.uslot[t][k++]; + e = s.ids[t][k]; + grp = left + (t * kTopK + k < G ? 1 : 0); // the first group of a slot below G is its own CTA's + }; + const int nf_c = before(slot), nf_g = before(G); + if (DIAG_ROUTE && lane == 0) stamp(p, 21); + const int xc = s.ids[my_t][my_k]; + const bool owner = s.first[xc] == slot; + const int D = G - nf_g, dc = slot - nf_c; // duplicate CTAs; this one's rank among them + // the whole warp walks the items (nth_unit is warp-collective); the lanes write each item's token list + int ni = 0; + auto add = [&](int e, int grp) { + if (ni >= kMaxItems) __trap(); + int before_g = 0; // the expert's tokens in earlier 32-token groups +#pragma unroll + for (int h = 0; h < kGroups; ++h) { + const int t = lane + 32 * h; + bool has = false; +#pragma unroll + for (int k = 0; k < kTopK; ++k) has |= t < M && s.ids[t][k] == e; + const uint32_t bal = __ballot_sync(0xffffffffu, has); + const int r = before_g + __popc(bal & ((1u << lane) - 1u)) - grp * kRows; // row within the group + if (has && r >= 0 && r < kRows) s.item_rows[ni][r] = (uint8_t)t; + before_g += __popc(bal); + } + if (lane == 0) { + s.item_e[ni] = e; + s.item_n[ni] = min(kRows, before_g - grp * kRows); + } + ++ni; + }; + if (owner) { + add(xc, 0); + } else if (dc < O) { + int e, grp; + nth_unit(dc, e, grp); + add(e, grp); + } + for (int o = D + slot; o < O; o += G) { + int e, grp; + nth_unit(o, e, grp); + add(e, grp); + } + __syncwarp(); + const int nall = ni; + if (lane == 0) { + s.nitems = nall; + s.spec_ok = owner ? 1 : 0; + stamp(p, 1); + if (DIAG_ROUTE && p.trace) { // DIAG values (not times): 22 = items, 23 = speculation kept + p.trace[blockIdx.x * kTraceStride + 22] = nall; + p.trace[blockIdx.x * kTraceStride + 23] = owner ? 1 : 0; + } + mbar_arrive(&s.route_done); + } + } + named_bar(4, 224); + } + const int nitems = s.nitems; + if (nitems == 0) { // a duplicate slot with no work: let the speculative loads land, then leave + if (warp == 1) { + if (lane == 0) + for (int b = 0; b < SPEC_STAGES; ++b) mbar_wait(&s.full[b], 0); + __syncwarp(); + } + named_bar(7, kThreads); + if (warp == 2) tmem_dealloc(tm, 512); + return; + } + tc_fence_after(); + const bool spec_ok = s.spec_ok != 0; + // ring stage of item 0's stage 0: the speculative fill (stages 0 .. SPEC_STAGES - 1) is item 0's own only if spec_ok + const int base0 = spec_ok ? 0 : SPEC_STAGES; + float ag = 0.f, sc = 0.f, a2 = 0.f; + auto load_scales = [&](int it) { + const int e = s.item_e[it]; + ag = p.g1_alpha[e]; + sc = p.g1_scale_c[e]; + a2 = p.g2_alpha[e]; + }; + // FC1 drain of quarter-0 accumulator columns [64 c0, 64 c0 + 128) (TMEM lanes 0..15 hold the item rows) into acc1 + auto drain_fc1 = [&](int c0) { +#pragma unroll 1 + for (int c = c0; c < c0 + 2; ++c) { + uint32_t r[64]; + tmem_ld_x64(tm + (uint32_t)(c * 64), r); + tmem_ld_wait(); + if (lane < 16) { +#pragma unroll + for (int i = 0; i < 16; ++i) + *reinterpret_cast(s.acc1 + lane * kARow + c * 64 + 4 * i) = + make_uint4(r[4 * i], r[4 * i + 1], r[4 * i + 2], r[4 * i + 3]); + } + } + tc_fence_before(); + __syncwarp(); + }; + // SwiGLU + NVFP4 epilogue (h = 0: the one FC1 pass), tasks (row, 16-wide block, half) = 16 x 8 x 2 = 256, + // task = et + k * nt. + // Accumulator columns of each 32-block: [0, 8) up 2m, [8, 16) gate 2m, [16, 24) up 2m + 1, [24, 32) gate 2m + 1 + auto epi_tasks = [&](int h, int et, int nt) { +#pragma unroll 1 + for (int task = et; task < 256; task += nt) { + const int row = task / 16, blk = (task / 2) % 8, hf = task % 2; + const float4* ar = reinterpret_cast(s.acc1 + row * kARow + 32 * blk + 4 * hf); + const float4 u_even = ar[0], g_even = ar[2], u_odd = ar[4], g_odd = ar[6]; + const float ga[8] = {g_even.x, g_odd.x, g_even.y, g_odd.y, g_even.z, g_odd.z, g_even.w, g_odd.w}; + const float ua[8] = {u_even.x, u_odd.x, u_even.y, u_odd.y, u_even.z, u_odd.z, u_even.w, u_odd.w}; + // trtllm-gen's SwiGLU + NVFP4 epilogue op for op (from its SASS); a = FC1 dequant, c = a x FC2 input scale: + // h' = (u a)(g sigmoid(g a)), sigmoid = rcp(ex2(-(g (a log2e))) + 1), block scale e4m3(max|h'| c / 6), FP4 (h' / scale) c + const float a_l2e = mul_ftz(ag, 1.4426950216293334961f); + float hv[8]; + float amax = 0.f; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const float ex = ex2_approx(-(ga[i] * a_l2e)); + const float sg = rcp_approx(ex + 1.0f); + const float v = fmaf(ua[i], ag, 0.0f) * (ga[i] * sg); + hv[i] = v; + amax = fmaxf(amax, fabsf(v)); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + const uint32_t sf = cvt_e4m3(fmaxf(0.f, fabsf(amax * sc)) * 0.16666667163372039795f); + const float sff = e4m3_to_f32(sf); + const float osc = sff == 0.f ? 0.f : rcp_approx(sff); + float qv[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) qv[i] = mul_ftz(hv[i], osc) * sc; + const uint32_t packed = cvt_e2m1x8(qv); + const int gblk = 8 * h + blk; // 16-wide block over the 128 intermediates; k-blocks 0, 1 + const int chunk = gblk / 2, r8 = row % 8; + const uint32_t off = (uint32_t)((row / 8) * 1024 + r8 * 128 + ((chunk ^ r8) * 16) + (gblk % 2) * 8 + hf * 4); +#pragma unroll + for (int cp = 0; cp < 8; ++cp) *reinterpret_cast(s.ha + cp * 2048 + off) = packed; + if (hf == 0) reinterpret_cast(&s.hsf[row][gblk / 4])[gblk % 4] = (uint8_t)sf; + } + }; + // Rows of item it (its tokens in increasing order, at most 16) and their x rows / scale bytes, by warps 0 and 3; + // row info is double-buffered by item parity (the previous item's FC2 drains still read theirs). + auto gather = [&](int it) { + const int rb = it & 1; + if (warp == 0) set_rows(rb, it); + named_bar(5, 64); + copy_rows(rb, (warp == 0 ? 0 : 32) + lane, 64); + named_bar(5, 64); + if (warp == 0 && lane == 0) { + mbar_arrive(&s.x_full); + if (DIAG_ROUTE && it == 0) stamp(p, 14); + } + }; + + if (warp == 1) { + if (lane == 0) { + // ---------------- TMA producer: every item's weight stream through the ring + const uint64_t pol = policy_evict_first(); + for (int it = 0; it < nitems; ++it) { + const int e = s.item_e[it], base = base0 + it * kTotStages; + const int st0 = (it == 0 && spec_ok) ? SPEC_STAGES : 0; + for (int st = st0; st < kTotStages; ++st) { + const int g = base + st; + if (FC2_PF > 0 && st == kFc1Stages - FC2_PF_LEAD) { + // the FC2 tiles beyond the ring go to L2 now, while the last FC1 stages and h are in progress + for (int j = kStages; j < kStages + FC2_PF; ++j) prefetch_stage(kFc1Stages + j, e); + } + if (PF_SET != 0 && ((st == st0 && (PF_SET & 1)) || (st == kFc1Stages + 4 && (PF_SET & 4)) || + (st == kFc1Stages + 6 && (PF_SET & 8)))) { + // these stages wait for MMAs that are slow to start after a phase change; let HBM fetch them meanwhile + const int n = st == st0 ? PF_N1 : st == kFc1Stages + 4 ? PF_N2 : 4; + for (int k = st; k < st + n && k < kTotStages; ++k) prefetch_stage(k, e); + } + if (g >= kStages) mbar_wait(&s.empty[g % kStages], ((g / kStages) - 1) & 1); + issue_stage(g, st, e, pol); + if (STAGE_STAMPS && it == 0) stamp(p, 64 + st); + } + } + } + __syncwarp(); + } else if (warp == 2) { + if (lane == 0) { + // ---------------- MMA issuer (tensor pipe: MMAs only) + const uint32_t idesc = make_idesc_nvf4(128u, 256u); + if (!spec_ok) { // the speculative fill was another CTA's expert: recycle its ring stages without MMAs + for (int g = 0; g < SPEC_STAGES; ++g) { + mbar_wait(&s.full[g % kStages], (g / kStages) & 1); + mbar_wait(&s.sf_full[g % kStages], (g / kStages) & 1); + mbar_arrive(&s.empty[g % kStages]); + } + } + for (int it = 0; it < nitems; ++it) { + const int base = base0 + it * kTotStages, par = it & 1; + if (it > 0) mbar_wait(&s.acc_drained, (it - 1) & 1); + mbar_wait(&s.x_full, par); + tc_fence_after(); + int n = 0; + for (int st = 0; st < kFc1Stages; ++st) { + const int g = base + st, b = g % kStages, ks = st; + mbar_wait(&s.full[b], (g / kStages) & 1); + mbar_wait(&s.sf_full[b], (g / kStages) & 1); + tc_fence_after(); + if (st == 0 && it == 0) stamp(p, 3); + if (st == 0 && it > 0 && it == nitems - 1) stamp(p, 11); + if (DIAG_ROUTE && st == 0 && it < 4) stamp(p, 24 + it); // DIAG: item it's first FC1 MMA + if (DIAG_ROUTE && st == kFc1Stages - 1 && it < 4) stamp(p, 28 + it); // DIAG: its last FC1 stage landed +#pragma unroll + for (int q = 0; q < 4; ++q, ++n) { + const int kb = ks * 4 + q; // k-blocks in K order: one accumulator, the stock summation order + umma_nvf4(tm, make_sdesc(smem_u32(s.xa + (kb / 4) * 2048) + (kb % 4) * 32, 0, 1024, 2), + make_sdesc(smem_u32(s.ring[b]) + q * 32, 0, 1024, 2), idesc, + tm + kColSfaX + (uint32_t)(kb % kSfaSlots) * 4, tm + kColSfb + (uint32_t)(n % kSfbSlots) * 8, + kb > 0); + } + tc_commit(&s.empty[b]); + } + tc_commit(&s.fc1_done); + mbar_wait(&s.h_ready, par); + tc_fence_after(); + if (it == 0) stamp(p, 6); + for (int st = kFc1Stages; st < kTotStages; ++st) { + const int g = base + st, b = g % kStages, f = st - kFc1Stages; + if (f == 4) { mbar_wait(&s.fc2_free, par); tc_fence_after(); } + mbar_wait(&s.full[b], (g / kStages) & 1); + mbar_wait(&s.sf_full[b], (g / kStages) & 1); + tc_fence_after(); +#pragma unroll + for (int tl = 0; tl < 2; ++tl) { + const int j = 2 * f + tl, jt = j % 8; +#pragma unroll + for (int kb = 0; kb < 2; ++kb) { + const int n2 = 64 + j * 2 + kb; // scale-factor slot, numbered on from FC1's 64 k-blocks + umma_nvf4(tm, make_sdesc(smem_u32(s.ha) + kb * 32, 0, 1024, 2), + make_sdesc(smem_u32(s.ring[b] + tl * 16384) + kb * 32, 0, 512, 4), idesc, + tm + kColSfa + (uint32_t)(jt * 2 + kb) * 4, tm + kColSfb + (uint32_t)(n2 % kSfbSlots) * 8, + (jt | kb) != 0); + } + } + tc_commit(&s.empty[b]); + if (f % 4 == 3) tc_commit(&s.fc2_pass[f / 4]); + } + } + } + __syncwarp(); + } else if (warp == 0 || warp == 3) { + // ---------------- x gather, FC1 drains (warp 0: columns [0, 128); warp 4 drains [128, 256)) and the epilogue + const int et = (warp == 0 ? 0 : 32) + lane; // 0..63 + gather(0); + for (int it = 0; it < nitems; ++it) { + const int par = it & 1; + load_scales(it); + if (warp == 0) { + mbar_wait(&s.fc1_done, par); + tc_fence_after(); + if (lane == 0 && it == 0) stamp(p, 4); + drain_fc1(0); + if (lane == 0) { + mbar_arrive(&s.acc_ready); + if (it == 0 && !DIAG_ROUTE) stamp(p, 13); + } + } + mbar_wait(&s.acc_ready, par); + epi_tasks(0, et, 192); + if (warp == 3 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 14); + fence_proxy_async_smem(); + named_bar(6, 192); // h and its scale words are complete for the TMEM warpgroup + if (it + 1 < nitems) gather(it + 1); // FC1 of item it is done: x and its scale bytes are free + } + } else { + // ---------------- TMEM-side warpgroup (warps 4..7): q = lane quarter (TMEM lanes 32 q .. 32 q + 31) + const int q = warp - 4; + const uint32_t lane_addr = (uint32_t)(q * 32) << 16; + if (!spec_ok) { // the speculative fill is recycled without MMAs: release its scale-factor side + for (int g = 0; g < SPEC_STAGES; ++g) { + mbar_wait(&s.full[g % kStages], (g / kStages) & 1); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[g % kStages]); + } + } + for (int it = 0; it < nitems; ++it) { + const int base = base0 + it * kTotStages, par = it & 1, rb = it & 1; + load_scales(it); + mbar_wait(&s.x_full, par); + const int nrows = s.nrows[rb]; + auto stage_sf2 = [&](int st) { // FC2 tiles 2 f, 2 f + 1: k-blocks 0, 1 of their row blocks (rows 0..127 | 128..255) + const int g = base + st, b = g % kStages; + mbar_wait(&s.full[b], (g / kStages) & 1); + if (STAGE_STAMPS && threadIdx.x == 128 && it == 0) stamp(p, 16 + st); + if (g >= 4) mbar_wait(&s.empty[(g - 4) % kStages], ((g - 4) / kStages) & 1); + const uint8_t* sf = s.ring[b] + 32768; + const int n0 = 64 + (st - kFc1Stages) * 4; + for (int i = 0; i < 4; ++i) { + const int tl = i / 2, kb = i % 2; + const uint4 v0 = *reinterpret_cast(sf + tl * 2048 + kb * 512 + lane * 16); + const uint4 v1 = *reinterpret_cast(sf + tl * 2048 + 1024 + kb * 512 + lane * 16); + const uint32_t r[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; + tmem_st_x8(tm + lane_addr + kColSfb + (uint32_t)((n0 + i) % kSfbSlots) * 8, r); + } + tmem_st_wait(); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[b]); + }; + // FC1 stage st: weight scale factors of k-blocks [4 ks, 4 ks + 4) (8 columns each: row blocks 2 h | 2 h + 1) and + // the x scale factors of the same k-blocks (row m -> lane m % 32, column m / 32 of each 4-column atom, every + // quarter holds a copy; rows >= nrows, including the A rows read past the 16 real ones, get zero scales) + for (int st = 0; st < kFc1Stages; ++st) { + const int g = base + st, b = g % kStages, ks = st; + mbar_wait(&s.full[b], (g / kStages) & 1); + if (STAGE_STAMPS && threadIdx.x == 128 && it == 0) stamp(p, 16 + st); + if (g >= 4) mbar_wait(&s.empty[(g - 4) % kStages], ((g - 4) / kStages) & 1); + const uint8_t* sf = s.ring[b] + 32768; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const uint4 v0 = *reinterpret_cast(sf + i * 512 + lane * 16); + const uint4 v1 = *reinterpret_cast(sf + 2048 + i * 512 + lane * 16); + const uint32_t r[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; + tmem_st_x8(tm + lane_addr + kColSfb + (uint32_t)((st * 4 + i) % kSfbSlots) * 8, r); + } + uint32_t xr[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int kb = ks * 4 + i / 4; + xr[i] = (i % 4 == 0 && lane < nrows) ? *reinterpret_cast(&s.xsf[lane][kb * 4]) : 0u; + } + tmem_st_x16(tm + lane_addr + kColSfaX + (uint32_t)((ks * 4) % kSfaSlots) * 4, xr); + tmem_st_wait(); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[b]); + } + // EARLY_SF2: stage the first FC2 scale factors before h (they wait for FC2 data to land, which delays h) + for (int st = kFc1Stages; st < kFc1Stages + EARLY_SF2; ++st) stage_sf2(st); + if (q == 0) { // drain the upper 128 columns + mbar_wait(&s.fc1_done, par); + tc_fence_after(); + drain_fc1(2); + if (lane == 0) mbar_arrive(&s.acc_ready); + } + mbar_wait(&s.acc_ready, par); + epi_tasks(0, 64 + (warp - 4) * 32 + lane, 192); + fence_proxy_async_smem(); + named_bar(6, 192); // h complete + if (threadIdx.x == 128 && it == 0) stamp(p, 5); + // h scale factors: 16 atoms (tile jt, k-block kb) of 4 columns; row m = 32 c + lane is nonzero only in its tile + { + uint32_t r[32]; +#pragma unroll 1 + for (int half = 0; half < 2; ++half) { +#pragma unroll + for (int i = 0; i < 32; ++i) { + const int a = half * 8 + i / 4, col = i % 4; // atom a = 2 jt + kb + const int jt = a / 2, kb = a % 2; + const int row = 32 * col + lane; + r[i] = (row / 16 == jt && row % 16 < nrows) ? s.hsf[row % 16][kb] : 0u; + } + tmem_st_x32(tm + lane_addr + kColSfa + half * 32, r); + } + tmem_st_wait(); + } + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.h_ready); + if (q == 0 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 15); + for (int st = kFc1Stages + EARLY_SF2; st < kTotStages; ++st) { + if (st == kFc1Stages + 4) { + // pass 0 is fully staged; drain it once its MMAs are done, then free the accumulator for pass 1 + mbar_wait(&s.fc2_pass[0], par); + tc_fence_after(); + if (q == 0 && lane == 0 && it == 0) stamp(p, 7); + fc2_drain16<0, 4>(p, s, tm, q, lane, 0, a2, rb); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.fc2_free); + } + stage_sf2(st); + } + mbar_wait(&s.fc2_pass[1], par); + tc_fence_after(); + if (q == 0 && lane == 0 && it == 0) stamp(p, 12); + if (TAIL8 && it + 1 == nitems) fc2_drain16<0, 2>(p, s, tm, q, lane, 1, a2, rb); + else fc2_drain16<0, 4>(p, s, tm, q, lane, 1, a2, rb); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.acc_drained); + } + } + if (TAIL8 && warp < 4) { // the last item's last pass, upper columns: every warp is free now (TMEM lane quarter = warp) + const int it = nitems - 1; + load_scales(it); + mbar_wait(&s.fc2_pass[1], it & 1); + tc_fence_after(); + fc2_drain16<2, 4>(p, s, tm, warp, lane, 1, a2, it & 1); + } + tc_fence_before(); + named_bar(7, kThreads); // every warp, including the drains, before TMEM is released + if (threadIdx.x == 128) stamp(p, 8); + if (warp == 2) tmem_dealloc(tm, 512); +} + +static_assert(sizeof(Smem) + 1024 == Q35_MKM_SMEM_BYTES, "the launcher's dynamic shared memory size"); + +} // namespace qwen_mkm_moe_m128 + + +} // namespace sglang diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m32.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m32.cu new file mode 100644 index 00000000..1e066a50 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m32.cu @@ -0,0 +1,1346 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Fused small-M NVFP4 MoE expert phase for Qwen3.5 TP8 decode verify (M <= 32 tokens, top-10, 512 experts, +// hidden 4096, intermediate 128 per rank): the top-10 of the router logits in the trtllm-gen routing's order and +// weights, FC1, SwiGLU, the NVFP4 requantization and FC2 in one launch over the trtllm-gen weight layout. Writes the +// deferred-finalize rows (row t * 10 + k, identity permutation) and the expert weights. +// +// The 16-token engine (mkm_moe_m16.cu) on min(M * 10, SMs) persistent CTAs, as the MKM line's 21-token TP4 engine +// runs them: CTA c speculates on slot c's expert, experts first seen past the grid go to the CTAs whose slot repeats +// an earlier expert, and the rest become further items of CTAs 0, 1, ... Each item +// gathers its own token rows (at most 16; more are split into row groups), so the per-element arithmetic and the +// summation order stay the 16-token engine's and the rows stay bit for bit trtllm-gen's. +// +// Built with NVRTC: the kernel has C linkage (q35_mkm_moe_m32), and the launcher's dynamic shared memory size +// (Q35_MKM_SMEM_BYTES) is checked against the layout at compile time. +#include +#include +#include + +// NVRTC has no , or +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef int int32_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef unsigned long long uintptr_t; +typedef struct __align__(64) { + uint64_t opaque[16]; +} CUtensorMap; +#ifndef INFINITY +#define INFINITY __int_as_float(0x7f800000) +#endif +// Minimal SM100 (tcgen05 / TMA / mbarrier) PTX helpers for the small-M NVFP4 MoE engine. + +namespace qwen35 { + +__device__ __forceinline__ uint32_t smem_u32(const void* p) { + return static_cast(__cvta_generic_to_shared(p)); +} + +// ---------------------------------------------------------------- mbarrier +__device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(smem_u32(bar)), "r"(count) : "memory"); +} +__device__ __forceinline__ void fence_mbar_init() { + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} +__device__ __forceinline__ void mbar_expect_tx(uint64_t* bar, uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smem_u32(bar)), "r"(bytes) : "memory"); +} +__device__ __forceinline__ void mbar_arrive(uint64_t* bar) { + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(smem_u32(bar)) : "memory"); +} +__device__ __forceinline__ bool mbar_try_wait(uint64_t* bar, uint32_t phase) { + uint32_t ok; + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ok) + : "r"(smem_u32(bar)), "r"(phase) + : "memory"); + return ok != 0; +} +__device__ __forceinline__ void mbar_wait(uint64_t* bar, uint32_t phase) { + while (!mbar_try_wait(bar, phase)) { + } +} + +// ---------------------------------------------------------------- TMA / bulk copies +__device__ __forceinline__ void tma_load_2d(void* dst, const void* tmap, uint64_t* bar, int c0, int c1) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes [%0], [%1, {%3, %4}], [%2];" :: + "r"(smem_u32(dst)), + "l"(tmap), "r"(smem_u32(bar)), "r"(c0), "r"(c1) + : "memory"); +} +__device__ __forceinline__ void bulk_load(void* dst, const void* src, uint32_t bytes, uint64_t* bar) { + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)) + : "memory"); +} +__device__ __forceinline__ void prefetch_tmap(const void* tmap) { + asm volatile("prefetch.tensormap [%0];" ::"l"(tmap) : "memory"); +} +// generic-proxy smem writes -> visible to the async proxy (UMMA / UTCCP operand reads) +__device__ __forceinline__ void fence_proxy_async_smem() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + +// ---------------------------------------------------------------- tcgen05 +__device__ __forceinline__ void tmem_alloc(uint32_t* dst_smem, uint32_t ncols) { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smem_u32(dst_smem)), + "r"(ncols) + : "memory"); +} +__device__ __forceinline__ void tmem_relinquish() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;" ::: "memory"); +} +__device__ __forceinline__ void tmem_dealloc(uint32_t taddr, uint32_t ncols) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(ncols) : "memory"); +} +__device__ __forceinline__ void tc_fence_before() { + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); +} +__device__ __forceinline__ void tc_fence_after() { + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); +} +// arrive on `bar` when all prior tcgen05 async ops issued by this thread complete +__device__ __forceinline__ void tc_commit(uint64_t* bar) { + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" ::"r"(smem_u32(bar)) + : "memory"); +} +// smem -> TMEM copy of one 512-byte scale-factor chunk (32 rows x 16 B), broadcast to the 4 lane quarters +__device__ __forceinline__ void utccp_32x128b_warpx4(uint32_t taddr, uint64_t sdesc) { + asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" ::"r"(taddr), "l"(sdesc) : "memory"); +} +// D(tmem) (+)= A(smem) * B(smem)^T with NVFP4 block-16 scales from TMEM +__device__ __forceinline__ void umma_nvf4(uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, + uint32_t tmem_sfa, uint32_t tmem_sfb, uint32_t accumulate) { + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %4, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%5], [%6], p;\n}\n" ::"r"( + tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(accumulate), "r"(tmem_sfa), "r"(tmem_sfb) + : "memory"); +} +// each warp reads its 32-lane quarter: 32 lanes x 16 consecutive 32-bit columns +__device__ __forceinline__ void tmem_ld_32x32b_x16(uint32_t taddr, uint32_t (&r)[16]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15}, [%16];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(taddr) + : "memory"); +} +// L2 prefetch of a contiguous global range (bytes multiple of 16, 16-byte aligned) +__device__ __forceinline__ void bulk_prefetch_l2(const void* gptr, uint32_t bytes) { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(gptr), "r"(bytes) : "memory"); +} +__device__ __forceinline__ void tmem_ld_wait() { asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); } +// each warp writes its 32-lane quarter: 32 lanes x 16 consecutive 32-bit columns +__device__ __forceinline__ void tmem_st_32x32b_x16(uint32_t taddr, const uint32_t (&r)[16]) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16};" ::"r"( + taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_wait() { asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); } + +// ---------------------------------------------------------------- descriptors +// SM100 shared-memory matrix descriptor (version 1). layout: 0 none, 2 SW128, 4 SW64, 6 SW32. +__device__ __forceinline__ uint64_t make_sdesc(uint32_t saddr, uint32_t lbo, uint32_t sbo, uint32_t layout) { + uint64_t d = 0; + d |= (uint64_t)((saddr >> 4) & 0x3FFF); + d |= (uint64_t)((lbo >> 4) & 0x3FFF) << 16; + d |= (uint64_t)((sbo >> 4) & 0x3FFF) << 32; + d |= (uint64_t)1 << 46; // version + d |= (uint64_t)(layout & 0x7) << 61; + return d; +} +// Block-scaled instruction descriptor: E2M1 x E2M1, UE4M3 scales, both K-major, F32 accumulate. +__host__ __device__ constexpr uint32_t make_idesc_nvf4(uint32_t M, uint32_t N) { + return (1u << 7) // a_format E2M1 + | (1u << 10) // b_format E2M1 + | ((N >> 3) << 17) // n_dim + | (0u << 23) // scale format UE4M3 + | ((M >> 4) << 24); // m_dim +} + +} // namespace qwen35 + +namespace sglang { + +namespace qwen_mkm_moe_m32 { + +using namespace qwen35; + + +constexpr int kH = 4096, kE = 512, kTopK = 10, kMaxTok = 32; +constexpr int kRows = 16; // A rows (tokens) per item +// per CTA: own / orphan expert, later orphans (at most 2 more at 32 tokens over >= 148 CTAs), a row-group split each +constexpr int kMaxItems = 6; +constexpr int kRounds = (kMaxTok + 6) / 7; // routing rounds of the 7 routing warps +constexpr int kI = 128; // intermediates per expert per rank (TP8) +constexpr int kN1 = 2 * kI; // W13 rows per expert (gate + up) +constexpr int kStageBytes = 36864; // 32 KB weights + 4 KB scale factors +constexpr int kStages = 4; +constexpr int kW13BoxRows = 128; // rows of one W13 TMA box; the host tensor map must use the same +constexpr int kFc1Stages = 16; // the expert's 256 W13 rows, 4 k-blocks each (full K = 4096) +constexpr int kFc2Stages = 8; // two tiles of 256 W2 rows x 64 B each +constexpr int kTotStages = kFc1Stages + kFc2Stages; +constexpr int kThreads = 256; +// TMEM columns. FC1: accumulator [0,256), x SFA ring [256,320) (16 slots x 4), SFB ring [384,512) (16 slots x 8). +// FC2: accumulator [0,256), h SFA [256,320) (16 atoms x 4), SFB ring [384,512) (16 slots x 8). +constexpr uint32_t kColSfaX = 256, kColSfa = 256, kColSfb = 384; +constexpr int kSfbSlots = 16, kSfaSlots = 16; +constexpr int kARow = 260; // FC1 accumulator row stride in floats: 16 B aligned, rows on distinct banks +#ifndef L2_EVICT_FIRST // stream expert weights with an L2 evict-first policy (they are read once per layer) +#define L2_EVICT_FIRST 1 +#endif +#ifndef TOPK_RANK +#define TOPK_RANK 0 +#endif +#ifndef SF_UTCCP_FC1 +#define SF_UTCCP_FC1 0 +#endif +#ifndef SPEC_STAGES // ring stages issued for the speculated expert before the routing verdict (rest: after it) +#define SPEC_STAGES 4 +#endif +#ifndef STAGE_STAMPS // DIAG: per-stage issue (slots 64 + st) and landing (16 + st, seen by warp 4) stamps +#define STAGE_STAMPS 0 +#endif +#ifndef DIAG_SPEC // DIAG: warp 1's speculative ranking, stamps 16 = logits loaded, 17 = keys, 18 = top-10 (stride 32) +#define DIAG_SPEC 0 +#endif +constexpr int kTraceStride = STAGE_STAMPS ? 128 : DIAG_SPEC ? 32 : 16; +#ifndef SMEM_SPACE // 1: shared-memory accesses through Smem compile to shared-space instructions +#define SMEM_SPACE 1 +#endif +#ifndef PF_SET // L2 prefetch of the stages that wait for a ring slot at: 1 = start, 4 = FC2, 8 = FC2 pass 1 +// (none: from 17 rows the CTAs' streams already keep HBM busy, so these prefetches only compete with the ring's loads) +#define PF_SET 0 +#endif +#ifndef PF_N1 // stages prefetched at the start trigger +#define PF_N1 4 +#endif +#ifndef PF_N2 // stages prefetched at the FC2 trigger +#define PF_N2 4 +#endif +#ifndef TAILV // FC2 drains: bit 0 = last pass on all 8 warps; bit 1 = TMEM loads double-buffered; bit 2 = coalesced stores +// (0: measured at 28 / 56 rows, the plain drains cost the least) +#define TAILV 0 +#endif +#define TAIL8 (TAILV & 1) +#define DRAIN_PIPE ((TAILV >> 1) & 1) +#define DRAIN_SMEM ((TAILV >> 2) & 1) +static_assert(DRAIN_SMEM == 0, "x's shared memory holds the next item's rows while the drains run"); +#ifndef EARLY_SF2 // FC2 stages whose scale factors are staged before the h epilogue (old order: 2) +#define EARLY_SF2 0 +#endif +#ifndef GATHER_BATCH // x row granules each gathering thread loads before storing them +#define GATHER_BATCH 8 +#endif +#ifndef SPEC_PF // stages of the speculated expert past the first ring fill prefetched to L2 while the routing runs +#define SPEC_PF 4 +#endif +#ifndef DIAG_ROUTE // DIAG: stamps 13 = warp 0's routing rounds done, 15 = their barrier, 14 = item 0's rows gathered +#define DIAG_ROUTE 0 +#endif +#ifndef FC2_PF // FC2 tiles (beyond the first kStages) prefetched to L2 ahead of the ring; 0 = off +#define FC2_PF 0 +#endif +#ifndef FC2_PF_LEAD // issue the FC2 prefetch when the producer reaches FC1 stage kFc1Stages - FC2_PF_LEAD +#define FC2_PF_LEAD 4 +#endif + +struct Params { + const uint16_t* logits; // [M, 512] bf16 router logits (router_gemm_quant output) + uint16_t* expert_weights; // [M, 10] bf16 out: DynBlock softmax-over-top-10 weights (the last CTA writes) + int32_t* topk_ids; // [M, 10] out (the last CTA writes) + const uint8_t* x_sf; // [M, 256] linear e4m3 + const uint8_t* w13_sf; // [E, 2 row blocks x 64 k blocks x 512 B] (128x4 interleaved) + const uint8_t* w2_sf; // [E, 32 row blocks x 2 k blocks x 512 B] + const float* g1_alpha; // [E] gate dequant + const float* g1_scale_c; // [E] up dequant * FC2 input quant scale + const float* g2_alpha; // [E] + __nv_bfloat16* out; // [M * 10, 4096] + long long* trace; // optional [grid, 16] %globaltimer stamps + int num_tokens; + const uint8_t* x_fp4; // [M, 2048] NVFP4 activations (each item gathers its rows) +}; + +struct __align__(1024) Smem { + uint8_t ring[kStages][kStageBytes]; + alignas(1024) uint8_t xa[32 * 1024]; // x: per 128-B K chunk c, 2 KB = SW128 atoms of item rows 0..7 | 8..15 + alignas(1024) uint8_t ha[16 * 1024]; // h: 8 copies of (rows 0..7 | rows 8..15) x 128 B (K = 128: 64 B), SW128 + float acc1[16 * kARow]; // FC1 accumulator rows 0..15 (item rows) x 256 columns + uint8_t xsf[16][256]; // x scale bytes of the item's rows, full K + uint32_t hsf[16][2]; // h scale words: row r, k-block kb (4 e4m3 bytes) + uint64_t full[kStages], empty[kStages], sf_full[kStages]; + uint64_t x_full, fc1_done, acc_ready, h_ready, fc2_pass[2], fc2_free, acc_drained, route_done; + uint32_t tmem; + int nitems, spec_ok; + int item_e[kMaxItems]; + uint32_t item_mask[kMaxItems]; // tokens of each item (bit t) + int nrows[2], rtok[2][kRows], rkslot[2][kRows]; // rows of the current / next item (by item parity) + int first[kE]; // first slot of each expert (atomicMin) + uint32_t fbits[kMaxTok]; // token t: its top-k slots that are their expert's first (bit k) + alignas(16) uint32_t cand[8][32]; // read as uint4 by the rank-based top-k + int32_t ids[kMaxTok][kTopK]; +}; + +__device__ __forceinline__ uint64_t gtime() { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)::"memory"); + return t; +} +__device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} +__device__ __forceinline__ void cluster_arrive() { + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); +} +// no release: the phase only orders DSMEM traffic that mbarriers already completed, so pending global stores +// (stamps, routing outputs, output rows) need not drain before the arrive +__device__ __forceinline__ void cluster_arrive_relaxed() { + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} +__device__ __forceinline__ void cluster_wait() { + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); +} +__device__ __forceinline__ uint32_t mapa(uint32_t saddr, uint32_t rank) { + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(r) : "r"(saddr), "r"(rank)); + return r; +} +// smem (this CTA) -> smem (peer CTA in the cluster), completes tx bytes on the peer's mbarrier +__device__ __forceinline__ void bulk_s2s(uint32_t dst_cluster, uint32_t src, uint32_t bytes, uint32_t bar_cluster) { + asm volatile("cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + dst_cluster), + "r"(src), "r"(bytes), "r"(bar_cluster) + : "memory"); +} +// registers -> the peer CTA's shared memory; completes the bytes on the peer's mbarrier +__device__ __forceinline__ void st_async_v4(uint32_t dst, const uint4& v, uint32_t bar) { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"(dst), + "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "r"(bar) + : "memory"); +} +__device__ __forceinline__ void st_async_b32(uint32_t dst, uint32_t v, uint32_t bar) { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(dst), "r"(v), "r"(bar) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x4(uint32_t taddr, const uint4& v) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x4.b32 [%0], {%1, %2, %3, %4};" ::"r"(taddr), "r"(v.x), "r"(v.y), + "r"(v.z), "r"(v.w) + : "memory"); +} +// bar.sync counts arrivals per warp execution: reconverge first, so a warp whose lane 0 just ran a stamp or an +// mbarrier arrive cannot arrive twice and release the barrier early +__device__ __forceinline__ void named_bar(int id, int n) { + __syncwarp(); + asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(n) : "memory"); +} +__device__ __forceinline__ void tmem_ld_x32(uint32_t taddr, uint32_t (&r)[32]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,%19," + "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), "=r"(r[16]), + "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), "=r"(r[24]), + "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr) + : "memory"); +} +__device__ __forceinline__ void tmem_ld_x64(uint32_t taddr, uint32_t (&r)[64]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x64.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,%19," + "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31,%32,%33,%34,%35,%36,%37,%38,%39,%40,%41,%42,%43,%44,%45,%46,%47," + "%48,%49,%50,%51,%52,%53,%54,%55,%56,%57,%58,%59,%60,%61,%62,%63}, [%64];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), "=r"(r[16]), + "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), "=r"(r[24]), + "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]), "=r"(r[32]), + "=r"(r[33]), "=r"(r[34]), "=r"(r[35]), "=r"(r[36]), "=r"(r[37]), "=r"(r[38]), "=r"(r[39]), "=r"(r[40]), + "=r"(r[41]), "=r"(r[42]), "=r"(r[43]), "=r"(r[44]), "=r"(r[45]), "=r"(r[46]), "=r"(r[47]), "=r"(r[48]), + "=r"(r[49]), "=r"(r[50]), "=r"(r[51]), "=r"(r[52]), "=r"(r[53]), "=r"(r[54]), "=r"(r[55]), "=r"(r[56]), + "=r"(r[57]), "=r"(r[58]), "=r"(r[59]), "=r"(r[60]), "=r"(r[61]), "=r"(r[62]), "=r"(r[63]) + : "r"(taddr) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x8(uint32_t taddr, const uint32_t (&r)[8]) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x8.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" ::"r"(taddr), "r"(r[0]), + "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x16(uint32_t taddr, const uint32_t (&r)[16]) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x16.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16};" :: + "r"(taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x32(uint32_t taddr, const uint32_t (&r)[32]) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x32.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18," + "%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31,%32};" ::"r"(taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]), "r"(r[16]), "r"(r[17]), "r"(r[18]), + "r"(r[19]), "r"(r[20]), "r"(r[21]), "r"(r[22]), "r"(r[23]), "r"(r[24]), "r"(r[25]), "r"(r[26]), "r"(r[27]), + "r"(r[28]), "r"(r[29]), "r"(r[30]), "r"(r[31]) + : "memory"); +} +__device__ __forceinline__ float ex2_approx(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float mul_ftz(float a, float b) { + float y; + asm("mul.ftz.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ float rcp_approx(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ uint32_t cvt_e4m3(float a) { + uint32_t out; + asm("{\n .reg .b16 pair;\n .reg .f32 zero;\n mov.f32 zero, 0f00000000;\n" + " cvt.rn.satfinite.e4m3x2.f32 pair, zero, %1;\n cvt.u32.u16 %0, pair;\n}" + : "=r"(out) + : "f"(a)); + return out & 0xFFu; +} +__device__ __forceinline__ float e4m3_to_f32(uint32_t v) { + float out; + asm("{\n .reg .b16 pair;\n .reg .b32 h2;\n .reg .b16 lo, hi;\n cvt.u16.u32 pair, %1;\n" + " cvt.rn.f16x2.e4m3x2 h2, pair;\n mov.b32 {lo, hi}, h2;\n cvt.f32.f16 %0, lo;\n}" + : "=f"(out) + : "r"(v)); + return out; +} +__device__ __forceinline__ uint32_t cvt_e2m1x8(const float* v) { + uint32_t out; + asm("{\n .reg .b8 b0, b1, b2, b3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\n cvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\n cvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" + " mov.b32 %0, {b0, b1, b2, b3};\n}" + : "=r"(out) + : "f"(v[0]), "f"(v[1]), "f"(v[2]), "f"(v[3]), "f"(v[4]), "f"(v[5]), "f"(v[6]), "f"(v[7])); + return out; +} +// ---- router top-k, bit-exact with FlashInfer's single-block routing (routingCustom, RenormalizeNaive on bf16): +// key = bf16 order (cub float twiddle) above 65535 - expert, i.e. the higher score first, the smaller expert on ties. +__device__ __forceinline__ uint32_t topk_key(uint32_t bits, uint32_t expert) { + const uint32_t twiddled = bits ^ ((bits & 0x8000u) ? 0xFFFFu : 0x8000u); + return (twiddled << 16) | (65535u - expert); +} +__device__ __forceinline__ void topk_unkey(uint32_t key, uint32_t& bits, uint32_t& expert) { + expert = 65535u - (key & 0xFFFFu); + const uint32_t twiddled = key >> 16; + bits = twiddled ^ ((twiddled & 0x8000u) ? 0x8000u : 0xFFFFu); +} +__device__ __forceinline__ uint32_t max16(const uint32_t (&k)[16]) { + const uint32_t a = max(max(max(k[0], k[1]), max(k[2], k[3])), max(max(k[4], k[5]), max(k[6], k[7]))); + const uint32_t b = max(max(max(k[8], k[9]), max(k[10], k[11])), max(max(k[12], k[13]), max(k[14], k[15]))); + return max(a, b); +} +__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_sum_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.add.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_topk_of_lanes(uint32_t v, uint32_t lane) { + uint32_t mine = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(v); + v = v == top ? 0u : v; + mine = lane == k ? top : mine; + } + return mine; +} +__device__ __forceinline__ uint32_t count_greater(const uint32_t* s_cand, uint32_t v) { + uint32_t r = 0; +#pragma unroll + for (int j = 0; j < 32; j += 4) { + const uint4 q = *reinterpret_cast(s_cand + j); + r += (q.x > v) + (q.y > v) + (q.z > v) + (q.w > v); + } + return r; +} +// Top-10 keys of 512 (16 per lane), lane k < 10 holding the k-th largest (keys unique and nonzero). Ranks come from +// all-pairs compares over shared memory rather than chains of warp reductions. +__device__ __forceinline__ uint32_t warp_top10(uint32_t (&keys)[16], uint32_t lane, uint32_t* s_cand) { + // bound = the 10th largest lane maximum; every top-10 key is >= bound + const uint32_t m = max16(keys); + s_cand[lane] = m; + __syncwarp(); + const uint32_t r = count_greater(s_cand, m); + const uint32_t bound = __shfl_sync(0xffffffffu, m, __ffs(__ballot_sync(0xffffffffu, r == kTopK - 1)) - 1); + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) count += keys[i] >= bound; + // exclusive prefix of count over the lanes, from the ballots of its bits (count <= 16) + const uint32_t lt = (1u << lane) - 1u; + uint32_t offset = 0, total = 0; +#pragma unroll + for (int b = 0; b < 5; ++b) { + const uint32_t bb = __ballot_sync(0xffffffffu, (count >> b) & 1u); + offset += (uint32_t)__popc(bb & lt) << b; + total += (uint32_t)__popc(bb) << b; + } + if (total <= 32) { + __syncwarp(); // every lane has read the lane maxima +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) { + if (keys[i] >= bound) s_cand[offset++] = keys[i]; + } + if (lane >= total) s_cand[lane] = 0u; + __syncwarp(); + const uint32_t cand = s_cand[lane]; + const uint32_t rk = count_greater(s_cand, cand); + __syncwarp(); + if (lane < total && rk < kTopK) s_cand[rk] = cand; + __syncwarp(); + return lane < kTopK ? s_cand[lane] : 0u; + } + uint32_t mine = 0; +#pragma unroll 1 + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(max16(keys)); +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) keys[i] = keys[i] == top ? 0u : keys[i]; + mine = lane == k ? top : mine; + } + return mine; +} +// Top-10 keys of 512 (16 per lane), lane k holding the k-th largest (keys unique and nonzero). +__device__ __forceinline__ uint32_t warp_top10_redux(uint32_t (&keys)[16], uint32_t lane, uint32_t* s_cand) { + uint32_t v = max16(keys); + uint32_t bound = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + bound = warp_max_u32(v); + v = v == bound ? 0u : v; + } + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) count += keys[i] >= bound; + const uint32_t total = warp_sum_u32(count); + if (total <= 32) { + uint32_t offset = count; +#pragma unroll + for (uint32_t d = 1; d < 32; d <<= 1) { + const uint32_t up = __shfl_up_sync(0xffffffffu, offset, d); + offset += lane >= d ? up : 0u; + } + offset -= count; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) { + if (keys[i] >= bound) s_cand[offset++] = keys[i]; + } + __syncwarp(); + const uint32_t cand = lane < total ? s_cand[lane] : 0u; + __syncwarp(); + return warp_topk_of_lanes(cand, lane); + } + uint32_t mine = 0; +#pragma unroll 1 + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(max16(keys)); +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) keys[i] = keys[i] == top ? 0u : keys[i]; + mine = lane == k ? top : mine; + } + return mine; +} +// The top-10 keys of one logits row (16 bf16 a lane in v0, v1), lane k < 10 holding the k-th largest +__device__ __forceinline__ uint32_t row_top10(const uint4 v0, const uint4 v1, uint32_t lane, uint32_t* s_cand) { + uint32_t keys[16]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t w = i < 4 ? (&v0.x)[i] : (&v1.x)[i - 4]; + keys[2 * i] = topk_key(w & 0xFFFFu, lane * 16 + 2 * i); + keys[2 * i + 1] = topk_key(w >> 16, lane * 16 + 2 * i + 1); + } + return TOPK_RANK ? warp_top10(keys, lane, s_cand) : warp_top10_redux(keys, lane, s_cand); +} +// FlashInfer calcSoftmax (bf16 scores) as built with -use_fast_math: expf -> ex2.approx.ftz of x * log2(e), +// division -> div.full.ftz; reductions are cooperative_groups::reduce like the reference. +__device__ __forceinline__ float fast_expf(float x) { + float y; + asm("{\n .reg .f32 t;\n mul.ftz.f32 t, %1, 0f3FB8AA3B;\n ex2.approx.ftz.f32 %0, t;\n}" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float fast_div(float a, float b) { + float y; + asm("div.full.ftz.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ uint16_t softmax_topk(uint16_t score_bits, uint32_t lane) { + namespace cg = cooperative_groups; + const auto warp = cg::tiled_partition<32>(cg::this_thread_block()); + const float si = __uint_as_float((uint32_t)score_bits << 16); + float max_score = -INFINITY; + if (lane < kTopK) max_score = si >= max_score ? si : max_score; + max_score = cg::reduce(warp, max_score, cg::greater()); + float sum_score = 0.f, new_score = 0.f; + if (lane < kTopK) { + new_score = fast_expf(__fsub_rn(si, max_score)); + sum_score = __fadd_rn(sum_score, new_score); + } + sum_score = cg::reduce(warp, sum_score, cg::plus()); + uint16_t out = score_bits; + if (lane < kTopK) { + const __nv_bfloat16 b = __float2bfloat16_rn(fast_div(new_score, sum_score)); + out = *reinterpret_cast(&b); + } + return out; +} + +__device__ __forceinline__ uint64_t policy_evict_first() { + uint64_t pol; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(pol)); + return pol; +} +__device__ __forceinline__ void tma_load_2d_w(void* dst, const void* tmap, uint64_t* bar, int c0, int c1, uint64_t pol) { + if (L2_EVICT_FIRST) + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%3, %4}], [%2], %5;" ::"r"(smem_u32(dst)), + "l"(tmap), "r"(smem_u32(bar)), "r"(c0), "r"(c1), "l"(pol) + : "memory"); + else + tma_load_2d(dst, tmap, bar, c0, c1); +} +__device__ __forceinline__ void bulk_load_w(void* dst, const void* src, uint32_t bytes, uint64_t* bar, uint64_t pol) { + if (L2_EVICT_FIRST) + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)), "l"(pol) + : "memory"); + else + bulk_load(dst, src, bytes, bar); +} + +__device__ __forceinline__ void stamp(const Params& p, int slot) { + if (p.trace) p.trace[blockIdx.x * kTraceStride + slot] = (long long)gtime(); +} +// stamp after `dep` is available (the branch needs its value; keys are never 0) +__device__ __forceinline__ void stamp_dep(const Params& p, int slot, uint32_t dep) { + if (dep == 0u) __trap(); + stamp(p, slot); +} +__device__ __forceinline__ void tma_prefetch_2d(const void* tmap, int c0, int c1) { + asm volatile("cp.async.bulk.prefetch.tensor.2d.L2.global.tile [%0, {%1, %2}];" ::"l"(tmap), "r"(c0), "r"(c1) : "memory"); +} +// Physical W13 column (accumulator column) holding reordered row q: block-32 shuffle of FlashInfer's +// get_shuffle_matrix_a_row_indices (epilogue_tile_m = 128); reordered row 2j is up_j, 2j + 1 is gate_j. +__device__ __forceinline__ int w13_col(int q) { return (q & ~31) + ((q & 31) % 4) * 8 + (q & 31) / 4; } +__host__ __device__ constexpr int w13_logical(int p) { + return ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) % 2 ? ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) / 2 + : 128 + ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) / 2; +} + + +// FC2 drain of one pass: TMEM lane quarter q (tiles 2 q, 2 q + 1 of the pass), 256 accumulator columns +__device__ __forceinline__ void reg_fence64(uint32_t (&r)[64]) { // keep uses of r after the tcgen05.wait::ld +#pragma unroll + for (int i = 0; i < 64; ++i) asm volatile("" : "+r"(r[i])); +} + +// FC2 accumulator columns [64 C0, 64 C1) of this warp's lane quarter -> bf16 gemm2_out rows (row r of a tile = the +// r-th row of the item, row info buffer rb) +template +__device__ __forceinline__ void fc2_drain16(const Params& p, Smem& s, uint32_t tm, int q, int lane, int pass, float a2, + int rb) { + const int row = q * 32 + lane, tile = pass * 8 + row / 16, r = row % 16; + const bool valid = r < s.nrows[rb]; + __nv_bfloat16* orow = valid ? p.out + (size_t)(s.rtok[rb][r] * kTopK + s.rkslot[rb][r]) * kH + tile * 256 : nullptr; + const uint32_t base = tm + ((uint32_t)(q * 32) << 16); + auto put = [&](const uint32_t (&r)[64], int c) { + if (valid) { +#pragma unroll + for (int b32 = 0; b32 < 2; ++b32) { + uint32_t pk[16]; +#pragma unroll + for (int l = 0; l < 32; l += 2) { + const float v0 = __uint_as_float(r[b32 * 32 + (l % 4) * 8 + l / 4]) * a2; + const float v1 = __uint_as_float(r[b32 * 32 + ((l + 1) % 4) * 8 + (l + 1) / 4]) * a2; + const __nv_bfloat162 b2 = __floats2bfloat162_rn(v0, v1); + pk[l / 2] = *reinterpret_cast(&b2); + } +#pragma unroll + for (int i = 0; i < 4; ++i) + *reinterpret_cast(orow + c * 64 + b32 * 32 + i * 8) = + make_uint4(pk[4 * i], pk[4 * i + 1], pk[4 * i + 2], pk[4 * i + 3]); + } + } + }; + if (DRAIN_PIPE) { // the next chunk's TMEM load is in flight while this chunk converts and stores + uint32_t ra[64], rb2[64]; + tmem_ld_x64(base + C0 * 64, ra); + tmem_ld_wait(); + reg_fence64(ra); +#pragma unroll + for (int c = C0; c < C1; c += 2) { + if (c + 1 < C1) tmem_ld_x64(base + (uint32_t)((c + 1) * 64), rb2); + put(ra, c); + if (c + 1 < C1) { + tmem_ld_wait(); + reg_fence64(rb2); + if (c + 2 < C1) tmem_ld_x64(base + (uint32_t)((c + 2) * 64), ra); + put(rb2, c + 1); + if (c + 2 < C1) { tmem_ld_wait(); reg_fence64(ra); } + } + } + } else { +#pragma unroll 1 + for (int c = C0; c < C1; ++c) { + uint32_t r[64]; + tmem_ld_x64(base + (uint32_t)(c * 64), r); + tmem_ld_wait(); + put(r, c); + } + } +} + +extern "C" __global__ void __launch_bounds__(kThreads, 1) +q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant__ CUtensorMap tm2, const Params p) { + extern __shared__ __align__(1024) uint8_t raw[]; +#if SMEM_SPACE // pointer arithmetic on the __shared__ array keeps the shared address space (STS/LDS, not generic) + Smem& s = *reinterpret_cast(raw + ((1024u - (smem_u32(raw) & 1023u)) & 1023u)); +#else + Smem& s = *reinterpret_cast((reinterpret_cast(raw) + 1023) & ~uintptr_t(1023)); +#endif + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + const int slot = blockIdx.x; // the (token, top-k) slot this CTA speculates on; gridDim.x <= M * 10 + const int M = p.num_tokens; + for (int i = threadIdx.x; i < kE; i += kThreads) s.first[i] = 0x7FFFFFFF; + if (threadIdx.x == 0) { + stamp(p, 0); + for (int i = 0; i < kStages; ++i) { + mbar_init(&s.full[i], 1); + mbar_init(&s.empty[i], 1); + mbar_init(&s.sf_full[i], 4); + } + // every barrier below completes once per item, so item i waits parity i & 1 + mbar_init(&s.x_full, 1); + mbar_init(&s.fc1_done, 1); + mbar_init(&s.acc_ready, 2); // warps 0 and 4 drained their columns into acc1 + mbar_init(&s.h_ready, 4); + mbar_init(&s.fc2_pass[0], 1); + mbar_init(&s.fc2_pass[1], 1); + mbar_init(&s.fc2_free, 4); + mbar_init(&s.acc_drained, 4); // warps 4..7 read the last FC2 pass: the next item's FC1 may overwrite columns 0..255 + mbar_init(&s.route_done, 1); + fence_mbar_init(); + } + if (warp == 2) { tmem_alloc(&s.tmem, 512); tmem_relinquish(); tc_fence_before(); } + if (warp == 1 && lane == 0) { prefetch_tmap(&tm13); prefetch_tmap(&tm2); } + __syncthreads(); + tc_fence_after(); + const uint32_t tm = s.tmem; + const int my_t = slot / kTopK, my_k = slot % kTopK; + // ring stage g (global over the CTA's items) of item-relative stage st of expert e + auto issue_stage = [&](int g, int st, int e, uint64_t pol) { + const int b = g % kStages; + mbar_expect_tx(&s.full[b], kStageBytes); + uint8_t* dst = s.ring[b]; + if (st < kFc1Stages) { + // the expert's rows as two boxes of 128 rows, K bytes [128 st, 128 st + 128) + tma_load_2d_w(dst, &tm13, &s.full[b], st * 128, e * kN1, pol); + tma_load_2d_w(dst + 16384, &tm13, &s.full[b], st * 128, e * kN1 + 128, pol); + // scale atoms of row blocks 0 and 1, k-blocks [4 st, 4 st + 4) + const uint8_t* sf = p.w13_sf + (size_t)e * (kN1 * kH / 16); + bulk_load_w(dst + 32768, sf + (size_t)(st * 4) * 512, 2048, &s.full[b], pol); + bulk_load_w(dst + 32768 + 2048, sf + (size_t)(64 + st * 4) * 512, 2048, &s.full[b], pol); + } else { + // W2 tiles 2 f and 2 f + 1 (256 rows x 64 B each) and their scale atoms (2 row blocks x 2 k-blocks) + const int f = st - kFc1Stages; +#pragma unroll + for (int tl = 0; tl < 2; ++tl) { + const int row0 = (2 * f + tl) * 256; + tma_load_2d_w(dst + tl * 16384, &tm2, &s.full[b], 0, e * kH + row0, pol); + bulk_load_w(dst + 32768 + tl * 2048, p.w2_sf + (size_t)e * (kH * 8) + (size_t)(row0 / 128) * 2 * 512, 2048, + &s.full[b], pol); + } + } + }; + auto prefetch_stage = [&](int st, int e) { // the same bytes as issue_stage, to L2 only + if (st < kFc1Stages) { + tma_prefetch_2d(&tm13, st * 128, e * kN1); + tma_prefetch_2d(&tm13, st * 128, e * kN1 + 128); + const uint8_t* sf = p.w13_sf + (size_t)e * (kN1 * kH / 16); + bulk_prefetch_l2(sf + (size_t)(st * 4) * 512, 2048); + bulk_prefetch_l2(sf + (size_t)(64 + st * 4) * 512, 2048); + } else { + const int row0 = (st - kFc1Stages) * 512; + tma_prefetch_2d(&tm2, 0, e * kH + row0); + tma_prefetch_2d(&tm2, 0, e * kH + row0 + 256); + bulk_prefetch_l2(p.w2_sf + (size_t)e * (kH * 8) + (size_t)(row0 / 128) * 2 * 512, 4096); + } + }; + // Rows of an item (the tokens of mask in increasing order, at most 16) into row buffer rb, by one warp + auto set_rows = [&](int rb, uint32_t mask, int ei) { + if (lane < M && ((mask >> lane) & 1u)) { + int ks = 0; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + if (s.ids[lane][k] == ei) ks = k; + const int r = __popc(mask & ((1u << lane) - 1u)); + s.rtok[rb][r] = lane; + s.rkslot[rb][r] = ks; + } + if (lane == 0) s.nrows[rb] = __popc(mask); + }; + // their x rows and scale bytes into xa / xsf, by threads tid = 0 .. nt - 1 + auto copy_rows = [&](int rb, int tid, int nt) { + const int n1 = s.nrows[rb] * 128, n = s.nrows[rb] * 144; // 16-B granules: x rows, then their scale bytes + // GATHER_BATCH loads a thread ahead of their stores: the global latency is paid once a batch + for (int i0 = tid; i0 < n; i0 += GATHER_BATCH * nt) { + uint4 v[GATHER_BATCH]; +#pragma unroll + for (int j = 0; j < GATHER_BATCH; ++j) { + const int idx = i0 + j * nt; + if (idx < n1) { // row r, K chunk c (128 B), 16-B granule q of the chunk + const int r = idx >> 7, c = (idx >> 3) & 15, q = idx & 7; + v[j] = *reinterpret_cast(p.x_fp4 + (size_t)s.rtok[rb][r] * (kH / 2) + c * 128 + q * 16); + } else if (idx < n) { + const int r = (idx - n1) >> 4, q = (idx - n1) & 15; + v[j] = *reinterpret_cast(p.x_sf + (size_t)s.rtok[rb][r] * (kH / 16) + q * 16); + } + } +#pragma unroll + for (int j = 0; j < GATHER_BATCH; ++j) { + const int idx = i0 + j * nt; + if (idx < n1) { + const int r = idx >> 7, c = (idx >> 3) & 15, q = idx & 7; + *reinterpret_cast(s.xa + c * 2048 + (r / 8) * 1024 + (r % 8) * 128 + ((q ^ (r % 8)) * 16)) = v[j]; + } else if (idx < n) { + const int r = (idx - n1) >> 4, q = (idx - n1) & 15; + *reinterpret_cast(&s.xsf[r][q * 16]) = v[j]; + } + } + } + fence_proxy_async_smem(); // the tensor core reads xa through the async proxy + }; + // routing: warps 0, 2..7 take tokens 0..31 (DynBlock-exact top-10 + softmax), up to kRounds rounds; warp 1 recomputes + // token my_t's top-10 and starts this CTA's weight stream at once (a wrong guess only wastes its first ring fill) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (threadIdx.x == 32) stamp(p, 10); + auto top10v = [&](const uint4 v0, const uint4 v1) -> uint32_t { + return row_top10(v0, v1, (uint32_t)lane, s.cand[warp]); + }; + auto top10 = [&](int rtok) -> uint32_t { + const uint4* src = reinterpret_cast(p.logits + (size_t)rtok * kE + lane * 16); + return top10v(src[0], src[1]); + }; + if (warp == 1) { + // speculate on this CTA's own slot: its expert's weight stream starts before the routing verdict + uint32_t key; + if (DIAG_SPEC) { + const uint4* src = reinterpret_cast(p.logits + (size_t)my_t * kE + lane * 16); + const uint4 v0 = src[0], v1 = src[1]; + if (lane == 0) stamp_dep(p, 16, (v0.x | v0.y | v0.z | v0.w | v1.x | v1.y | v1.z | v1.w) | 1u); + key = top10v(v0, v1); + if (lane == 0) stamp_dep(p, 18, key | 1u); + } else { + key = top10(my_t); + } + uint32_t bits, espec; + topk_unkey(__shfl_sync(0xffffffffu, key, my_k), bits, espec); + if (lane == 0) { + stamp(p, 9); + const uint64_t pol = policy_evict_first(); + for (int st = 0; st < SPEC_STAGES; ++st) { + issue_stage(st, st, (int)espec, pol); + if (STAGE_STAMPS) stamp(p, 64 + st); + } + // HBM is otherwise idle until the routing is done: the next stages go to L2 now (another CTA streams the + // expert when this slot is not its first) + for (int st = SPEC_STAGES; st < SPEC_STAGES + SPEC_PF && st < kTotStages; ++st) prefetch_stage(st, (int)espec); + stamp(p, 2); + } + mbar_wait(&s.route_done, 0); + } else { + const int base = warp == 0 ? 0 : warp - 1; + // every round's logits are requested before the first ranking, so later rounds do not pay the load latency again + uint4 lv[kRounds][2]; +#pragma unroll + for (int r = 0; r < kRounds; ++r) { + const int rtok = base + 7 * r; + if (rtok < M) { + const uint4* src = reinterpret_cast(p.logits + (size_t)rtok * kE + lane * 16); + lv[r][0] = src[0]; + lv[r][1] = src[1]; + } + } + // the last CTA writes the routing outputs (the further items start at CTA 0); the others need the ids only + const bool writer = blockIdx.x == gridDim.x - 1; +#pragma unroll + for (int r = 0; r < kRounds; ++r) { + const int rtok = base + 7 * r; + if (rtok >= M) break; + const uint32_t key = top10v(lv[r][0], lv[r][1]); + uint32_t bits, expert; + topk_unkey(key, bits, expert); + if (lane < kTopK) { + s.ids[rtok][lane] = (int32_t)expert; + atomicMin(&s.first[expert], rtok * kTopK + lane); + } + if (writer) { + const uint16_t wbits = softmax_topk((uint16_t)bits, (uint32_t)lane); + if (lane < kTopK) { + p.expert_weights[rtok * kTopK + lane] = wbits; + p.topk_ids[rtok * kTopK + lane] = (int32_t)expert; + } + } + __syncwarp(); + } + if (DIAG_ROUTE && threadIdx.x == 0) stamp(p, 13); + named_bar(4, 224); + if (DIAG_ROUTE && threadIdx.x == 0) stamp(p, 15); + // each routing warp marks its tokens' slots that are their expert's first (s.first is final after the barrier) +#pragma unroll + for (int r = 0; r < kRounds; ++r) { + const int rtok = base + 7 * r; + if (rtok >= M) break; + const uint32_t f = __ballot_sync(0xffffffffu, lane < kTopK && s.first[s.ids[rtok][lane]] == rtok * kTopK + lane); + if (lane == 0) s.fbits[rtok] = f; + } + named_bar(4, 224); + if (warp == 0) { + // Work items. Distinct experts in order of their first slot; slot c's first-occurrence expert belongs to CTA c + // (whose speculative stream already carries it). Experts first seen at slots >= G ("orphans") go to the CTAs + // whose own slot repeats an earlier expert, in order; orphans beyond those CTAs become further items of CTAs + // 0, 1, ... Every CTA derives the same assignment from the same routing. + const int G = (int)gridDim.x; + const uint32_t fb = lane < M ? s.fbits[lane] : 0u; + int incl = __popc(fb); // first-occurrence slots up to token `lane`'s, inclusive +#pragma unroll + for (int d = 1; d < 32; d <<= 1) { + const int up = __shfl_up_sync(0xffffffffu, incl, d); + incl += lane >= d ? up : 0; + } + const int excl = incl - __popc(fb); + const int U = __shfl_sync(0xffffffffu, incl, M - 1); // distinct experts + auto before = [&](int s2) -> int { // first-occurrence slots before slot s2 (s2 <= M * 10) + const int t = s2 / kTopK, tt = t < M ? t : M - 1; + const int ex = __shfl_sync(0xffffffffu, excl, tt); + const uint32_t ft = __shfl_sync(0xffffffffu, fb, tt); + return t < M ? ex + __popc(ft & ((1u << (s2 % kTopK)) - 1u)) : U; + }; + auto nth_first = [&](int n) -> int { // the slot of the n-th (0-based) first occurrence, n < U + const int t = __popc(__ballot_sync(0xffffffffu, lane < M && incl <= n)); + uint32_t ft = __shfl_sync(0xffffffffu, fb, t); + for (int i = __shfl_sync(0xffffffffu, excl, t); i < n; ++i) ft &= ft - 1u; + return t * kTopK + __ffs(ft) - 1; + }; + const int nf_c = before(slot), nf_g = before(G); + const int xc = s.ids[my_t][my_k]; + const bool owner = s.first[xc] == slot; + const int D = G - nf_g, dc = slot - nf_c, O = U - nf_g; // duplicate CTAs, this one's rank among them; orphans + // the whole warp walks the items (nth_first is warp-collective); lane 0 writes them + int ni = 0; + if (owner || dc < O) { + const int s2 = owner ? slot : nth_first(nf_g + dc); + if (lane == 0) s.item_e[ni] = s.ids[s2 / kTopK][s2 % kTopK]; + ++ni; + } + for (int o = D + slot; o < O; o += G) { + if (ni >= kMaxItems) __trap(); + const int s2 = nth_first(nf_g + o); + if (lane == 0) s.item_e[ni] = s.ids[s2 / kTopK][s2 % kTopK]; + ++ni; + } + if (lane == 0) + for (int i = ni; i < kMaxItems; ++i) s.item_e[i] = -1; + __syncwarp(); + // token masks; an expert with more than 16 tokens is split into row groups (the second one appended) + int n = 0; + for (int i = 0; i < kMaxItems && s.item_e[i] >= 0; ++i) n = i + 1; + int nall = n; + for (int i = 0; i < n; ++i) { + const int ei = s.item_e[i]; + bool has = false; + if (lane < M) { +#pragma unroll + for (int k = 0; k < kTopK; ++k) has |= s.ids[lane][k] == ei; + } + const uint32_t mask = __ballot_sync(0xffffffffu, has); + uint32_t lo = mask; // the 16 lowest tokens; the rest are a further item + for (int k = __popc(mask) - kRows; k > 0; --k) lo &= ~(0x80000000u >> __clz(lo)); + if (lane == 0) s.item_mask[i] = lo; + if (mask != lo) { + if (nall >= kMaxItems) __trap(); + if (lane == 0) { s.item_e[nall] = ei; s.item_mask[nall] = mask & ~lo; } + ++nall; + } + __syncwarp(); + } + if (lane == 0) { + s.nitems = nall; + s.spec_ok = owner ? 1 : 0; + stamp(p, 1); + mbar_arrive(&s.route_done); + } + } + named_bar(4, 224); + } + const int nitems = s.nitems; + if (nitems == 0) { // a duplicate slot with no work: let the speculative loads land, then leave + if (warp == 1) { + if (lane == 0) + for (int b = 0; b < SPEC_STAGES; ++b) mbar_wait(&s.full[b], 0); + __syncwarp(); + } + named_bar(7, kThreads); + if (warp == 2) tmem_dealloc(tm, 512); + return; + } + tc_fence_after(); + const bool spec_ok = s.spec_ok != 0; + // ring stage of item 0's stage 0: the speculative fill (stages 0 .. SPEC_STAGES - 1) is item 0's own only if spec_ok + const int base0 = spec_ok ? 0 : SPEC_STAGES; + float ag = 0.f, sc = 0.f, a2 = 0.f; + auto load_scales = [&](int it) { + const int e = s.item_e[it]; + ag = p.g1_alpha[e]; + sc = p.g1_scale_c[e]; + a2 = p.g2_alpha[e]; + }; + // FC1 drain of quarter-0 accumulator columns [64 c0, 64 c0 + 128) (TMEM lanes 0..15 hold the item rows) into acc1 + auto drain_fc1 = [&](int c0) { +#pragma unroll 1 + for (int c = c0; c < c0 + 2; ++c) { + uint32_t r[64]; + tmem_ld_x64(tm + (uint32_t)(c * 64), r); + tmem_ld_wait(); + if (lane < 16) { +#pragma unroll + for (int i = 0; i < 16; ++i) + *reinterpret_cast(s.acc1 + lane * kARow + c * 64 + 4 * i) = + make_uint4(r[4 * i], r[4 * i + 1], r[4 * i + 2], r[4 * i + 3]); + } + } + tc_fence_before(); + __syncwarp(); + }; + // SwiGLU + NVFP4 epilogue (h = 0: the one FC1 pass), tasks (row, 16-wide block, half) = 16 x 8 x 2 = 256, + // task = et + k * nt. + // Accumulator columns of each 32-block: [0, 8) up 2m, [8, 16) gate 2m, [16, 24) up 2m + 1, [24, 32) gate 2m + 1 + auto epi_tasks = [&](int h, int et, int nt) { +#pragma unroll 1 + for (int task = et; task < 256; task += nt) { + const int row = task / 16, blk = (task / 2) % 8, hf = task % 2; + const float4* ar = reinterpret_cast(s.acc1 + row * kARow + 32 * blk + 4 * hf); + const float4 u_even = ar[0], g_even = ar[2], u_odd = ar[4], g_odd = ar[6]; + const float ga[8] = {g_even.x, g_odd.x, g_even.y, g_odd.y, g_even.z, g_odd.z, g_even.w, g_odd.w}; + const float ua[8] = {u_even.x, u_odd.x, u_even.y, u_odd.y, u_even.z, u_odd.z, u_even.w, u_odd.w}; + // trtllm-gen's SwiGLU + NVFP4 epilogue op for op (from its SASS); a = FC1 dequant, c = a x FC2 input scale: + // h' = (u a)(g sigmoid(g a)), sigmoid = rcp(ex2(-(g (a log2e))) + 1), block scale e4m3(max|h'| c / 6), FP4 (h' / scale) c + const float a_l2e = mul_ftz(ag, 1.4426950216293334961f); + float hv[8]; + float amax = 0.f; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const float ex = ex2_approx(-(ga[i] * a_l2e)); + const float sg = rcp_approx(ex + 1.0f); + const float v = fmaf(ua[i], ag, 0.0f) * (ga[i] * sg); + hv[i] = v; + amax = fmaxf(amax, fabsf(v)); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + const uint32_t sf = cvt_e4m3(fmaxf(0.f, fabsf(amax * sc)) * 0.16666667163372039795f); + const float sff = e4m3_to_f32(sf); + const float osc = sff == 0.f ? 0.f : rcp_approx(sff); + float qv[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) qv[i] = mul_ftz(hv[i], osc) * sc; + const uint32_t packed = cvt_e2m1x8(qv); + const int gblk = 8 * h + blk; // 16-wide block over the 128 intermediates; k-blocks 0, 1 + const int chunk = gblk / 2, r8 = row % 8; + const uint32_t off = (uint32_t)((row / 8) * 1024 + r8 * 128 + ((chunk ^ r8) * 16) + (gblk % 2) * 8 + hf * 4); +#pragma unroll + for (int cp = 0; cp < 8; ++cp) *reinterpret_cast(s.ha + cp * 2048 + off) = packed; + if (hf == 0) reinterpret_cast(&s.hsf[row][gblk / 4])[gblk % 4] = (uint8_t)sf; + } + }; + // Rows of item it (its tokens in increasing order, at most 16) and their x rows / scale bytes, by warps 0 and 3; + // row info is double-buffered by item parity (the previous item's FC2 drains still read theirs). + auto gather = [&](int it) { + const int rb = it & 1; + if (warp == 0) set_rows(rb, s.item_mask[it], s.item_e[it]); + named_bar(5, 64); + copy_rows(rb, (warp == 0 ? 0 : 32) + lane, 64); + named_bar(5, 64); + if (warp == 0 && lane == 0) { + mbar_arrive(&s.x_full); + if (DIAG_ROUTE && it == 0) stamp(p, 14); + } + }; + + if (warp == 1) { + if (lane == 0) { + // ---------------- TMA producer: every item's weight stream through the ring + const uint64_t pol = policy_evict_first(); + for (int it = 0; it < nitems; ++it) { + const int e = s.item_e[it], base = base0 + it * kTotStages; + const int st0 = (it == 0 && spec_ok) ? SPEC_STAGES : 0; + for (int st = st0; st < kTotStages; ++st) { + const int g = base + st; + if (FC2_PF > 0 && st == kFc1Stages - FC2_PF_LEAD) { + // the FC2 tiles beyond the ring go to L2 now, while the last FC1 stages and h are in progress + for (int j = kStages; j < kStages + FC2_PF; ++j) prefetch_stage(kFc1Stages + j, e); + } + if (PF_SET != 0 && ((st == st0 && (PF_SET & 1)) || (st == kFc1Stages + 4 && (PF_SET & 4)) || + (st == kFc1Stages + 6 && (PF_SET & 8)))) { + // these stages wait for MMAs that are slow to start after a phase change; let HBM fetch them meanwhile + const int n = st == st0 ? PF_N1 : st == kFc1Stages + 4 ? PF_N2 : 4; + for (int k = st; k < st + n && k < kTotStages; ++k) prefetch_stage(k, e); + } + if (g >= kStages) mbar_wait(&s.empty[g % kStages], ((g / kStages) - 1) & 1); + issue_stage(g, st, e, pol); + if (STAGE_STAMPS && it == 0) stamp(p, 64 + st); + } + } + } + __syncwarp(); + } else if (warp == 2) { + if (lane == 0) { + // ---------------- MMA issuer (tensor pipe: MMAs only) + const uint32_t idesc = make_idesc_nvf4(128u, 256u); + if (!spec_ok) { // the speculative fill was another CTA's expert: recycle its ring stages without MMAs + for (int g = 0; g < SPEC_STAGES; ++g) { + mbar_wait(&s.full[g % kStages], (g / kStages) & 1); + mbar_wait(&s.sf_full[g % kStages], (g / kStages) & 1); + mbar_arrive(&s.empty[g % kStages]); + } + } + for (int it = 0; it < nitems; ++it) { + const int base = base0 + it * kTotStages, par = it & 1; + if (it > 0) mbar_wait(&s.acc_drained, (it - 1) & 1); + mbar_wait(&s.x_full, par); + tc_fence_after(); + int n = 0; + for (int st = 0; st < kFc1Stages; ++st) { + const int g = base + st, b = g % kStages, ks = st; + mbar_wait(&s.full[b], (g / kStages) & 1); + mbar_wait(&s.sf_full[b], (g / kStages) & 1); + tc_fence_after(); + if (st == 0 && it == 0) stamp(p, 3); + if (st == 0 && it > 0 && it == nitems - 1) stamp(p, 11); +#pragma unroll + for (int q = 0; q < 4; ++q, ++n) { + const int kb = ks * 4 + q; // k-blocks in K order: one accumulator, the stock summation order + umma_nvf4(tm, make_sdesc(smem_u32(s.xa + (kb / 4) * 2048) + (kb % 4) * 32, 0, 1024, 2), + make_sdesc(smem_u32(s.ring[b]) + q * 32, 0, 1024, 2), idesc, + tm + kColSfaX + (uint32_t)(kb % kSfaSlots) * 4, tm + kColSfb + (uint32_t)(n % kSfbSlots) * 8, + kb > 0); + } + tc_commit(&s.empty[b]); + } + tc_commit(&s.fc1_done); + mbar_wait(&s.h_ready, par); + tc_fence_after(); + if (it == 0) stamp(p, 6); + for (int st = kFc1Stages; st < kTotStages; ++st) { + const int g = base + st, b = g % kStages, f = st - kFc1Stages; + if (f == 4) { mbar_wait(&s.fc2_free, par); tc_fence_after(); } + mbar_wait(&s.full[b], (g / kStages) & 1); + mbar_wait(&s.sf_full[b], (g / kStages) & 1); + tc_fence_after(); +#pragma unroll + for (int tl = 0; tl < 2; ++tl) { + const int j = 2 * f + tl, jt = j % 8; +#pragma unroll + for (int kb = 0; kb < 2; ++kb) { + const int n2 = 64 + j * 2 + kb; // scale-factor slot, numbered on from FC1's 64 k-blocks + umma_nvf4(tm, make_sdesc(smem_u32(s.ha) + kb * 32, 0, 1024, 2), + make_sdesc(smem_u32(s.ring[b] + tl * 16384) + kb * 32, 0, 512, 4), idesc, + tm + kColSfa + (uint32_t)(jt * 2 + kb) * 4, tm + kColSfb + (uint32_t)(n2 % kSfbSlots) * 8, + (jt | kb) != 0); + } + } + tc_commit(&s.empty[b]); + if (f % 4 == 3) tc_commit(&s.fc2_pass[f / 4]); + } + } + } + __syncwarp(); + } else if (warp == 0 || warp == 3) { + // ---------------- x gather, FC1 drains (warp 0: columns [0, 128); warp 4 drains [128, 256)) and the epilogue + const int et = (warp == 0 ? 0 : 32) + lane; // 0..63 + gather(0); + for (int it = 0; it < nitems; ++it) { + const int par = it & 1; + load_scales(it); + if (warp == 0) { + mbar_wait(&s.fc1_done, par); + tc_fence_after(); + if (lane == 0 && it == 0) stamp(p, 4); + drain_fc1(0); + if (lane == 0) { + mbar_arrive(&s.acc_ready); + if (it == 0 && !DIAG_ROUTE) stamp(p, 13); + } + } + mbar_wait(&s.acc_ready, par); + epi_tasks(0, et, 192); + if (warp == 3 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 14); + fence_proxy_async_smem(); + named_bar(6, 192); // h and its scale words are complete for the TMEM warpgroup + if (it + 1 < nitems) gather(it + 1); // FC1 of item it is done: x and its scale bytes are free + } + } else { + // ---------------- TMEM-side warpgroup (warps 4..7): q = lane quarter (TMEM lanes 32 q .. 32 q + 31) + const int q = warp - 4; + const uint32_t lane_addr = (uint32_t)(q * 32) << 16; + if (!spec_ok) { // the speculative fill is recycled without MMAs: release its scale-factor side + for (int g = 0; g < SPEC_STAGES; ++g) { + mbar_wait(&s.full[g % kStages], (g / kStages) & 1); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[g % kStages]); + } + } + for (int it = 0; it < nitems; ++it) { + const int base = base0 + it * kTotStages, par = it & 1, rb = it & 1; + load_scales(it); + mbar_wait(&s.x_full, par); + const int nrows = s.nrows[rb]; + auto stage_sf2 = [&](int st) { // FC2 tiles 2 f, 2 f + 1: k-blocks 0, 1 of their row blocks (rows 0..127 | 128..255) + const int g = base + st, b = g % kStages; + mbar_wait(&s.full[b], (g / kStages) & 1); + if (STAGE_STAMPS && threadIdx.x == 128 && it == 0) stamp(p, 16 + st); + if (g >= 4) mbar_wait(&s.empty[(g - 4) % kStages], ((g - 4) / kStages) & 1); + const uint8_t* sf = s.ring[b] + 32768; + const int n0 = 64 + (st - kFc1Stages) * 4; + for (int i = 0; i < 4; ++i) { + const int tl = i / 2, kb = i % 2; + const uint4 v0 = *reinterpret_cast(sf + tl * 2048 + kb * 512 + lane * 16); + const uint4 v1 = *reinterpret_cast(sf + tl * 2048 + 1024 + kb * 512 + lane * 16); + const uint32_t r[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; + tmem_st_x8(tm + lane_addr + kColSfb + (uint32_t)((n0 + i) % kSfbSlots) * 8, r); + } + tmem_st_wait(); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[b]); + }; + // FC1 stage st: weight scale factors of k-blocks [4 ks, 4 ks + 4) (8 columns each: row blocks 2 h | 2 h + 1) and + // the x scale factors of the same k-blocks (row m -> lane m % 32, column m / 32 of each 4-column atom, every + // quarter holds a copy; rows >= nrows, including the A rows read past the 16 real ones, get zero scales) + for (int st = 0; st < kFc1Stages; ++st) { + const int g = base + st, b = g % kStages, ks = st; + mbar_wait(&s.full[b], (g / kStages) & 1); + if (STAGE_STAMPS && threadIdx.x == 128 && it == 0) stamp(p, 16 + st); + if (g >= 4) mbar_wait(&s.empty[(g - 4) % kStages], ((g - 4) / kStages) & 1); + const uint8_t* sf = s.ring[b] + 32768; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const uint4 v0 = *reinterpret_cast(sf + i * 512 + lane * 16); + const uint4 v1 = *reinterpret_cast(sf + 2048 + i * 512 + lane * 16); + const uint32_t r[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; + tmem_st_x8(tm + lane_addr + kColSfb + (uint32_t)((st * 4 + i) % kSfbSlots) * 8, r); + } + uint32_t xr[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int kb = ks * 4 + i / 4; + xr[i] = (i % 4 == 0 && lane < nrows) ? *reinterpret_cast(&s.xsf[lane][kb * 4]) : 0u; + } + tmem_st_x16(tm + lane_addr + kColSfaX + (uint32_t)((ks * 4) % kSfaSlots) * 4, xr); + tmem_st_wait(); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[b]); + } + // EARLY_SF2: stage the first FC2 scale factors before h (they wait for FC2 data to land, which delays h) + for (int st = kFc1Stages; st < kFc1Stages + EARLY_SF2; ++st) stage_sf2(st); + if (q == 0) { // drain the upper 128 columns + mbar_wait(&s.fc1_done, par); + tc_fence_after(); + drain_fc1(2); + if (lane == 0) mbar_arrive(&s.acc_ready); + } + mbar_wait(&s.acc_ready, par); + epi_tasks(0, 64 + (warp - 4) * 32 + lane, 192); + fence_proxy_async_smem(); + named_bar(6, 192); // h complete + if (threadIdx.x == 128 && it == 0) stamp(p, 5); + // h scale factors: 16 atoms (tile jt, k-block kb) of 4 columns; row m = 32 c + lane is nonzero only in its tile + { + uint32_t r[32]; +#pragma unroll 1 + for (int half = 0; half < 2; ++half) { +#pragma unroll + for (int i = 0; i < 32; ++i) { + const int a = half * 8 + i / 4, col = i % 4; // atom a = 2 jt + kb + const int jt = a / 2, kb = a % 2; + const int row = 32 * col + lane; + r[i] = (row / 16 == jt && row % 16 < nrows) ? s.hsf[row % 16][kb] : 0u; + } + tmem_st_x32(tm + lane_addr + kColSfa + half * 32, r); + } + tmem_st_wait(); + } + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.h_ready); + if (q == 0 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 15); + for (int st = kFc1Stages + EARLY_SF2; st < kTotStages; ++st) { + if (st == kFc1Stages + 4) { + // pass 0 is fully staged; drain it once its MMAs are done, then free the accumulator for pass 1 + mbar_wait(&s.fc2_pass[0], par); + tc_fence_after(); + if (q == 0 && lane == 0 && it == 0) stamp(p, 7); + fc2_drain16<0, 4>(p, s, tm, q, lane, 0, a2, rb); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.fc2_free); + } + stage_sf2(st); + } + mbar_wait(&s.fc2_pass[1], par); + tc_fence_after(); + if (q == 0 && lane == 0 && it == 0) stamp(p, 12); + if (TAIL8 && it + 1 == nitems) fc2_drain16<0, 2>(p, s, tm, q, lane, 1, a2, rb); + else fc2_drain16<0, 4>(p, s, tm, q, lane, 1, a2, rb); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.acc_drained); + } + } + if (TAIL8 && warp < 4) { // the last item's last pass, upper columns: every warp is free now (TMEM lane quarter = warp) + const int it = nitems - 1; + load_scales(it); + mbar_wait(&s.fc2_pass[1], it & 1); + tc_fence_after(); + fc2_drain16<2, 4>(p, s, tm, warp, lane, 1, a2, it & 1); + } + tc_fence_before(); + named_bar(7, kThreads); // every warp, including the drains, before TMEM is released + if (threadIdx.x == 128) stamp(p, 8); + if (warp == 2) tmem_dealloc(tm, 512); +} + +static_assert(sizeof(Smem) + 1024 == Q35_MKM_SMEM_BYTES, "the launcher's dynamic shared memory size"); + +} // namespace qwen_mkm_moe_m32 + + +} // namespace sglang diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m64.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m64.cu new file mode 100644 index 00000000..544757e9 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m64.cu @@ -0,0 +1,1448 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Fused small-M NVFP4 MoE expert phase for Qwen3.5 TP8 decode verify (M <= 64 tokens, top-10, 512 experts, +// hidden 4096, intermediate 128 per rank): the top-10 of the router logits in the trtllm-gen routing's order and +// weights, FC1, SwiGLU, the NVFP4 requantization and FC2 in one launch over the trtllm-gen weight layout. Writes the +// deferred-finalize rows (row t * 10 + k, identity permutation) and the expert weights. +// +// The 32-token engine (mkm_moe_m32.cu) for up to 64 tokens: the work is (expert, row group of at most 16 of its +// tokens) units, so an expert that more than 16 tokens chose runs as several units on different CTAs. CTA c speculates +// on slot c's expert and runs its first unit when slot c is that expert's first; every other unit (experts first seen +// past the grid, further row groups) goes to the CTAs whose slot repeats an earlier expert, then round robin to +// CTAs 0, 1, ... Per unit the arithmetic and summation order are the 16-token engine's: rows bit for bit +// trtllm-gen's. +// +// Built with NVRTC: the kernel has C linkage (q35_mkm_moe_m64), and the launcher's dynamic shared memory size +// (Q35_MKM_SMEM_BYTES) is checked against the layout at compile time. +#include +#include +#include + +// NVRTC has no , or +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef int int32_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef unsigned long long uintptr_t; +typedef struct __align__(64) { + uint64_t opaque[16]; +} CUtensorMap; +#ifndef INFINITY +#define INFINITY __int_as_float(0x7f800000) +#endif +// Minimal SM100 (tcgen05 / TMA / mbarrier) PTX helpers for the small-M NVFP4 MoE engine. + +namespace qwen35 { + +__device__ __forceinline__ uint32_t smem_u32(const void* p) { + return static_cast(__cvta_generic_to_shared(p)); +} + +// ---------------------------------------------------------------- mbarrier +__device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(smem_u32(bar)), "r"(count) : "memory"); +} +__device__ __forceinline__ void fence_mbar_init() { + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} +__device__ __forceinline__ void mbar_expect_tx(uint64_t* bar, uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smem_u32(bar)), "r"(bytes) : "memory"); +} +__device__ __forceinline__ void mbar_arrive(uint64_t* bar) { + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(smem_u32(bar)) : "memory"); +} +__device__ __forceinline__ bool mbar_try_wait(uint64_t* bar, uint32_t phase) { + uint32_t ok; + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ok) + : "r"(smem_u32(bar)), "r"(phase) + : "memory"); + return ok != 0; +} +__device__ __forceinline__ void mbar_wait(uint64_t* bar, uint32_t phase) { + while (!mbar_try_wait(bar, phase)) { + } +} + +// ---------------------------------------------------------------- TMA / bulk copies +__device__ __forceinline__ void tma_load_2d(void* dst, const void* tmap, uint64_t* bar, int c0, int c1) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes [%0], [%1, {%3, %4}], [%2];" :: + "r"(smem_u32(dst)), + "l"(tmap), "r"(smem_u32(bar)), "r"(c0), "r"(c1) + : "memory"); +} +__device__ __forceinline__ void bulk_load(void* dst, const void* src, uint32_t bytes, uint64_t* bar) { + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)) + : "memory"); +} +__device__ __forceinline__ void prefetch_tmap(const void* tmap) { + asm volatile("prefetch.tensormap [%0];" ::"l"(tmap) : "memory"); +} +// generic-proxy smem writes -> visible to the async proxy (UMMA / UTCCP operand reads) +__device__ __forceinline__ void fence_proxy_async_smem() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + +// ---------------------------------------------------------------- tcgen05 +__device__ __forceinline__ void tmem_alloc(uint32_t* dst_smem, uint32_t ncols) { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smem_u32(dst_smem)), + "r"(ncols) + : "memory"); +} +__device__ __forceinline__ void tmem_relinquish() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;" ::: "memory"); +} +__device__ __forceinline__ void tmem_dealloc(uint32_t taddr, uint32_t ncols) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(ncols) : "memory"); +} +__device__ __forceinline__ void tc_fence_before() { + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); +} +__device__ __forceinline__ void tc_fence_after() { + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); +} +// arrive on `bar` when all prior tcgen05 async ops issued by this thread complete +__device__ __forceinline__ void tc_commit(uint64_t* bar) { + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" ::"r"(smem_u32(bar)) + : "memory"); +} +// smem -> TMEM copy of one 512-byte scale-factor chunk (32 rows x 16 B), broadcast to the 4 lane quarters +__device__ __forceinline__ void utccp_32x128b_warpx4(uint32_t taddr, uint64_t sdesc) { + asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" ::"r"(taddr), "l"(sdesc) : "memory"); +} +// D(tmem) (+)= A(smem) * B(smem)^T with NVFP4 block-16 scales from TMEM +__device__ __forceinline__ void umma_nvf4(uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, + uint32_t tmem_sfa, uint32_t tmem_sfb, uint32_t accumulate) { + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %4, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%5], [%6], p;\n}\n" ::"r"( + tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(accumulate), "r"(tmem_sfa), "r"(tmem_sfb) + : "memory"); +} +// each warp reads its 32-lane quarter: 32 lanes x 16 consecutive 32-bit columns +__device__ __forceinline__ void tmem_ld_32x32b_x16(uint32_t taddr, uint32_t (&r)[16]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15}, [%16];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(taddr) + : "memory"); +} +// L2 prefetch of a contiguous global range (bytes multiple of 16, 16-byte aligned) +__device__ __forceinline__ void bulk_prefetch_l2(const void* gptr, uint32_t bytes) { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(gptr), "r"(bytes) : "memory"); +} +__device__ __forceinline__ void tmem_ld_wait() { asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); } +// each warp writes its 32-lane quarter: 32 lanes x 16 consecutive 32-bit columns +__device__ __forceinline__ void tmem_st_32x32b_x16(uint32_t taddr, const uint32_t (&r)[16]) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16};" ::"r"( + taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_wait() { asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); } + +// ---------------------------------------------------------------- descriptors +// SM100 shared-memory matrix descriptor (version 1). layout: 0 none, 2 SW128, 4 SW64, 6 SW32. +__device__ __forceinline__ uint64_t make_sdesc(uint32_t saddr, uint32_t lbo, uint32_t sbo, uint32_t layout) { + uint64_t d = 0; + d |= (uint64_t)((saddr >> 4) & 0x3FFF); + d |= (uint64_t)((lbo >> 4) & 0x3FFF) << 16; + d |= (uint64_t)((sbo >> 4) & 0x3FFF) << 32; + d |= (uint64_t)1 << 46; // version + d |= (uint64_t)(layout & 0x7) << 61; + return d; +} +// Block-scaled instruction descriptor: E2M1 x E2M1, UE4M3 scales, both K-major, F32 accumulate. +__host__ __device__ constexpr uint32_t make_idesc_nvf4(uint32_t M, uint32_t N) { + return (1u << 7) // a_format E2M1 + | (1u << 10) // b_format E2M1 + | ((N >> 3) << 17) // n_dim + | (0u << 23) // scale format UE4M3 + | ((M >> 4) << 24); // m_dim +} + +} // namespace qwen35 + +namespace sglang { + +namespace qwen_mkm_moe_m64 { + +using namespace qwen35; + + +constexpr int kH = 4096, kE = 512, kTopK = 10, kMaxTok = 64; +constexpr int kRows = 16; // A rows (tokens) per item +// per CTA: at most 1 + ceil(units / CTAs) items; units <= 520 at 64 tokens (each expert's ceil(tokens / 16)) over +// >= 107 CTAs (the launcher's bound) +constexpr int kMaxItems = 8; +constexpr int kRounds = (kMaxTok + 6) / 7; // routing rounds of the 7 routing warps +constexpr int kI = 128; // intermediates per expert per rank (TP8) +constexpr int kN1 = 2 * kI; // W13 rows per expert (gate + up) +constexpr int kStageBytes = 36864; // 32 KB weights + 4 KB scale factors +constexpr int kStages = 4; +constexpr int kW13BoxRows = 128; // rows of one W13 TMA box; the host tensor map must use the same +constexpr int kFc1Stages = 16; // the expert's 256 W13 rows, 4 k-blocks each (full K = 4096) +constexpr int kFc2Stages = 8; // two tiles of 256 W2 rows x 64 B each +constexpr int kTotStages = kFc1Stages + kFc2Stages; +constexpr int kThreads = 256; +// TMEM columns. FC1: accumulator [0,256), x SFA ring [256,320) (16 slots x 4), SFB ring [384,512) (16 slots x 8). +// FC2: accumulator [0,256), h SFA [256,320) (16 atoms x 4), SFB ring [384,512) (16 slots x 8). +constexpr uint32_t kColSfaX = 256, kColSfa = 256, kColSfb = 384; +constexpr int kSfbSlots = 16, kSfaSlots = 16; +constexpr int kARow = 260; // FC1 accumulator row stride in floats: 16 B aligned, rows on distinct banks +#ifndef L2_EVICT_FIRST // stream expert weights with an L2 evict-first policy (they are read once per layer) +#define L2_EVICT_FIRST 1 +#endif +#ifndef TOPK_RANK +#define TOPK_RANK 0 +#endif +#ifndef SF_UTCCP_FC1 +#define SF_UTCCP_FC1 0 +#endif +#ifndef SPEC_STAGES // ring stages issued for the speculated expert before the routing verdict (rest: after it) +#define SPEC_STAGES 4 +#endif +#ifndef STAGE_STAMPS // DIAG: per-stage issue (slots 64 + st) and landing (16 + st, seen by warp 4) stamps +#define STAGE_STAMPS 0 +#endif +#ifndef DIAG_SPEC // DIAG: warp 1's speculative ranking, stamps 16 = logits loaded, 17 = keys, 18 = top-10 (stride 32) +#define DIAG_SPEC 0 +#endif +#ifndef DIAG_ROUTE // DIAG: stamps 13 = warp 0's routing rounds done, 15 = their barrier, 14 = item 0's rows gathered +#define DIAG_ROUTE 0 +#endif +constexpr int kTraceStride = STAGE_STAMPS ? 128 : (DIAG_SPEC || DIAG_ROUTE) ? 32 : 16; +#ifndef SMEM_SPACE // 1: shared-memory accesses through Smem compile to shared-space instructions +#define SMEM_SPACE 1 +#endif +#ifndef PF_SET // L2 prefetch of the stages that wait for a ring slot at: 1 = start, 4 = FC2, 8 = FC2 pass 1 +// (none: from 56 rows every CTA streams and HBM is the bound, so these prefetches only compete with the ring's loads) +#define PF_SET 0 +#endif +#ifndef PF_N1 // stages prefetched at the start trigger +#define PF_N1 4 +#endif +#ifndef PF_N2 // stages prefetched at the FC2 trigger +#define PF_N2 4 +#endif +#ifndef TAILV // FC2 drains: bit 0 = last pass on all 8 warps; bit 1 = TMEM loads double-buffered; bit 2 = coalesced stores +// (0: measured at 28 / 56 rows, the plain drains cost the least) +#define TAILV 0 +#endif +#define TAIL8 (TAILV & 1) +#define DRAIN_PIPE ((TAILV >> 1) & 1) +#define DRAIN_SMEM ((TAILV >> 2) & 1) +static_assert(DRAIN_SMEM == 0, "x's shared memory holds the next item's rows while the drains run"); +#ifndef EARLY_SF2 // FC2 stages whose scale factors are staged before the h epilogue (old order: 2) +#define EARLY_SF2 0 +#endif +#ifndef GATHER_BATCH // x row granules each gathering thread loads before storing them +#define GATHER_BATCH 8 +#endif +#ifndef GATHER_ASYNC // 1: the gather copies by cp.async (every granule in flight at once, no register staging) +#define GATHER_ASYNC 1 +#endif +#ifndef SPEC_PF // stages of the speculated expert past the first ring fill prefetched to L2 while the routing runs +#define SPEC_PF 4 +#endif +#ifndef FC2_PF // FC2 tiles (beyond the first kStages) prefetched to L2 ahead of the ring; 0 = off +#define FC2_PF 0 +#endif +#ifndef FC2_PF_LEAD // issue the FC2 prefetch when the producer reaches FC1 stage kFc1Stages - FC2_PF_LEAD +#define FC2_PF_LEAD 4 +#endif + +struct Params { + const uint16_t* logits; // [M, 512] bf16 router logits (router_gemm_quant output) + uint16_t* expert_weights; // [M, 10] bf16 out: DynBlock softmax-over-top-10 weights (the last CTA writes) + int32_t* topk_ids; // [M, 10] out (the last CTA writes) + const uint8_t* x_sf; // [M, 256] linear e4m3 + const uint8_t* w13_sf; // [E, 2 row blocks x 64 k blocks x 512 B] (128x4 interleaved) + const uint8_t* w2_sf; // [E, 32 row blocks x 2 k blocks x 512 B] + const float* g1_alpha; // [E] gate dequant + const float* g1_scale_c; // [E] up dequant * FC2 input quant scale + const float* g2_alpha; // [E] + __nv_bfloat16* out; // [M * 10, 4096] + long long* trace; // optional [grid, 16] %globaltimer stamps + int num_tokens; + const uint8_t* x_fp4; // [M, 2048] NVFP4 activations (each item gathers its rows) + // [M, 10] the routing's expert ids (route_topk's, bit for bit the trtllm-gen routing): when set, the engine takes them + // in place of ranking the logits (and writes neither expert_weights nor topk_ids, which route_topk wrote) + const int32_t* ids_in; +}; + +struct __align__(1024) Smem { + uint8_t ring[kStages][kStageBytes]; + alignas(1024) uint8_t xa[32 * 1024]; // x: per 128-B K chunk c, 2 KB = SW128 atoms of item rows 0..7 | 8..15 + alignas(1024) uint8_t ha[16 * 1024]; // h: 8 copies of (rows 0..7 | rows 8..15) x 128 B (K = 128: 64 B), SW128 + float acc1[16 * kARow]; // FC1 accumulator rows 0..15 (item rows) x 256 columns + uint8_t xsf[16][256]; // x scale bytes of the item's rows, full K + uint32_t hsf[16][2]; // h scale words: row r, k-block kb (4 e4m3 bytes) + uint64_t full[kStages], empty[kStages], sf_full[kStages]; + uint64_t x_full, fc1_done, acc_ready, h_ready, fc2_pass[2], fc2_free, acc_drained, route_done; + uint32_t tmem; + int nitems, spec_ok; + int item_e[kMaxItems]; + uint64_t item_mask[kMaxItems]; // tokens of each item (bit t) + int nrows[2], rtok[2][kRows], rkslot[2][kRows]; // rows of the current / next item (by item parity) + int first[kE]; // first slot of each expert (atomicMin) + int cnt[kE]; // tokens of each expert (atomicAdd) + uint32_t fbits[kMaxTok]; // token t: its top-k slots that are their expert's first (bit k) + int ounits[kMaxTok]; // token t: the units its first-occurrence slots hand out (not a CTA's own) + uint8_t uslot[kMaxTok][kTopK]; // per slot: those units (0 unless a first occurrence) + alignas(16) uint32_t cand[8][32]; // read as uint4 by the rank-based top-k + int32_t ids[kMaxTok][kTopK]; +}; + +__device__ __forceinline__ uint64_t gtime() { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)::"memory"); + return t; +} +__device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} +__device__ __forceinline__ void cluster_arrive() { + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); +} +// no release: the phase only orders DSMEM traffic that mbarriers already completed, so pending global stores +// (stamps, routing outputs, output rows) need not drain before the arrive +__device__ __forceinline__ void cluster_arrive_relaxed() { + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} +__device__ __forceinline__ void cluster_wait() { + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); +} +__device__ __forceinline__ uint32_t mapa(uint32_t saddr, uint32_t rank) { + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(r) : "r"(saddr), "r"(rank)); + return r; +} +// smem (this CTA) -> smem (peer CTA in the cluster), completes tx bytes on the peer's mbarrier +__device__ __forceinline__ void bulk_s2s(uint32_t dst_cluster, uint32_t src, uint32_t bytes, uint32_t bar_cluster) { + asm volatile("cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + dst_cluster), + "r"(src), "r"(bytes), "r"(bar_cluster) + : "memory"); +} +// registers -> the peer CTA's shared memory; completes the bytes on the peer's mbarrier +__device__ __forceinline__ void st_async_v4(uint32_t dst, const uint4& v, uint32_t bar) { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"(dst), + "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "r"(bar) + : "memory"); +} +__device__ __forceinline__ void st_async_b32(uint32_t dst, uint32_t v, uint32_t bar) { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(dst), "r"(v), "r"(bar) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x4(uint32_t taddr, const uint4& v) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x4.b32 [%0], {%1, %2, %3, %4};" ::"r"(taddr), "r"(v.x), "r"(v.y), + "r"(v.z), "r"(v.w) + : "memory"); +} +// bar.sync counts arrivals per warp execution: reconverge first, so a warp whose lane 0 just ran a stamp or an +// mbarrier arrive cannot arrive twice and release the barrier early +__device__ __forceinline__ void named_bar(int id, int n) { + __syncwarp(); + asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(n) : "memory"); +} +__device__ __forceinline__ void tmem_ld_x32(uint32_t taddr, uint32_t (&r)[32]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,%19," + "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), "=r"(r[16]), + "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), "=r"(r[24]), + "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr) + : "memory"); +} +__device__ __forceinline__ void tmem_ld_x64(uint32_t taddr, uint32_t (&r)[64]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x64.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,%19," + "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31,%32,%33,%34,%35,%36,%37,%38,%39,%40,%41,%42,%43,%44,%45,%46,%47," + "%48,%49,%50,%51,%52,%53,%54,%55,%56,%57,%58,%59,%60,%61,%62,%63}, [%64];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), "=r"(r[16]), + "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), "=r"(r[24]), + "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]), "=r"(r[32]), + "=r"(r[33]), "=r"(r[34]), "=r"(r[35]), "=r"(r[36]), "=r"(r[37]), "=r"(r[38]), "=r"(r[39]), "=r"(r[40]), + "=r"(r[41]), "=r"(r[42]), "=r"(r[43]), "=r"(r[44]), "=r"(r[45]), "=r"(r[46]), "=r"(r[47]), "=r"(r[48]), + "=r"(r[49]), "=r"(r[50]), "=r"(r[51]), "=r"(r[52]), "=r"(r[53]), "=r"(r[54]), "=r"(r[55]), "=r"(r[56]), + "=r"(r[57]), "=r"(r[58]), "=r"(r[59]), "=r"(r[60]), "=r"(r[61]), "=r"(r[62]), "=r"(r[63]) + : "r"(taddr) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x8(uint32_t taddr, const uint32_t (&r)[8]) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x8.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" ::"r"(taddr), "r"(r[0]), + "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x16(uint32_t taddr, const uint32_t (&r)[16]) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x16.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16};" :: + "r"(taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x32(uint32_t taddr, const uint32_t (&r)[32]) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x32.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18," + "%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31,%32};" ::"r"(taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]), "r"(r[16]), "r"(r[17]), "r"(r[18]), + "r"(r[19]), "r"(r[20]), "r"(r[21]), "r"(r[22]), "r"(r[23]), "r"(r[24]), "r"(r[25]), "r"(r[26]), "r"(r[27]), + "r"(r[28]), "r"(r[29]), "r"(r[30]), "r"(r[31]) + : "memory"); +} +__device__ __forceinline__ float ex2_approx(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float mul_ftz(float a, float b) { + float y; + asm("mul.ftz.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ float rcp_approx(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ uint32_t cvt_e4m3(float a) { + uint32_t out; + asm("{\n .reg .b16 pair;\n .reg .f32 zero;\n mov.f32 zero, 0f00000000;\n" + " cvt.rn.satfinite.e4m3x2.f32 pair, zero, %1;\n cvt.u32.u16 %0, pair;\n}" + : "=r"(out) + : "f"(a)); + return out & 0xFFu; +} +__device__ __forceinline__ float e4m3_to_f32(uint32_t v) { + float out; + asm("{\n .reg .b16 pair;\n .reg .b32 h2;\n .reg .b16 lo, hi;\n cvt.u16.u32 pair, %1;\n" + " cvt.rn.f16x2.e4m3x2 h2, pair;\n mov.b32 {lo, hi}, h2;\n cvt.f32.f16 %0, lo;\n}" + : "=f"(out) + : "r"(v)); + return out; +} +__device__ __forceinline__ uint32_t cvt_e2m1x8(const float* v) { + uint32_t out; + asm("{\n .reg .b8 b0, b1, b2, b3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\n cvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\n cvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" + " mov.b32 %0, {b0, b1, b2, b3};\n}" + : "=r"(out) + : "f"(v[0]), "f"(v[1]), "f"(v[2]), "f"(v[3]), "f"(v[4]), "f"(v[5]), "f"(v[6]), "f"(v[7])); + return out; +} +// ---- router top-k, bit-exact with FlashInfer's single-block routing (routingCustom, RenormalizeNaive on bf16): +// key = bf16 order (cub float twiddle) above 65535 - expert, i.e. the higher score first, the smaller expert on ties. +__device__ __forceinline__ uint32_t topk_key(uint32_t bits, uint32_t expert) { + const uint32_t twiddled = bits ^ ((bits & 0x8000u) ? 0xFFFFu : 0x8000u); + return (twiddled << 16) | (65535u - expert); +} +__device__ __forceinline__ void topk_unkey(uint32_t key, uint32_t& bits, uint32_t& expert) { + expert = 65535u - (key & 0xFFFFu); + const uint32_t twiddled = key >> 16; + bits = twiddled ^ ((twiddled & 0x8000u) ? 0x8000u : 0xFFFFu); +} +__device__ __forceinline__ uint32_t max16(const uint32_t (&k)[16]) { + const uint32_t a = max(max(max(k[0], k[1]), max(k[2], k[3])), max(max(k[4], k[5]), max(k[6], k[7]))); + const uint32_t b = max(max(max(k[8], k[9]), max(k[10], k[11])), max(max(k[12], k[13]), max(k[14], k[15]))); + return max(a, b); +} +__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_sum_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.add.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_topk_of_lanes(uint32_t v, uint32_t lane) { + uint32_t mine = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(v); + v = v == top ? 0u : v; + mine = lane == k ? top : mine; + } + return mine; +} +__device__ __forceinline__ uint32_t count_greater(const uint32_t* s_cand, uint32_t v) { + uint32_t r = 0; +#pragma unroll + for (int j = 0; j < 32; j += 4) { + const uint4 q = *reinterpret_cast(s_cand + j); + r += (q.x > v) + (q.y > v) + (q.z > v) + (q.w > v); + } + return r; +} +// Top-10 keys of 512 (16 per lane), lane k < 10 holding the k-th largest (keys unique and nonzero). Ranks come from +// all-pairs compares over shared memory rather than chains of warp reductions. +__device__ __forceinline__ uint32_t warp_top10(uint32_t (&keys)[16], uint32_t lane, uint32_t* s_cand) { + // bound = the 10th largest lane maximum; every top-10 key is >= bound + const uint32_t m = max16(keys); + s_cand[lane] = m; + __syncwarp(); + const uint32_t r = count_greater(s_cand, m); + const uint32_t bound = __shfl_sync(0xffffffffu, m, __ffs(__ballot_sync(0xffffffffu, r == kTopK - 1)) - 1); + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) count += keys[i] >= bound; + // exclusive prefix of count over the lanes, from the ballots of its bits (count <= 16) + const uint32_t lt = (1u << lane) - 1u; + uint32_t offset = 0, total = 0; +#pragma unroll + for (int b = 0; b < 5; ++b) { + const uint32_t bb = __ballot_sync(0xffffffffu, (count >> b) & 1u); + offset += (uint32_t)__popc(bb & lt) << b; + total += (uint32_t)__popc(bb) << b; + } + if (total <= 32) { + __syncwarp(); // every lane has read the lane maxima +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) { + if (keys[i] >= bound) s_cand[offset++] = keys[i]; + } + if (lane >= total) s_cand[lane] = 0u; + __syncwarp(); + const uint32_t cand = s_cand[lane]; + const uint32_t rk = count_greater(s_cand, cand); + __syncwarp(); + if (lane < total && rk < kTopK) s_cand[rk] = cand; + __syncwarp(); + return lane < kTopK ? s_cand[lane] : 0u; + } + uint32_t mine = 0; +#pragma unroll 1 + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(max16(keys)); +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) keys[i] = keys[i] == top ? 0u : keys[i]; + mine = lane == k ? top : mine; + } + return mine; +} +// Top-10 keys of 512 (16 per lane), lane k holding the k-th largest (keys unique and nonzero). +__device__ __forceinline__ uint32_t warp_top10_redux(uint32_t (&keys)[16], uint32_t lane, uint32_t* s_cand) { + uint32_t v = max16(keys); + uint32_t bound = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + bound = warp_max_u32(v); + v = v == bound ? 0u : v; + } + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) count += keys[i] >= bound; + const uint32_t total = warp_sum_u32(count); + if (total <= 32) { + uint32_t offset = count; +#pragma unroll + for (uint32_t d = 1; d < 32; d <<= 1) { + const uint32_t up = __shfl_up_sync(0xffffffffu, offset, d); + offset += lane >= d ? up : 0u; + } + offset -= count; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) { + if (keys[i] >= bound) s_cand[offset++] = keys[i]; + } + __syncwarp(); + const uint32_t cand = lane < total ? s_cand[lane] : 0u; + __syncwarp(); + return warp_topk_of_lanes(cand, lane); + } + uint32_t mine = 0; +#pragma unroll 1 + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(max16(keys)); +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) keys[i] = keys[i] == top ? 0u : keys[i]; + mine = lane == k ? top : mine; + } + return mine; +} +// The top-10 keys of one logits row (16 bf16 a lane in v0, v1), lane k < 10 holding the k-th largest +__device__ __forceinline__ uint32_t row_top10(const uint4 v0, const uint4 v1, uint32_t lane, uint32_t* s_cand) { + uint32_t keys[16]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t w = i < 4 ? (&v0.x)[i] : (&v1.x)[i - 4]; + keys[2 * i] = topk_key(w & 0xFFFFu, lane * 16 + 2 * i); + keys[2 * i + 1] = topk_key(w >> 16, lane * 16 + 2 * i + 1); + } + return TOPK_RANK ? warp_top10(keys, lane, s_cand) : warp_top10_redux(keys, lane, s_cand); +} +// FlashInfer calcSoftmax (bf16 scores) as built with -use_fast_math: expf -> ex2.approx.ftz of x * log2(e), +// division -> div.full.ftz; reductions are cooperative_groups::reduce like the reference. +__device__ __forceinline__ float fast_expf(float x) { + float y; + asm("{\n .reg .f32 t;\n mul.ftz.f32 t, %1, 0f3FB8AA3B;\n ex2.approx.ftz.f32 %0, t;\n}" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float fast_div(float a, float b) { + float y; + asm("div.full.ftz.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ uint16_t softmax_topk(uint16_t score_bits, uint32_t lane) { + namespace cg = cooperative_groups; + const auto warp = cg::tiled_partition<32>(cg::this_thread_block()); + const float si = __uint_as_float((uint32_t)score_bits << 16); + float max_score = -INFINITY; + if (lane < kTopK) max_score = si >= max_score ? si : max_score; + max_score = cg::reduce(warp, max_score, cg::greater()); + float sum_score = 0.f, new_score = 0.f; + if (lane < kTopK) { + new_score = fast_expf(__fsub_rn(si, max_score)); + sum_score = __fadd_rn(sum_score, new_score); + } + sum_score = cg::reduce(warp, sum_score, cg::plus()); + uint16_t out = score_bits; + if (lane < kTopK) { + const __nv_bfloat16 b = __float2bfloat16_rn(fast_div(new_score, sum_score)); + out = *reinterpret_cast(&b); + } + return out; +} + +__device__ __forceinline__ uint64_t policy_evict_first() { + uint64_t pol; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(pol)); + return pol; +} +__device__ __forceinline__ void tma_load_2d_w(void* dst, const void* tmap, uint64_t* bar, int c0, int c1, uint64_t pol) { + if (L2_EVICT_FIRST) + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%3, %4}], [%2], %5;" ::"r"(smem_u32(dst)), + "l"(tmap), "r"(smem_u32(bar)), "r"(c0), "r"(c1), "l"(pol) + : "memory"); + else + tma_load_2d(dst, tmap, bar, c0, c1); +} +__device__ __forceinline__ void bulk_load_w(void* dst, const void* src, uint32_t bytes, uint64_t* bar, uint64_t pol) { + if (L2_EVICT_FIRST) + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)), "l"(pol) + : "memory"); + else + bulk_load(dst, src, bytes, bar); +} + +__device__ __forceinline__ void stamp(const Params& p, int slot) { + if (p.trace) p.trace[blockIdx.x * kTraceStride + slot] = (long long)gtime(); +} +// stamp after `dep` is available (the branch needs its value; keys are never 0) +__device__ __forceinline__ void stamp_dep(const Params& p, int slot, uint32_t dep) { + if (dep == 0u) __trap(); + stamp(p, slot); +} +__device__ __forceinline__ void tma_prefetch_2d(const void* tmap, int c0, int c1) { + asm volatile("cp.async.bulk.prefetch.tensor.2d.L2.global.tile [%0, {%1, %2}];" ::"l"(tmap), "r"(c0), "r"(c1) : "memory"); +} +// Physical W13 column (accumulator column) holding reordered row q: block-32 shuffle of FlashInfer's +// get_shuffle_matrix_a_row_indices (epilogue_tile_m = 128); reordered row 2j is up_j, 2j + 1 is gate_j. +__device__ __forceinline__ int w13_col(int q) { return (q & ~31) + ((q & 31) % 4) * 8 + (q & 31) / 4; } +__host__ __device__ constexpr int w13_logical(int p) { + return ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) % 2 ? ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) / 2 + : 128 + ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) / 2; +} + + +// FC2 drain of one pass: TMEM lane quarter q (tiles 2 q, 2 q + 1 of the pass), 256 accumulator columns +__device__ __forceinline__ void reg_fence64(uint32_t (&r)[64]) { // keep uses of r after the tcgen05.wait::ld +#pragma unroll + for (int i = 0; i < 64; ++i) asm volatile("" : "+r"(r[i])); +} + +// FC2 accumulator columns [64 C0, 64 C1) of this warp's lane quarter -> bf16 gemm2_out rows (row r of a tile = the +// r-th row of the item, row info buffer rb) +template +__device__ __forceinline__ void fc2_drain16(const Params& p, Smem& s, uint32_t tm, int q, int lane, int pass, float a2, + int rb) { + const int row = q * 32 + lane, tile = pass * 8 + row / 16, r = row % 16; + const bool valid = r < s.nrows[rb]; + __nv_bfloat16* orow = valid ? p.out + (size_t)(s.rtok[rb][r] * kTopK + s.rkslot[rb][r]) * kH + tile * 256 : nullptr; + const uint32_t base = tm + ((uint32_t)(q * 32) << 16); + auto put = [&](const uint32_t (&r)[64], int c) { + if (valid) { +#pragma unroll + for (int b32 = 0; b32 < 2; ++b32) { + uint32_t pk[16]; +#pragma unroll + for (int l = 0; l < 32; l += 2) { + const float v0 = __uint_as_float(r[b32 * 32 + (l % 4) * 8 + l / 4]) * a2; + const float v1 = __uint_as_float(r[b32 * 32 + ((l + 1) % 4) * 8 + (l + 1) / 4]) * a2; + const __nv_bfloat162 b2 = __floats2bfloat162_rn(v0, v1); + pk[l / 2] = *reinterpret_cast(&b2); + } +#pragma unroll + for (int i = 0; i < 4; ++i) + *reinterpret_cast(orow + c * 64 + b32 * 32 + i * 8) = + make_uint4(pk[4 * i], pk[4 * i + 1], pk[4 * i + 2], pk[4 * i + 3]); + } + } + }; + if (DRAIN_PIPE) { // the next chunk's TMEM load is in flight while this chunk converts and stores + uint32_t ra[64], rb2[64]; + tmem_ld_x64(base + C0 * 64, ra); + tmem_ld_wait(); + reg_fence64(ra); +#pragma unroll + for (int c = C0; c < C1; c += 2) { + if (c + 1 < C1) tmem_ld_x64(base + (uint32_t)((c + 1) * 64), rb2); + put(ra, c); + if (c + 1 < C1) { + tmem_ld_wait(); + reg_fence64(rb2); + if (c + 2 < C1) tmem_ld_x64(base + (uint32_t)((c + 2) * 64), ra); + put(rb2, c + 1); + if (c + 2 < C1) { tmem_ld_wait(); reg_fence64(ra); } + } + } + } else { +#pragma unroll 1 + for (int c = C0; c < C1; ++c) { + uint32_t r[64]; + tmem_ld_x64(base + (uint32_t)(c * 64), r); + tmem_ld_wait(); + put(r, c); + } + } +} + +extern "C" __global__ void __launch_bounds__(kThreads, 1) +q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant__ CUtensorMap tm2, const Params p) { + extern __shared__ __align__(1024) uint8_t raw[]; +#if SMEM_SPACE // pointer arithmetic on the __shared__ array keeps the shared address space (STS/LDS, not generic) + Smem& s = *reinterpret_cast(raw + ((1024u - (smem_u32(raw) & 1023u)) & 1023u)); +#else + Smem& s = *reinterpret_cast((reinterpret_cast(raw) + 1023) & ~uintptr_t(1023)); +#endif + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + const int slot = blockIdx.x; // the (token, top-k) slot this CTA speculates on; gridDim.x <= M * 10 + const int M = p.num_tokens; + for (int i = threadIdx.x; i < kE; i += kThreads) { + s.first[i] = 0x7FFFFFFF; + s.cnt[i] = 0; + } + if (threadIdx.x < kMaxTok) { + s.fbits[threadIdx.x] = 0u; + s.ounits[threadIdx.x] = 0; + } + if (threadIdx.x == 0) { + stamp(p, 0); + for (int i = 0; i < kStages; ++i) { + mbar_init(&s.full[i], 1); + mbar_init(&s.empty[i], 1); + mbar_init(&s.sf_full[i], 4); + } + // every barrier below completes once per item, so item i waits parity i & 1 + mbar_init(&s.x_full, 1); + mbar_init(&s.fc1_done, 1); + mbar_init(&s.acc_ready, 2); // warps 0 and 4 drained their columns into acc1 + mbar_init(&s.h_ready, 4); + mbar_init(&s.fc2_pass[0], 1); + mbar_init(&s.fc2_pass[1], 1); + mbar_init(&s.fc2_free, 4); + mbar_init(&s.acc_drained, 4); // warps 4..7 read the last FC2 pass: the next item's FC1 may overwrite columns 0..255 + mbar_init(&s.route_done, 1); + fence_mbar_init(); + } + if (warp == 2) { tmem_alloc(&s.tmem, 512); tmem_relinquish(); tc_fence_before(); } + if (warp == 1 && lane == 0) { prefetch_tmap(&tm13); prefetch_tmap(&tm2); } + __syncthreads(); + tc_fence_after(); + const uint32_t tm = s.tmem; + const int my_t = slot / kTopK, my_k = slot % kTopK; + // ring stage g (global over the CTA's items) of item-relative stage st of expert e + auto issue_stage = [&](int g, int st, int e, uint64_t pol) { + const int b = g % kStages; + mbar_expect_tx(&s.full[b], kStageBytes); + uint8_t* dst = s.ring[b]; + if (st < kFc1Stages) { + // the expert's rows as two boxes of 128 rows, K bytes [128 st, 128 st + 128) + tma_load_2d_w(dst, &tm13, &s.full[b], st * 128, e * kN1, pol); + tma_load_2d_w(dst + 16384, &tm13, &s.full[b], st * 128, e * kN1 + 128, pol); + // scale atoms of row blocks 0 and 1, k-blocks [4 st, 4 st + 4) + const uint8_t* sf = p.w13_sf + (size_t)e * (kN1 * kH / 16); + bulk_load_w(dst + 32768, sf + (size_t)(st * 4) * 512, 2048, &s.full[b], pol); + bulk_load_w(dst + 32768 + 2048, sf + (size_t)(64 + st * 4) * 512, 2048, &s.full[b], pol); + } else { + // W2 tiles 2 f and 2 f + 1 (256 rows x 64 B each) and their scale atoms (2 row blocks x 2 k-blocks) + const int f = st - kFc1Stages; +#pragma unroll + for (int tl = 0; tl < 2; ++tl) { + const int row0 = (2 * f + tl) * 256; + tma_load_2d_w(dst + tl * 16384, &tm2, &s.full[b], 0, e * kH + row0, pol); + bulk_load_w(dst + 32768 + tl * 2048, p.w2_sf + (size_t)e * (kH * 8) + (size_t)(row0 / 128) * 2 * 512, 2048, + &s.full[b], pol); + } + } + }; + auto prefetch_stage = [&](int st, int e) { // the same bytes as issue_stage, to L2 only + if (st < kFc1Stages) { + tma_prefetch_2d(&tm13, st * 128, e * kN1); + tma_prefetch_2d(&tm13, st * 128, e * kN1 + 128); + const uint8_t* sf = p.w13_sf + (size_t)e * (kN1 * kH / 16); + bulk_prefetch_l2(sf + (size_t)(st * 4) * 512, 2048); + bulk_prefetch_l2(sf + (size_t)(64 + st * 4) * 512, 2048); + } else { + const int row0 = (st - kFc1Stages) * 512; + tma_prefetch_2d(&tm2, 0, e * kH + row0); + tma_prefetch_2d(&tm2, 0, e * kH + row0 + 256); + bulk_prefetch_l2(p.w2_sf + (size_t)e * (kH * 8) + (size_t)(row0 / 128) * 2 * 512, 4096); + } + }; + // Rows of an item (the tokens of mask in increasing order, at most 16) into row buffer rb, by one warp + auto set_rows = [&](int rb, uint64_t mask, int ei) { +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int t = lane + 32 * h; + if (t < M && ((mask >> t) & 1ull)) { + int ks = 0; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + if (s.ids[t][k] == ei) ks = k; + const int r = __popcll(mask & ((1ull << t) - 1ull)); + s.rtok[rb][r] = t; + s.rkslot[rb][r] = ks; + } + } + if (lane == 0) s.nrows[rb] = __popcll(mask); + }; + // their x rows and scale bytes into xa / xsf, by threads tid = 0 .. nt - 1 + auto copy_rows = [&](int rb, int tid, int nt) { + const int n1 = s.nrows[rb] * 128, n = s.nrows[rb] * 144; // 16-B granules: x rows, then their scale bytes + if (GATHER_ASYNC) { + for (int idx = tid; idx < n; idx += nt) { + const uint8_t* src; + uint8_t* dst; + if (idx < n1) { + const int r = idx >> 7, c = (idx >> 3) & 15, q = idx & 7; + src = p.x_fp4 + (size_t)s.rtok[rb][r] * (kH / 2) + c * 128 + q * 16; + dst = s.xa + c * 2048 + (r / 8) * 1024 + (r % 8) * 128 + ((q ^ (r % 8)) * 16); + } else { + const int r = (idx - n1) >> 4, q = (idx - n1) & 15; + src = p.x_sf + (size_t)s.rtok[rb][r] * (kH / 16) + q * 16; + dst = &s.xsf[r][q * 16]; + } + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(smem_u32(dst)), "l"(src) : "memory"); + } + asm volatile("cp.async.wait_all;" ::: "memory"); + fence_proxy_async_smem(); + return; + } + // GATHER_BATCH loads a thread ahead of their stores: the global latency is paid once a batch + for (int i0 = tid; i0 < n; i0 += GATHER_BATCH * nt) { + uint4 v[GATHER_BATCH]; +#pragma unroll + for (int j = 0; j < GATHER_BATCH; ++j) { + const int idx = i0 + j * nt; + if (idx < n1) { // row r, K chunk c (128 B), 16-B granule q of the chunk + const int r = idx >> 7, c = (idx >> 3) & 15, q = idx & 7; + v[j] = *reinterpret_cast(p.x_fp4 + (size_t)s.rtok[rb][r] * (kH / 2) + c * 128 + q * 16); + } else if (idx < n) { + const int r = (idx - n1) >> 4, q = (idx - n1) & 15; + v[j] = *reinterpret_cast(p.x_sf + (size_t)s.rtok[rb][r] * (kH / 16) + q * 16); + } + } +#pragma unroll + for (int j = 0; j < GATHER_BATCH; ++j) { + const int idx = i0 + j * nt; + if (idx < n1) { + const int r = idx >> 7, c = (idx >> 3) & 15, q = idx & 7; + *reinterpret_cast(s.xa + c * 2048 + (r / 8) * 1024 + (r % 8) * 128 + ((q ^ (r % 8)) * 16)) = v[j]; + } else if (idx < n) { + const int r = (idx - n1) >> 4, q = (idx - n1) & 15; + *reinterpret_cast(&s.xsf[r][q * 16]) = v[j]; + } + } + } + fence_proxy_async_smem(); // the tensor core reads xa through the async proxy + }; + // routing: warps 0, 2..7 take tokens 0..31 (DynBlock-exact top-10 + softmax), up to kRounds rounds; warp 1 recomputes + // token my_t's top-10 and starts this CTA's weight stream at once (a wrong guess only wastes its first ring fill) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (threadIdx.x == 32) stamp(p, 10); + auto top10v = [&](const uint4 v0, const uint4 v1) -> uint32_t { + return row_top10(v0, v1, (uint32_t)lane, s.cand[warp]); + }; + auto top10 = [&](int rtok) -> uint32_t { + const uint4* src = reinterpret_cast(p.logits + (size_t)rtok * kE + lane * 16); + return top10v(src[0], src[1]); + }; + if (warp == 1) { + // speculate on this CTA's own slot: its expert's weight stream starts before the routing verdict + uint32_t key = 0; + if (p.ids_in) { + } else if (DIAG_SPEC) { + const uint4* src = reinterpret_cast(p.logits + (size_t)my_t * kE + lane * 16); + const uint4 v0 = src[0], v1 = src[1]; + if (lane == 0) stamp_dep(p, 16, (v0.x | v0.y | v0.z | v0.w | v1.x | v1.y | v1.z | v1.w) | 1u); + key = top10v(v0, v1); + if (lane == 0) stamp_dep(p, 18, key | 1u); + } else { + key = top10(my_t); + } + uint32_t bits, espec; + if (p.ids_in) espec = (uint32_t)p.ids_in[slot]; + else topk_unkey(__shfl_sync(0xffffffffu, key, my_k), bits, espec); + if (lane == 0) { + stamp(p, 9); + const uint64_t pol = policy_evict_first(); + for (int st = 0; st < SPEC_STAGES; ++st) { + issue_stage(st, st, (int)espec, pol); + if (STAGE_STAMPS) stamp(p, 64 + st); + } + // HBM is otherwise idle until the routing is done: the next stages go to L2 now (another CTA streams the + // expert when this slot is not its first) + for (int st = SPEC_STAGES; st < SPEC_STAGES + SPEC_PF && st < kTotStages; ++st) prefetch_stage(st, (int)espec); + stamp(p, 2); + } + mbar_wait(&s.route_done, 0); + } else if (p.ids_in) { + // the routing as given: the 224 routing threads load every slot's expert at once, then record and count them + const int tid = (warp == 0 ? 0 : warp - 1) * 32 + lane, S = M * kTopK; + constexpr int kPer = (kMaxTok * kTopK + 223) / 224; + int ex[kPer]; +#pragma unroll + for (int j = 0; j < kPer; ++j) ex[j] = tid + 224 * j < S ? p.ids_in[tid + 224 * j] : 0; +#pragma unroll + for (int j = 0; j < kPer; ++j) { + const int s2 = tid + 224 * j; + if (s2 < S) { + s.ids[s2 / kTopK][s2 % kTopK] = ex[j]; + atomicMin(&s.first[ex[j]], s2); + atomicAdd(&s.cnt[ex[j]], 1); + } + } + } else { + const int base = warp == 0 ? 0 : warp - 1; + // every round's logits are requested before the first ranking, so later rounds do not pay the load latency again + uint4 lv[kRounds][2]; +#pragma unroll + for (int r = 0; r < kRounds; ++r) { + const int rtok = base + 7 * r; + if (rtok < M) { + const uint4* src = reinterpret_cast(p.logits + (size_t)rtok * kE + lane * 16); + lv[r][0] = src[0]; + lv[r][1] = src[1]; + } + } + // the last CTA writes the routing outputs (the further items start at CTA 0); the others need the ids only + const bool writer = blockIdx.x == gridDim.x - 1; +#pragma unroll + for (int r = 0; r < kRounds; ++r) { + const int rtok = base + 7 * r; + if (rtok >= M) break; + const uint32_t key = top10v(lv[r][0], lv[r][1]); + uint32_t bits, expert; + topk_unkey(key, bits, expert); + if (lane < kTopK) { + s.ids[rtok][lane] = (int32_t)expert; + atomicMin(&s.first[expert], rtok * kTopK + lane); + atomicAdd(&s.cnt[expert], 1); + } + if (writer) { + const uint16_t wbits = softmax_topk((uint16_t)bits, (uint32_t)lane); + if (lane < kTopK) { + p.expert_weights[rtok * kTopK + lane] = wbits; + p.topk_ids[rtok * kTopK + lane] = (int32_t)expert; + } + } + __syncwarp(); + } + } + if (warp != 1) { + const int base = warp == 0 ? 0 : warp - 1; + if (DIAG_ROUTE && threadIdx.x == 0) stamp(p, 13); + named_bar(4, 224); + if (DIAG_ROUTE && threadIdx.x == 0) stamp(p, 15); + // each routing warp marks its tokens' slots that are their expert's first and the units they hand out (s.first and + // s.cnt are final after the barrier): a first occurrence at slot s2 of an expert with n tokens hands out its + // ceil(n / 16) row groups, less the first one when s2 < G (that CTA runs it) + const int G = (int)gridDim.x; + { // every slot at once over the 224 routing threads (per-token sums by shared-memory atomics) + const int tid = base * 32 + lane, S = M * kTopK; +#pragma unroll + for (int j = 0; j < (kMaxTok * kTopK + 223) / 224; ++j) { + const int s2 = tid + 224 * j; + if (s2 < S) { + const int t = s2 / kTopK, k = s2 % kTopK, e2 = s.ids[t][k]; + const bool f = s.first[e2] == s2; + const int u = f ? (s.cnt[e2] + kRows - 1) / kRows - (s2 < G ? 1 : 0) : 0; + s.uslot[t][k] = (uint8_t)u; + if (f) { + atomicOr(&s.fbits[t], 1u << k); + atomicAdd(&s.ounits[t], u); + } + } + } + } + named_bar(4, 224); + if (DIAG_ROUTE && threadIdx.x == 0) stamp(p, 20); + if (warp == 0) { + // Work items: (expert, row group) units. Slot c's expert's first unit belongs to CTA c when c is the expert's first + // slot (its speculative stream already carries it). The other units, in order of their expert's first slot then + // row group, go to the CTAs whose own slot repeats an earlier expert, in order; units beyond those CTAs become + // further items of CTAs 0, 1, ... Every CTA derives the same assignment from the same routing. + // Per-token prefixes over 64 tokens: lane l holds tokens l and l + 32. + auto scan = [&](int v) -> int { // inclusive prefix over the lanes +#pragma unroll + for (int d = 1; d < 32; d <<= 1) { + const int up = __shfl_up_sync(0xffffffffu, v, d); + v += lane >= d ? up : 0; + } + return v; + }; + const uint32_t fb0 = lane < M ? s.fbits[lane] : 0u, fb1 = lane + 32 < M ? s.fbits[lane + 32] : 0u; + const int ou0 = lane < M ? s.ounits[lane] : 0, ou1 = lane + 32 < M ? s.ounits[lane + 32] : 0; + const int fi0 = scan(__popc(fb0)), fi1 = scan(__popc(fb1)) + __shfl_sync(0xffffffffu, fi0, 31); + const int oi0 = scan(ou0), oi1 = scan(ou1) + __shfl_sync(0xffffffffu, oi0, 31); + // token t's exclusive prefixes and first bits (t < M) + auto tok = [&](int t, int& fex, int& oex, uint32_t& fb) { + const int l = t & 31; + const int fin = __shfl_sync(0xffffffffu, t < 32 ? fi0 : fi1, l); + const int oin = __shfl_sync(0xffffffffu, t < 32 ? oi0 : oi1, l); + fb = __shfl_sync(0xffffffffu, t < 32 ? fb0 : fb1, l); + oex = oin - __shfl_sync(0xffffffffu, t < 32 ? ou0 : ou1, l); + fex = fin - __popc(fb); + }; + const int Uf = __shfl_sync(0xffffffffu, fi1, 31); // distinct experts + const int O = __shfl_sync(0xffffffffu, oi1, 31); // units handed out + auto before = [&](int s2) -> int { // first-occurrence slots before slot s2 (s2 <= M * 10) + const int t = s2 / kTopK; + if (t >= M) return Uf; + int fex, oex; + uint32_t fb; + tok(t, fex, oex, fb); + return fex + __popc(fb & ((1u << (s2 % kTopK)) - 1u)); + }; + // the n-th (0-based) handed-out unit: (expert, row group) + auto nth_unit = [&](int n, int& e, int& grp) { + const int t = __popc(__ballot_sync(0xffffffffu, lane < M && oi0 <= n)) + + __popc(__ballot_sync(0xffffffffu, lane + 32 < M && oi1 <= n)); + int fex, oex; + uint32_t fb; + tok(t, fex, oex, fb); + int k = 0, left = n - oex; + while (left >= (int)s.uslot[t][k]) left -= s.uslot[t][k++]; + e = s.ids[t][k]; + grp = left + (t * kTopK + k < G ? 1 : 0); // the first group of a slot below G is its own CTA's + }; + const int nf_c = before(slot), nf_g = before(G); + if (DIAG_ROUTE && lane == 0) stamp(p, 21); + const int xc = s.ids[my_t][my_k]; + const bool owner = s.first[xc] == slot; + const int D = G - nf_g, dc = slot - nf_c; // duplicate CTAs; this one's rank among them + // the whole warp walks the items (nth_unit is warp-collective); lane 0 writes them + int ni = 0; + auto add = [&](int e, int grp) { + if (ni >= kMaxItems) __trap(); + bool h0 = false, h1 = false; +#pragma unroll + for (int k = 0; k < kTopK; ++k) { + h0 |= lane < M && s.ids[lane][k] == e; + h1 |= lane + 32 < M && s.ids[lane + 32][k] == e; + } + uint64_t mask = (uint64_t)__ballot_sync(0xffffffffu, h0) | ((uint64_t)__ballot_sync(0xffffffffu, h1) << 32); + for (int j = 0; j < grp * kRows; ++j) mask &= mask - 1ull; // skip the earlier groups' tokens + uint64_t rows = mask; + for (int j = __popcll(mask) - kRows; j > 0; --j) rows &= ~(1ull << (63 - __clzll(rows))); // at most 16 + if (lane == 0) { + s.item_e[ni] = e; + s.item_mask[ni] = rows; + } + ++ni; + }; + if (owner) { + add(xc, 0); + } else if (dc < O) { + int e, grp; + nth_unit(dc, e, grp); + add(e, grp); + } + for (int o = D + slot; o < O; o += G) { + int e, grp; + nth_unit(o, e, grp); + add(e, grp); + } + __syncwarp(); + const int nall = ni; + if (lane == 0) { + s.nitems = nall; + s.spec_ok = owner ? 1 : 0; + stamp(p, 1); + if (DIAG_ROUTE && p.trace) { // DIAG values (not times): 22 = items, 23 = speculation kept + p.trace[blockIdx.x * kTraceStride + 22] = nall; + p.trace[blockIdx.x * kTraceStride + 23] = owner ? 1 : 0; + } + mbar_arrive(&s.route_done); + } + } + named_bar(4, 224); + } + const int nitems = s.nitems; + if (nitems == 0) { // a duplicate slot with no work: let the speculative loads land, then leave + if (warp == 1) { + if (lane == 0) + for (int b = 0; b < SPEC_STAGES; ++b) mbar_wait(&s.full[b], 0); + __syncwarp(); + } + named_bar(7, kThreads); + if (warp == 2) tmem_dealloc(tm, 512); + return; + } + tc_fence_after(); + const bool spec_ok = s.spec_ok != 0; + // ring stage of item 0's stage 0: the speculative fill (stages 0 .. SPEC_STAGES - 1) is item 0's own only if spec_ok + const int base0 = spec_ok ? 0 : SPEC_STAGES; + float ag = 0.f, sc = 0.f, a2 = 0.f; + auto load_scales = [&](int it) { + const int e = s.item_e[it]; + ag = p.g1_alpha[e]; + sc = p.g1_scale_c[e]; + a2 = p.g2_alpha[e]; + }; + // FC1 drain of quarter-0 accumulator columns [64 c0, 64 c0 + 128) (TMEM lanes 0..15 hold the item rows) into acc1 + auto drain_fc1 = [&](int c0) { +#pragma unroll 1 + for (int c = c0; c < c0 + 2; ++c) { + uint32_t r[64]; + tmem_ld_x64(tm + (uint32_t)(c * 64), r); + tmem_ld_wait(); + if (lane < 16) { +#pragma unroll + for (int i = 0; i < 16; ++i) + *reinterpret_cast(s.acc1 + lane * kARow + c * 64 + 4 * i) = + make_uint4(r[4 * i], r[4 * i + 1], r[4 * i + 2], r[4 * i + 3]); + } + } + tc_fence_before(); + __syncwarp(); + }; + // SwiGLU + NVFP4 epilogue (h = 0: the one FC1 pass), tasks (row, 16-wide block, half) = 16 x 8 x 2 = 256, + // task = et + k * nt. + // Accumulator columns of each 32-block: [0, 8) up 2m, [8, 16) gate 2m, [16, 24) up 2m + 1, [24, 32) gate 2m + 1 + auto epi_tasks = [&](int h, int et, int nt) { +#pragma unroll 1 + for (int task = et; task < 256; task += nt) { + const int row = task / 16, blk = (task / 2) % 8, hf = task % 2; + const float4* ar = reinterpret_cast(s.acc1 + row * kARow + 32 * blk + 4 * hf); + const float4 u_even = ar[0], g_even = ar[2], u_odd = ar[4], g_odd = ar[6]; + const float ga[8] = {g_even.x, g_odd.x, g_even.y, g_odd.y, g_even.z, g_odd.z, g_even.w, g_odd.w}; + const float ua[8] = {u_even.x, u_odd.x, u_even.y, u_odd.y, u_even.z, u_odd.z, u_even.w, u_odd.w}; + // trtllm-gen's SwiGLU + NVFP4 epilogue op for op (from its SASS); a = FC1 dequant, c = a x FC2 input scale: + // h' = (u a)(g sigmoid(g a)), sigmoid = rcp(ex2(-(g (a log2e))) + 1), block scale e4m3(max|h'| c / 6), FP4 (h' / scale) c + const float a_l2e = mul_ftz(ag, 1.4426950216293334961f); + float hv[8]; + float amax = 0.f; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const float ex = ex2_approx(-(ga[i] * a_l2e)); + const float sg = rcp_approx(ex + 1.0f); + const float v = fmaf(ua[i], ag, 0.0f) * (ga[i] * sg); + hv[i] = v; + amax = fmaxf(amax, fabsf(v)); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + const uint32_t sf = cvt_e4m3(fmaxf(0.f, fabsf(amax * sc)) * 0.16666667163372039795f); + const float sff = e4m3_to_f32(sf); + const float osc = sff == 0.f ? 0.f : rcp_approx(sff); + float qv[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) qv[i] = mul_ftz(hv[i], osc) * sc; + const uint32_t packed = cvt_e2m1x8(qv); + const int gblk = 8 * h + blk; // 16-wide block over the 128 intermediates; k-blocks 0, 1 + const int chunk = gblk / 2, r8 = row % 8; + const uint32_t off = (uint32_t)((row / 8) * 1024 + r8 * 128 + ((chunk ^ r8) * 16) + (gblk % 2) * 8 + hf * 4); +#pragma unroll + for (int cp = 0; cp < 8; ++cp) *reinterpret_cast(s.ha + cp * 2048 + off) = packed; + if (hf == 0) reinterpret_cast(&s.hsf[row][gblk / 4])[gblk % 4] = (uint8_t)sf; + } + }; + // Rows of item it (its tokens in increasing order, at most 16) and their x rows / scale bytes, by warps 0 and 3; + // row info is double-buffered by item parity (the previous item's FC2 drains still read theirs). + auto gather = [&](int it) { + const int rb = it & 1; + if (warp == 0) set_rows(rb, s.item_mask[it], s.item_e[it]); + named_bar(5, 64); + copy_rows(rb, (warp == 0 ? 0 : 32) + lane, 64); + named_bar(5, 64); + if (warp == 0 && lane == 0) { + mbar_arrive(&s.x_full); + if (DIAG_ROUTE && it == 0) stamp(p, 14); + } + }; + + if (warp == 1) { + if (lane == 0) { + // ---------------- TMA producer: every item's weight stream through the ring + const uint64_t pol = policy_evict_first(); + for (int it = 0; it < nitems; ++it) { + const int e = s.item_e[it], base = base0 + it * kTotStages; + const int st0 = (it == 0 && spec_ok) ? SPEC_STAGES : 0; + for (int st = st0; st < kTotStages; ++st) { + const int g = base + st; + if (FC2_PF > 0 && st == kFc1Stages - FC2_PF_LEAD) { + // the FC2 tiles beyond the ring go to L2 now, while the last FC1 stages and h are in progress + for (int j = kStages; j < kStages + FC2_PF; ++j) prefetch_stage(kFc1Stages + j, e); + } + if (PF_SET != 0 && ((st == st0 && (PF_SET & 1)) || (st == kFc1Stages + 4 && (PF_SET & 4)) || + (st == kFc1Stages + 6 && (PF_SET & 8)))) { + // these stages wait for MMAs that are slow to start after a phase change; let HBM fetch them meanwhile + const int n = st == st0 ? PF_N1 : st == kFc1Stages + 4 ? PF_N2 : 4; + for (int k = st; k < st + n && k < kTotStages; ++k) prefetch_stage(k, e); + } + if (g >= kStages) mbar_wait(&s.empty[g % kStages], ((g / kStages) - 1) & 1); + issue_stage(g, st, e, pol); + if (STAGE_STAMPS && it == 0) stamp(p, 64 + st); + } + } + } + __syncwarp(); + } else if (warp == 2) { + if (lane == 0) { + // ---------------- MMA issuer (tensor pipe: MMAs only) + const uint32_t idesc = make_idesc_nvf4(128u, 256u); + if (!spec_ok) { // the speculative fill was another CTA's expert: recycle its ring stages without MMAs + for (int g = 0; g < SPEC_STAGES; ++g) { + mbar_wait(&s.full[g % kStages], (g / kStages) & 1); + mbar_wait(&s.sf_full[g % kStages], (g / kStages) & 1); + mbar_arrive(&s.empty[g % kStages]); + } + } + for (int it = 0; it < nitems; ++it) { + const int base = base0 + it * kTotStages, par = it & 1; + if (it > 0) mbar_wait(&s.acc_drained, (it - 1) & 1); + mbar_wait(&s.x_full, par); + tc_fence_after(); + int n = 0; + for (int st = 0; st < kFc1Stages; ++st) { + const int g = base + st, b = g % kStages, ks = st; + mbar_wait(&s.full[b], (g / kStages) & 1); + mbar_wait(&s.sf_full[b], (g / kStages) & 1); + tc_fence_after(); + if (st == 0 && it == 0) stamp(p, 3); + if (st == 0 && it > 0 && it == nitems - 1) stamp(p, 11); + if (DIAG_ROUTE && st == 0 && it < 4) stamp(p, 24 + it); // DIAG: item it's first FC1 MMA + if (DIAG_ROUTE && st == kFc1Stages - 1 && it < 4) stamp(p, 28 + it); // DIAG: its last FC1 stage landed +#pragma unroll + for (int q = 0; q < 4; ++q, ++n) { + const int kb = ks * 4 + q; // k-blocks in K order: one accumulator, the stock summation order + umma_nvf4(tm, make_sdesc(smem_u32(s.xa + (kb / 4) * 2048) + (kb % 4) * 32, 0, 1024, 2), + make_sdesc(smem_u32(s.ring[b]) + q * 32, 0, 1024, 2), idesc, + tm + kColSfaX + (uint32_t)(kb % kSfaSlots) * 4, tm + kColSfb + (uint32_t)(n % kSfbSlots) * 8, + kb > 0); + } + tc_commit(&s.empty[b]); + } + tc_commit(&s.fc1_done); + mbar_wait(&s.h_ready, par); + tc_fence_after(); + if (it == 0) stamp(p, 6); + for (int st = kFc1Stages; st < kTotStages; ++st) { + const int g = base + st, b = g % kStages, f = st - kFc1Stages; + if (f == 4) { mbar_wait(&s.fc2_free, par); tc_fence_after(); } + mbar_wait(&s.full[b], (g / kStages) & 1); + mbar_wait(&s.sf_full[b], (g / kStages) & 1); + tc_fence_after(); +#pragma unroll + for (int tl = 0; tl < 2; ++tl) { + const int j = 2 * f + tl, jt = j % 8; +#pragma unroll + for (int kb = 0; kb < 2; ++kb) { + const int n2 = 64 + j * 2 + kb; // scale-factor slot, numbered on from FC1's 64 k-blocks + umma_nvf4(tm, make_sdesc(smem_u32(s.ha) + kb * 32, 0, 1024, 2), + make_sdesc(smem_u32(s.ring[b] + tl * 16384) + kb * 32, 0, 512, 4), idesc, + tm + kColSfa + (uint32_t)(jt * 2 + kb) * 4, tm + kColSfb + (uint32_t)(n2 % kSfbSlots) * 8, + (jt | kb) != 0); + } + } + tc_commit(&s.empty[b]); + if (f % 4 == 3) tc_commit(&s.fc2_pass[f / 4]); + } + } + } + __syncwarp(); + } else if (warp == 0 || warp == 3) { + // ---------------- x gather, FC1 drains (warp 0: columns [0, 128); warp 4 drains [128, 256)) and the epilogue + const int et = (warp == 0 ? 0 : 32) + lane; // 0..63 + gather(0); + for (int it = 0; it < nitems; ++it) { + const int par = it & 1; + load_scales(it); + if (warp == 0) { + mbar_wait(&s.fc1_done, par); + tc_fence_after(); + if (lane == 0 && it == 0) stamp(p, 4); + drain_fc1(0); + if (lane == 0) { + mbar_arrive(&s.acc_ready); + if (it == 0 && !DIAG_ROUTE) stamp(p, 13); + } + } + mbar_wait(&s.acc_ready, par); + epi_tasks(0, et, 192); + if (warp == 3 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 14); + fence_proxy_async_smem(); + named_bar(6, 192); // h and its scale words are complete for the TMEM warpgroup + if (it + 1 < nitems) gather(it + 1); // FC1 of item it is done: x and its scale bytes are free + } + } else { + // ---------------- TMEM-side warpgroup (warps 4..7): q = lane quarter (TMEM lanes 32 q .. 32 q + 31) + const int q = warp - 4; + const uint32_t lane_addr = (uint32_t)(q * 32) << 16; + if (!spec_ok) { // the speculative fill is recycled without MMAs: release its scale-factor side + for (int g = 0; g < SPEC_STAGES; ++g) { + mbar_wait(&s.full[g % kStages], (g / kStages) & 1); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[g % kStages]); + } + } + for (int it = 0; it < nitems; ++it) { + const int base = base0 + it * kTotStages, par = it & 1, rb = it & 1; + load_scales(it); + mbar_wait(&s.x_full, par); + const int nrows = s.nrows[rb]; + auto stage_sf2 = [&](int st) { // FC2 tiles 2 f, 2 f + 1: k-blocks 0, 1 of their row blocks (rows 0..127 | 128..255) + const int g = base + st, b = g % kStages; + mbar_wait(&s.full[b], (g / kStages) & 1); + if (STAGE_STAMPS && threadIdx.x == 128 && it == 0) stamp(p, 16 + st); + if (g >= 4) mbar_wait(&s.empty[(g - 4) % kStages], ((g - 4) / kStages) & 1); + const uint8_t* sf = s.ring[b] + 32768; + const int n0 = 64 + (st - kFc1Stages) * 4; + for (int i = 0; i < 4; ++i) { + const int tl = i / 2, kb = i % 2; + const uint4 v0 = *reinterpret_cast(sf + tl * 2048 + kb * 512 + lane * 16); + const uint4 v1 = *reinterpret_cast(sf + tl * 2048 + 1024 + kb * 512 + lane * 16); + const uint32_t r[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; + tmem_st_x8(tm + lane_addr + kColSfb + (uint32_t)((n0 + i) % kSfbSlots) * 8, r); + } + tmem_st_wait(); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[b]); + }; + // FC1 stage st: weight scale factors of k-blocks [4 ks, 4 ks + 4) (8 columns each: row blocks 2 h | 2 h + 1) and + // the x scale factors of the same k-blocks (row m -> lane m % 32, column m / 32 of each 4-column atom, every + // quarter holds a copy; rows >= nrows, including the A rows read past the 16 real ones, get zero scales) + for (int st = 0; st < kFc1Stages; ++st) { + const int g = base + st, b = g % kStages, ks = st; + mbar_wait(&s.full[b], (g / kStages) & 1); + if (STAGE_STAMPS && threadIdx.x == 128 && it == 0) stamp(p, 16 + st); + if (g >= 4) mbar_wait(&s.empty[(g - 4) % kStages], ((g - 4) / kStages) & 1); + const uint8_t* sf = s.ring[b] + 32768; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const uint4 v0 = *reinterpret_cast(sf + i * 512 + lane * 16); + const uint4 v1 = *reinterpret_cast(sf + 2048 + i * 512 + lane * 16); + const uint32_t r[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; + tmem_st_x8(tm + lane_addr + kColSfb + (uint32_t)((st * 4 + i) % kSfbSlots) * 8, r); + } + uint32_t xr[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int kb = ks * 4 + i / 4; + xr[i] = (i % 4 == 0 && lane < nrows) ? *reinterpret_cast(&s.xsf[lane][kb * 4]) : 0u; + } + tmem_st_x16(tm + lane_addr + kColSfaX + (uint32_t)((ks * 4) % kSfaSlots) * 4, xr); + tmem_st_wait(); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[b]); + } + // EARLY_SF2: stage the first FC2 scale factors before h (they wait for FC2 data to land, which delays h) + for (int st = kFc1Stages; st < kFc1Stages + EARLY_SF2; ++st) stage_sf2(st); + if (q == 0) { // drain the upper 128 columns + mbar_wait(&s.fc1_done, par); + tc_fence_after(); + drain_fc1(2); + if (lane == 0) mbar_arrive(&s.acc_ready); + } + mbar_wait(&s.acc_ready, par); + epi_tasks(0, 64 + (warp - 4) * 32 + lane, 192); + fence_proxy_async_smem(); + named_bar(6, 192); // h complete + if (threadIdx.x == 128 && it == 0) stamp(p, 5); + // h scale factors: 16 atoms (tile jt, k-block kb) of 4 columns; row m = 32 c + lane is nonzero only in its tile + { + uint32_t r[32]; +#pragma unroll 1 + for (int half = 0; half < 2; ++half) { +#pragma unroll + for (int i = 0; i < 32; ++i) { + const int a = half * 8 + i / 4, col = i % 4; // atom a = 2 jt + kb + const int jt = a / 2, kb = a % 2; + const int row = 32 * col + lane; + r[i] = (row / 16 == jt && row % 16 < nrows) ? s.hsf[row % 16][kb] : 0u; + } + tmem_st_x32(tm + lane_addr + kColSfa + half * 32, r); + } + tmem_st_wait(); + } + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.h_ready); + if (q == 0 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 15); + for (int st = kFc1Stages + EARLY_SF2; st < kTotStages; ++st) { + if (st == kFc1Stages + 4) { + // pass 0 is fully staged; drain it once its MMAs are done, then free the accumulator for pass 1 + mbar_wait(&s.fc2_pass[0], par); + tc_fence_after(); + if (q == 0 && lane == 0 && it == 0) stamp(p, 7); + fc2_drain16<0, 4>(p, s, tm, q, lane, 0, a2, rb); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.fc2_free); + } + stage_sf2(st); + } + mbar_wait(&s.fc2_pass[1], par); + tc_fence_after(); + if (q == 0 && lane == 0 && it == 0) stamp(p, 12); + if (TAIL8 && it + 1 == nitems) fc2_drain16<0, 2>(p, s, tm, q, lane, 1, a2, rb); + else fc2_drain16<0, 4>(p, s, tm, q, lane, 1, a2, rb); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.acc_drained); + } + } + if (TAIL8 && warp < 4) { // the last item's last pass, upper columns: every warp is free now (TMEM lane quarter = warp) + const int it = nitems - 1; + load_scales(it); + mbar_wait(&s.fc2_pass[1], it & 1); + tc_fence_after(); + fc2_drain16<2, 4>(p, s, tm, warp, lane, 1, a2, it & 1); + } + tc_fence_before(); + named_bar(7, kThreads); // every warp, including the drains, before TMEM is released + if (threadIdx.x == 128) stamp(p, 8); + if (warp == 2) tmem_dealloc(tm, 512); +} + +static_assert(sizeof(Smem) + 1024 == Q35_MKM_SMEM_BYTES, "the launcher's dynamic shared memory size"); + +} // namespace qwen_mkm_moe_m64 + + +} // namespace sglang diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_wide.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_wide.cu new file mode 100644 index 00000000..de1a7d2a --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_wide.cu @@ -0,0 +1,1331 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Fused NVFP4 MoE expert phase for Qwen3.5 TP8 decode verify at 129 to Q35_MKM_MAX_TOK (256 or 448) tokens, top-10, +// 512 experts, hidden 4096, intermediate 128 per rank, on a given routing (route_topk's ids): FC1, SwiGLU, the NVFP4 +// requantization and FC2 in one launch over the trtllm-gen weight layout. Writes the deferred-finalize rows (row +// t * 10 + k, identity permutation). +// +// The 128-token engine's (mkm_moe_m128.cu) units, (expert, row group of at most 16 of its tokens), each with its +// arithmetic and summation order: rows bit for bit trtllm-gen's. At these widths nearly every expert has tokens, so the +// work split is static: CTA c runs experts c, c + G, c + 2 G, ... (those with tokens), and streams expert c's weights +// before the grid dependency wait (the weights do not depend on the preceding kernel). The further row groups of +// experts with more than 16 tokens go, in expert order, to CTAs G - 1, G - 2, ... (the CTAs with one expert fewer +// first). Every CTA derives the same split from the per-expert token counts. +// +// Built with NVRTC: the kernel has C linkage (q35_mkm_moe_wide), and the launcher's dynamic shared memory size +// (Q35_MKM_SMEM_BYTES) is checked against the layout at compile time. +#include +#include +#include + +// NVRTC has no , or +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef int int32_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef unsigned long long uintptr_t; +typedef struct __align__(64) { + uint64_t opaque[16]; +} CUtensorMap; +#ifndef INFINITY +#define INFINITY __int_as_float(0x7f800000) +#endif +// Minimal SM100 (tcgen05 / TMA / mbarrier) PTX helpers for the small-M NVFP4 MoE engine. + +namespace qwen35 { + +__device__ __forceinline__ uint32_t smem_u32(const void* p) { + return static_cast(__cvta_generic_to_shared(p)); +} + +// ---------------------------------------------------------------- mbarrier +__device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(smem_u32(bar)), "r"(count) : "memory"); +} +__device__ __forceinline__ void fence_mbar_init() { + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} +__device__ __forceinline__ void mbar_expect_tx(uint64_t* bar, uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smem_u32(bar)), "r"(bytes) : "memory"); +} +__device__ __forceinline__ void mbar_arrive(uint64_t* bar) { + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(smem_u32(bar)) : "memory"); +} +__device__ __forceinline__ bool mbar_try_wait(uint64_t* bar, uint32_t phase) { + uint32_t ok; + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ok) + : "r"(smem_u32(bar)), "r"(phase) + : "memory"); + return ok != 0; +} +__device__ __forceinline__ void mbar_wait(uint64_t* bar, uint32_t phase) { + while (!mbar_try_wait(bar, phase)) { + } +} + +// ---------------------------------------------------------------- TMA / bulk copies +__device__ __forceinline__ void tma_load_2d(void* dst, const void* tmap, uint64_t* bar, int c0, int c1) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes [%0], [%1, {%3, %4}], [%2];" :: + "r"(smem_u32(dst)), + "l"(tmap), "r"(smem_u32(bar)), "r"(c0), "r"(c1) + : "memory"); +} +__device__ __forceinline__ void bulk_load(void* dst, const void* src, uint32_t bytes, uint64_t* bar) { + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)) + : "memory"); +} +__device__ __forceinline__ void prefetch_tmap(const void* tmap) { + asm volatile("prefetch.tensormap [%0];" ::"l"(tmap) : "memory"); +} +// generic-proxy smem writes -> visible to the async proxy (UMMA / UTCCP operand reads) +__device__ __forceinline__ void fence_proxy_async_smem() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + +// ---------------------------------------------------------------- tcgen05 +__device__ __forceinline__ void tmem_alloc(uint32_t* dst_smem, uint32_t ncols) { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smem_u32(dst_smem)), + "r"(ncols) + : "memory"); +} +__device__ __forceinline__ void tmem_relinquish() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;" ::: "memory"); +} +__device__ __forceinline__ void tmem_dealloc(uint32_t taddr, uint32_t ncols) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(ncols) : "memory"); +} +__device__ __forceinline__ void tc_fence_before() { + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); +} +__device__ __forceinline__ void tc_fence_after() { + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); +} +// arrive on `bar` when all prior tcgen05 async ops issued by this thread complete +__device__ __forceinline__ void tc_commit(uint64_t* bar) { + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" ::"r"(smem_u32(bar)) + : "memory"); +} +// smem -> TMEM copy of one 512-byte scale-factor chunk (32 rows x 16 B), broadcast to the 4 lane quarters +__device__ __forceinline__ void utccp_32x128b_warpx4(uint32_t taddr, uint64_t sdesc) { + asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" ::"r"(taddr), "l"(sdesc) : "memory"); +} +// D(tmem) (+)= A(smem) * B(smem)^T with NVFP4 block-16 scales from TMEM +__device__ __forceinline__ void umma_nvf4(uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, + uint32_t tmem_sfa, uint32_t tmem_sfb, uint32_t accumulate) { + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %4, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%5], [%6], p;\n}\n" ::"r"( + tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(accumulate), "r"(tmem_sfa), "r"(tmem_sfb) + : "memory"); +} +// each warp reads its 32-lane quarter: 32 lanes x 16 consecutive 32-bit columns +__device__ __forceinline__ void tmem_ld_32x32b_x16(uint32_t taddr, uint32_t (&r)[16]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15}, [%16];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(taddr) + : "memory"); +} +// L2 prefetch of a contiguous global range (bytes multiple of 16, 16-byte aligned) +__device__ __forceinline__ void bulk_prefetch_l2(const void* gptr, uint32_t bytes) { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(gptr), "r"(bytes) : "memory"); +} +__device__ __forceinline__ void tmem_ld_wait() { asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); } +// each warp writes its 32-lane quarter: 32 lanes x 16 consecutive 32-bit columns +__device__ __forceinline__ void tmem_st_32x32b_x16(uint32_t taddr, const uint32_t (&r)[16]) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16};" ::"r"( + taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_wait() { asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); } + +// ---------------------------------------------------------------- descriptors +// SM100 shared-memory matrix descriptor (version 1). layout: 0 none, 2 SW128, 4 SW64, 6 SW32. +__device__ __forceinline__ uint64_t make_sdesc(uint32_t saddr, uint32_t lbo, uint32_t sbo, uint32_t layout) { + uint64_t d = 0; + d |= (uint64_t)((saddr >> 4) & 0x3FFF); + d |= (uint64_t)((lbo >> 4) & 0x3FFF) << 16; + d |= (uint64_t)((sbo >> 4) & 0x3FFF) << 32; + d |= (uint64_t)1 << 46; // version + d |= (uint64_t)(layout & 0x7) << 61; + return d; +} +// Block-scaled instruction descriptor: E2M1 x E2M1, UE4M3 scales, both K-major, F32 accumulate. +__host__ __device__ constexpr uint32_t make_idesc_nvf4(uint32_t M, uint32_t N) { + return (1u << 7) // a_format E2M1 + | (1u << 10) // b_format E2M1 + | ((N >> 3) << 17) // n_dim + | (0u << 23) // scale format UE4M3 + | ((M >> 4) << 24); // m_dim +} + +} // namespace qwen35 + +namespace sglang { + +namespace qwen_mkm_moe_wide { + +using namespace qwen35; + + +#ifndef Q35_MKM_MAX_TOK +#define Q35_MKM_MAX_TOK 256 +#endif +constexpr int kH = 4096, kE = 512, kTopK = 10, kMaxTok = Q35_MKM_MAX_TOK; +static_assert(kMaxTok % 32 == 0 && kMaxTok <= 512, "tokens in whole 32-token groups"); +constexpr int kGroups = kMaxTok / 32; // 32-token groups: lane l holds tokens l + 32 h +constexpr int kRows = 16; // A rows (tokens) per item +// per CTA (G >= 128, the launcher's bound): at most 4 own experts (512 / 128) and ceil(kMaxTok * 10 / 16 / 128) <= 3 +// further row groups +constexpr int kMaxItems = 8; +constexpr int kI = 128; // intermediates per expert per rank (TP8) +constexpr int kN1 = 2 * kI; // W13 rows per expert (gate + up) +constexpr int kStageBytes = 36864; // 32 KB weights + 4 KB scale factors +constexpr int kStages = 4; +constexpr int kW13BoxRows = 128; // rows of one W13 TMA box; the host tensor map must use the same +constexpr int kFc1Stages = 16; // the expert's 256 W13 rows, 4 k-blocks each (full K = 4096) +constexpr int kFc2Stages = 8; // two tiles of 256 W2 rows x 64 B each +constexpr int kTotStages = kFc1Stages + kFc2Stages; +constexpr int kThreads = 256; +// TMEM columns. FC1: accumulator [0,256), x SFA ring [256,320) (16 slots x 4), SFB ring [384,512) (16 slots x 8). +// FC2: accumulator [0,256), h SFA [256,320) (16 atoms x 4), SFB ring [384,512) (16 slots x 8). +constexpr uint32_t kColSfaX = 256, kColSfa = 256, kColSfb = 384; +constexpr int kSfbSlots = 16, kSfaSlots = 16; +constexpr int kARow = 260; // FC1 accumulator row stride in floats: 16 B aligned, rows on distinct banks +#ifndef L2_EVICT_FIRST // stream expert weights with an L2 evict-first policy (they are read once per layer) +#define L2_EVICT_FIRST 1 +#endif +#ifndef TOPK_RANK +#define TOPK_RANK 0 +#endif +#ifndef SF_UTCCP_FC1 +#define SF_UTCCP_FC1 0 +#endif +#ifndef SPEC_STAGES // ring stages issued for the speculated expert before the routing verdict (rest: after it) +#define SPEC_STAGES 4 +#endif +#ifndef STAGE_STAMPS // DIAG: per-stage issue (slots 64 + st) and landing (16 + st, seen by warp 4) stamps +#define STAGE_STAMPS 0 +#endif +#ifndef DIAG_SPEC // DIAG: warp 1's speculative ranking, stamps 16 = logits loaded, 17 = keys, 18 = top-10 (stride 32) +#define DIAG_SPEC 0 +#endif +#ifndef DIAG_ROUTE // DIAG: stamps 13 = warp 0's routing rounds done, 15 = their barrier, 14 = item 0's rows gathered +#define DIAG_ROUTE 0 +#endif +constexpr int kTraceStride = STAGE_STAMPS ? 128 : (DIAG_SPEC || DIAG_ROUTE) ? 32 : 16; +#ifndef SMEM_SPACE // 1: shared-memory accesses through Smem compile to shared-space instructions +#define SMEM_SPACE 1 +#endif +#ifndef PF_SET // L2 prefetch of the stages that wait for a ring slot at: 1 = start, 4 = FC2, 8 = FC2 pass 1 +// (none: from 56 rows every CTA streams and HBM is the bound, so these prefetches only compete with the ring's loads) +#define PF_SET 0 +#endif +#ifndef PF_N1 // stages prefetched at the start trigger +#define PF_N1 4 +#endif +#ifndef PF_N2 // stages prefetched at the FC2 trigger +#define PF_N2 4 +#endif +#ifndef TAILV // FC2 drains: bit 0 = last pass on all 8 warps; bit 1 = TMEM loads double-buffered; bit 2 = coalesced stores +#define TAILV 3 +#endif +#define TAIL8 (TAILV & 1) +#define DRAIN_PIPE ((TAILV >> 1) & 1) +#define DRAIN_SMEM ((TAILV >> 2) & 1) +static_assert(DRAIN_SMEM == 0, "x's shared memory holds the next item's rows while the drains run"); +#ifndef EARLY_SF2 // FC2 stages whose scale factors are staged before the h epilogue (old order: 2) +#define EARLY_SF2 0 +#endif +#ifndef GATHER_BATCH // x row granules each gathering thread loads before storing them +#define GATHER_BATCH 8 +#endif +#ifndef GATHER_ASYNC // 1: the gather copies by cp.async (every granule in flight at once, no register staging) +#define GATHER_ASYNC 1 +#endif +#ifndef SPEC_PF // stages of the speculated expert past the first ring fill prefetched to L2 while the routing runs +#define SPEC_PF 4 +#endif +#ifndef FC2_PF // FC2 tiles (beyond the first kStages) prefetched to L2 ahead of the ring; 0 = off +#define FC2_PF 0 +#endif +#ifndef FC2_PF_LEAD // issue the FC2 prefetch when the producer reaches FC1 stage kFc1Stages - FC2_PF_LEAD +#define FC2_PF_LEAD 4 +#endif + +struct Params { + const uint16_t* logits; // [M, 512] bf16 router logits (router_gemm_quant output) + uint16_t* expert_weights; // [M, 10] bf16 out: DynBlock softmax-over-top-10 weights (the last CTA writes) + int32_t* topk_ids; // [M, 10] out (the last CTA writes) + const uint8_t* x_sf; // [M, 256] linear e4m3 + const uint8_t* w13_sf; // [E, 2 row blocks x 64 k blocks x 512 B] (128x4 interleaved) + const uint8_t* w2_sf; // [E, 32 row blocks x 2 k blocks x 512 B] + const float* g1_alpha; // [E] gate dequant + const float* g1_scale_c; // [E] up dequant * FC2 input quant scale + const float* g2_alpha; // [E] + __nv_bfloat16* out; // [M * 10, 4096] + long long* trace; // optional [grid, 16] %globaltimer stamps + int num_tokens; + const uint8_t* x_fp4; // [M, 2048] NVFP4 activations (each item gathers its rows) + // [M, 10] the routing's expert ids (route_topk's, bit for bit the trtllm-gen routing): when set, the engine takes them + // in place of ranking the logits (and writes neither expert_weights nor topk_ids, which route_topk wrote) + const int32_t* ids_in; +}; + +struct __align__(1024) Smem { + uint8_t ring[kStages][kStageBytes]; + alignas(1024) uint8_t xa[32 * 1024]; // x: per 128-B K chunk c, 2 KB = SW128 atoms of item rows 0..7 | 8..15 + alignas(1024) uint8_t ha[16 * 1024]; // h: 8 copies of (rows 0..7 | rows 8..15) x 128 B (K = 128: 64 B), SW128 + float acc1[16 * kARow]; // FC1 accumulator rows 0..15 (item rows) x 256 columns + uint8_t xsf[16][256]; // x scale bytes of the item's rows, full K + uint32_t hsf[16][2]; // h scale words: row r, k-block kb (4 e4m3 bytes) + uint64_t full[kStages], empty[kStages], sf_full[kStages]; + uint64_t x_full, fc1_done, acc_ready, h_ready, fc2_pass[2], fc2_free, acc_drained, route_done; + uint32_t tmem; + int nitems, spec_ok; + int item_e[kMaxItems]; + int item_g[kMaxItems]; // the item's row group: its expert's tokens [16 g, 16 g + 16) in token order + int nrows[2], rtok[2][kRows], rkslot[2][kRows]; // rows of the current / next item (by item parity) + int cnt[kE]; // tokens of each expert (atomicAdd) + alignas(16) uint16_t ids[kMaxTok][kTopK]; // the routing (16-B aligned: stored 4 ids at a time) +}; + +__device__ __forceinline__ uint64_t gtime() { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)::"memory"); + return t; +} +__device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} +__device__ __forceinline__ void cluster_arrive() { + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); +} +// no release: the phase only orders DSMEM traffic that mbarriers already completed, so pending global stores +// (stamps, routing outputs, output rows) need not drain before the arrive +__device__ __forceinline__ void cluster_arrive_relaxed() { + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} +__device__ __forceinline__ void cluster_wait() { + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); +} +__device__ __forceinline__ uint32_t mapa(uint32_t saddr, uint32_t rank) { + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(r) : "r"(saddr), "r"(rank)); + return r; +} +// smem (this CTA) -> smem (peer CTA in the cluster), completes tx bytes on the peer's mbarrier +__device__ __forceinline__ void bulk_s2s(uint32_t dst_cluster, uint32_t src, uint32_t bytes, uint32_t bar_cluster) { + asm volatile("cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + dst_cluster), + "r"(src), "r"(bytes), "r"(bar_cluster) + : "memory"); +} +// registers -> the peer CTA's shared memory; completes the bytes on the peer's mbarrier +__device__ __forceinline__ void st_async_v4(uint32_t dst, const uint4& v, uint32_t bar) { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"(dst), + "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "r"(bar) + : "memory"); +} +__device__ __forceinline__ void st_async_b32(uint32_t dst, uint32_t v, uint32_t bar) { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(dst), "r"(v), "r"(bar) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x4(uint32_t taddr, const uint4& v) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x4.b32 [%0], {%1, %2, %3, %4};" ::"r"(taddr), "r"(v.x), "r"(v.y), + "r"(v.z), "r"(v.w) + : "memory"); +} +// bar.sync counts arrivals per warp execution: reconverge first, so a warp whose lane 0 just ran a stamp or an +// mbarrier arrive cannot arrive twice and release the barrier early +__device__ __forceinline__ void named_bar(int id, int n) { + __syncwarp(); + asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(n) : "memory"); +} +__device__ __forceinline__ void tmem_ld_x32(uint32_t taddr, uint32_t (&r)[32]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,%19," + "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), "=r"(r[16]), + "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), "=r"(r[24]), + "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr) + : "memory"); +} +__device__ __forceinline__ void tmem_ld_x64(uint32_t taddr, uint32_t (&r)[64]) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x64.b32 {%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,%19," + "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31,%32,%33,%34,%35,%36,%37,%38,%39,%40,%41,%42,%43,%44,%45,%46,%47," + "%48,%49,%50,%51,%52,%53,%54,%55,%56,%57,%58,%59,%60,%61,%62,%63}, [%64];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), "=r"(r[16]), + "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), "=r"(r[24]), + "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]), "=r"(r[32]), + "=r"(r[33]), "=r"(r[34]), "=r"(r[35]), "=r"(r[36]), "=r"(r[37]), "=r"(r[38]), "=r"(r[39]), "=r"(r[40]), + "=r"(r[41]), "=r"(r[42]), "=r"(r[43]), "=r"(r[44]), "=r"(r[45]), "=r"(r[46]), "=r"(r[47]), "=r"(r[48]), + "=r"(r[49]), "=r"(r[50]), "=r"(r[51]), "=r"(r[52]), "=r"(r[53]), "=r"(r[54]), "=r"(r[55]), "=r"(r[56]), + "=r"(r[57]), "=r"(r[58]), "=r"(r[59]), "=r"(r[60]), "=r"(r[61]), "=r"(r[62]), "=r"(r[63]) + : "r"(taddr) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x8(uint32_t taddr, const uint32_t (&r)[8]) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x8.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" ::"r"(taddr), "r"(r[0]), + "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x16(uint32_t taddr, const uint32_t (&r)[16]) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x16.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16};" :: + "r"(taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]) + : "memory"); +} +__device__ __forceinline__ void tmem_st_x32(uint32_t taddr, const uint32_t (&r)[32]) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x32.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18," + "%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31,%32};" ::"r"(taddr), + "r"(r[0]), "r"(r[1]), "r"(r[2]), "r"(r[3]), "r"(r[4]), "r"(r[5]), "r"(r[6]), "r"(r[7]), "r"(r[8]), "r"(r[9]), + "r"(r[10]), "r"(r[11]), "r"(r[12]), "r"(r[13]), "r"(r[14]), "r"(r[15]), "r"(r[16]), "r"(r[17]), "r"(r[18]), + "r"(r[19]), "r"(r[20]), "r"(r[21]), "r"(r[22]), "r"(r[23]), "r"(r[24]), "r"(r[25]), "r"(r[26]), "r"(r[27]), + "r"(r[28]), "r"(r[29]), "r"(r[30]), "r"(r[31]) + : "memory"); +} +__device__ __forceinline__ float ex2_approx(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float mul_ftz(float a, float b) { + float y; + asm("mul.ftz.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ float rcp_approx(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ uint32_t cvt_e4m3(float a) { + uint32_t out; + asm("{\n .reg .b16 pair;\n .reg .f32 zero;\n mov.f32 zero, 0f00000000;\n" + " cvt.rn.satfinite.e4m3x2.f32 pair, zero, %1;\n cvt.u32.u16 %0, pair;\n}" + : "=r"(out) + : "f"(a)); + return out & 0xFFu; +} +__device__ __forceinline__ float e4m3_to_f32(uint32_t v) { + float out; + asm("{\n .reg .b16 pair;\n .reg .b32 h2;\n .reg .b16 lo, hi;\n cvt.u16.u32 pair, %1;\n" + " cvt.rn.f16x2.e4m3x2 h2, pair;\n mov.b32 {lo, hi}, h2;\n cvt.f32.f16 %0, lo;\n}" + : "=f"(out) + : "r"(v)); + return out; +} +__device__ __forceinline__ uint32_t cvt_e2m1x8(const float* v) { + uint32_t out; + asm("{\n .reg .b8 b0, b1, b2, b3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\n cvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\n cvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" + " mov.b32 %0, {b0, b1, b2, b3};\n}" + : "=r"(out) + : "f"(v[0]), "f"(v[1]), "f"(v[2]), "f"(v[3]), "f"(v[4]), "f"(v[5]), "f"(v[6]), "f"(v[7])); + return out; +} +// ---- router top-k, bit-exact with FlashInfer's single-block routing (routingCustom, RenormalizeNaive on bf16): +// key = bf16 order (cub float twiddle) above 65535 - expert, i.e. the higher score first, the smaller expert on ties. +__device__ __forceinline__ uint32_t topk_key(uint32_t bits, uint32_t expert) { + const uint32_t twiddled = bits ^ ((bits & 0x8000u) ? 0xFFFFu : 0x8000u); + return (twiddled << 16) | (65535u - expert); +} +__device__ __forceinline__ void topk_unkey(uint32_t key, uint32_t& bits, uint32_t& expert) { + expert = 65535u - (key & 0xFFFFu); + const uint32_t twiddled = key >> 16; + bits = twiddled ^ ((twiddled & 0x8000u) ? 0x8000u : 0xFFFFu); +} +__device__ __forceinline__ uint32_t max16(const uint32_t (&k)[16]) { + const uint32_t a = max(max(max(k[0], k[1]), max(k[2], k[3])), max(max(k[4], k[5]), max(k[6], k[7]))); + const uint32_t b = max(max(max(k[8], k[9]), max(k[10], k[11])), max(max(k[12], k[13]), max(k[14], k[15]))); + return max(a, b); +} +__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_sum_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.add.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_topk_of_lanes(uint32_t v, uint32_t lane) { + uint32_t mine = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(v); + v = v == top ? 0u : v; + mine = lane == k ? top : mine; + } + return mine; +} +__device__ __forceinline__ uint32_t count_greater(const uint32_t* s_cand, uint32_t v) { + uint32_t r = 0; +#pragma unroll + for (int j = 0; j < 32; j += 4) { + const uint4 q = *reinterpret_cast(s_cand + j); + r += (q.x > v) + (q.y > v) + (q.z > v) + (q.w > v); + } + return r; +} +// Top-10 keys of 512 (16 per lane), lane k < 10 holding the k-th largest (keys unique and nonzero). Ranks come from +// all-pairs compares over shared memory rather than chains of warp reductions. +__device__ __forceinline__ uint32_t warp_top10(uint32_t (&keys)[16], uint32_t lane, uint32_t* s_cand) { + // bound = the 10th largest lane maximum; every top-10 key is >= bound + const uint32_t m = max16(keys); + s_cand[lane] = m; + __syncwarp(); + const uint32_t r = count_greater(s_cand, m); + const uint32_t bound = __shfl_sync(0xffffffffu, m, __ffs(__ballot_sync(0xffffffffu, r == kTopK - 1)) - 1); + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) count += keys[i] >= bound; + // exclusive prefix of count over the lanes, from the ballots of its bits (count <= 16) + const uint32_t lt = (1u << lane) - 1u; + uint32_t offset = 0, total = 0; +#pragma unroll + for (int b = 0; b < 5; ++b) { + const uint32_t bb = __ballot_sync(0xffffffffu, (count >> b) & 1u); + offset += (uint32_t)__popc(bb & lt) << b; + total += (uint32_t)__popc(bb) << b; + } + if (total <= 32) { + __syncwarp(); // every lane has read the lane maxima +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) { + if (keys[i] >= bound) s_cand[offset++] = keys[i]; + } + if (lane >= total) s_cand[lane] = 0u; + __syncwarp(); + const uint32_t cand = s_cand[lane]; + const uint32_t rk = count_greater(s_cand, cand); + __syncwarp(); + if (lane < total && rk < kTopK) s_cand[rk] = cand; + __syncwarp(); + return lane < kTopK ? s_cand[lane] : 0u; + } + uint32_t mine = 0; +#pragma unroll 1 + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(max16(keys)); +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) keys[i] = keys[i] == top ? 0u : keys[i]; + mine = lane == k ? top : mine; + } + return mine; +} +// Top-10 keys of 512 (16 per lane), lane k holding the k-th largest (keys unique and nonzero). +__device__ __forceinline__ uint32_t warp_top10_redux(uint32_t (&keys)[16], uint32_t lane, uint32_t* s_cand) { + uint32_t v = max16(keys); + uint32_t bound = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + bound = warp_max_u32(v); + v = v == bound ? 0u : v; + } + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) count += keys[i] >= bound; + const uint32_t total = warp_sum_u32(count); + if (total <= 32) { + uint32_t offset = count; +#pragma unroll + for (uint32_t d = 1; d < 32; d <<= 1) { + const uint32_t up = __shfl_up_sync(0xffffffffu, offset, d); + offset += lane >= d ? up : 0u; + } + offset -= count; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) { + if (keys[i] >= bound) s_cand[offset++] = keys[i]; + } + __syncwarp(); + const uint32_t cand = lane < total ? s_cand[lane] : 0u; + __syncwarp(); + return warp_topk_of_lanes(cand, lane); + } + uint32_t mine = 0; +#pragma unroll 1 + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(max16(keys)); +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) keys[i] = keys[i] == top ? 0u : keys[i]; + mine = lane == k ? top : mine; + } + return mine; +} +// The top-10 keys of one logits row (16 bf16 a lane in v0, v1), lane k < 10 holding the k-th largest +__device__ __forceinline__ uint32_t row_top10(const uint4 v0, const uint4 v1, uint32_t lane, uint32_t* s_cand) { + uint32_t keys[16]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t w = i < 4 ? (&v0.x)[i] : (&v1.x)[i - 4]; + keys[2 * i] = topk_key(w & 0xFFFFu, lane * 16 + 2 * i); + keys[2 * i + 1] = topk_key(w >> 16, lane * 16 + 2 * i + 1); + } + return TOPK_RANK ? warp_top10(keys, lane, s_cand) : warp_top10_redux(keys, lane, s_cand); +} +// FlashInfer calcSoftmax (bf16 scores) as built with -use_fast_math: expf -> ex2.approx.ftz of x * log2(e), +// division -> div.full.ftz; reductions are cooperative_groups::reduce like the reference. +__device__ __forceinline__ float fast_expf(float x) { + float y; + asm("{\n .reg .f32 t;\n mul.ftz.f32 t, %1, 0f3FB8AA3B;\n ex2.approx.ftz.f32 %0, t;\n}" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float fast_div(float a, float b) { + float y; + asm("div.full.ftz.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ uint16_t softmax_topk(uint16_t score_bits, uint32_t lane) { + namespace cg = cooperative_groups; + const auto warp = cg::tiled_partition<32>(cg::this_thread_block()); + const float si = __uint_as_float((uint32_t)score_bits << 16); + float max_score = -INFINITY; + if (lane < kTopK) max_score = si >= max_score ? si : max_score; + max_score = cg::reduce(warp, max_score, cg::greater()); + float sum_score = 0.f, new_score = 0.f; + if (lane < kTopK) { + new_score = fast_expf(__fsub_rn(si, max_score)); + sum_score = __fadd_rn(sum_score, new_score); + } + sum_score = cg::reduce(warp, sum_score, cg::plus()); + uint16_t out = score_bits; + if (lane < kTopK) { + const __nv_bfloat16 b = __float2bfloat16_rn(fast_div(new_score, sum_score)); + out = *reinterpret_cast(&b); + } + return out; +} + +__device__ __forceinline__ uint64_t policy_evict_first() { + uint64_t pol; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(pol)); + return pol; +} +__device__ __forceinline__ void tma_load_2d_w(void* dst, const void* tmap, uint64_t* bar, int c0, int c1, uint64_t pol) { + if (L2_EVICT_FIRST) + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%3, %4}], [%2], %5;" ::"r"(smem_u32(dst)), + "l"(tmap), "r"(smem_u32(bar)), "r"(c0), "r"(c1), "l"(pol) + : "memory"); + else + tma_load_2d(dst, tmap, bar, c0, c1); +} +__device__ __forceinline__ void bulk_load_w(void* dst, const void* src, uint32_t bytes, uint64_t* bar, uint64_t pol) { + if (L2_EVICT_FIRST) + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)), "l"(pol) + : "memory"); + else + bulk_load(dst, src, bytes, bar); +} + +__device__ __forceinline__ void stamp(const Params& p, int slot) { + if (p.trace) p.trace[blockIdx.x * kTraceStride + slot] = (long long)gtime(); +} +// stamp after `dep` is available (the branch needs its value; keys are never 0) +__device__ __forceinline__ void stamp_dep(const Params& p, int slot, uint32_t dep) { + if (dep == 0u) __trap(); + stamp(p, slot); +} +__device__ __forceinline__ void tma_prefetch_2d(const void* tmap, int c0, int c1) { + asm volatile("cp.async.bulk.prefetch.tensor.2d.L2.global.tile [%0, {%1, %2}];" ::"l"(tmap), "r"(c0), "r"(c1) : "memory"); +} +// Physical W13 column (accumulator column) holding reordered row q: block-32 shuffle of FlashInfer's +// get_shuffle_matrix_a_row_indices (epilogue_tile_m = 128); reordered row 2j is up_j, 2j + 1 is gate_j. +__device__ __forceinline__ int w13_col(int q) { return (q & ~31) + ((q & 31) % 4) * 8 + (q & 31) / 4; } +__host__ __device__ constexpr int w13_logical(int p) { + return ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) % 2 ? ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) / 2 + : 128 + ((p & ~31) + (p % 8) * 4 + (p & 31) / 8) / 2; +} + + +// FC2 drain of one pass: TMEM lane quarter q (tiles 2 q, 2 q + 1 of the pass), 256 accumulator columns +__device__ __forceinline__ void reg_fence64(uint32_t (&r)[64]) { // keep uses of r after the tcgen05.wait::ld +#pragma unroll + for (int i = 0; i < 64; ++i) asm volatile("" : "+r"(r[i])); +} + +// FC2 accumulator columns [64 C0, 64 C1) of this warp's lane quarter -> bf16 gemm2_out rows (row r of a tile = the +// r-th row of the item, row info buffer rb) +template +__device__ __forceinline__ void fc2_drain16(const Params& p, Smem& s, uint32_t tm, int q, int lane, int pass, float a2, + int rb) { + const int row = q * 32 + lane, tile = pass * 8 + row / 16, r = row % 16; + const bool valid = r < s.nrows[rb]; + __nv_bfloat16* orow = valid ? p.out + (size_t)(s.rtok[rb][r] * kTopK + s.rkslot[rb][r]) * kH + tile * 256 : nullptr; + const uint32_t base = tm + ((uint32_t)(q * 32) << 16); + auto put = [&](const uint32_t (&r)[64], int c) { + if (valid) { +#pragma unroll + for (int b32 = 0; b32 < 2; ++b32) { + uint32_t pk[16]; +#pragma unroll + for (int l = 0; l < 32; l += 2) { + const float v0 = __uint_as_float(r[b32 * 32 + (l % 4) * 8 + l / 4]) * a2; + const float v1 = __uint_as_float(r[b32 * 32 + ((l + 1) % 4) * 8 + (l + 1) / 4]) * a2; + const __nv_bfloat162 b2 = __floats2bfloat162_rn(v0, v1); + pk[l / 2] = *reinterpret_cast(&b2); + } +#pragma unroll + for (int i = 0; i < 4; ++i) + *reinterpret_cast(orow + c * 64 + b32 * 32 + i * 8) = + make_uint4(pk[4 * i], pk[4 * i + 1], pk[4 * i + 2], pk[4 * i + 3]); + } + } + }; + if (DRAIN_PIPE) { // the next chunk's TMEM load is in flight while this chunk converts and stores + uint32_t ra[64], rb2[64]; + tmem_ld_x64(base + C0 * 64, ra); + tmem_ld_wait(); + reg_fence64(ra); +#pragma unroll + for (int c = C0; c < C1; c += 2) { + if (c + 1 < C1) tmem_ld_x64(base + (uint32_t)((c + 1) * 64), rb2); + put(ra, c); + if (c + 1 < C1) { + tmem_ld_wait(); + reg_fence64(rb2); + if (c + 2 < C1) tmem_ld_x64(base + (uint32_t)((c + 2) * 64), ra); + put(rb2, c + 1); + if (c + 2 < C1) { tmem_ld_wait(); reg_fence64(ra); } + } + } + } else { +#pragma unroll 1 + for (int c = C0; c < C1; ++c) { + uint32_t r[64]; + tmem_ld_x64(base + (uint32_t)(c * 64), r); + tmem_ld_wait(); + put(r, c); + } + } +} + +extern "C" __global__ void __launch_bounds__(kThreads, 1) +q35_mkm_moe_wide(const __grid_constant__ CUtensorMap tm13, const __grid_constant__ CUtensorMap tm2, const Params p) { + extern __shared__ __align__(1024) uint8_t raw[]; +#if SMEM_SPACE // pointer arithmetic on the __shared__ array keeps the shared address space (STS/LDS, not generic) + Smem& s = *reinterpret_cast(raw + ((1024u - (smem_u32(raw) & 1023u)) & 1023u)); +#else + Smem& s = *reinterpret_cast((reinterpret_cast(raw) + 1023) & ~uintptr_t(1023)); +#endif + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + const int cta = blockIdx.x, G = gridDim.x; // G >= 128 (the launcher's bound) + const int M = p.num_tokens; + for (int i = threadIdx.x; i < kE; i += kThreads) s.cnt[i] = 0; + if (threadIdx.x == 0) { + stamp(p, 0); + for (int i = 0; i < kStages; ++i) { + mbar_init(&s.full[i], 1); + mbar_init(&s.empty[i], 1); + mbar_init(&s.sf_full[i], 4); + } + // every barrier below completes once per item, so item i waits parity i & 1 + mbar_init(&s.x_full, 1); + mbar_init(&s.fc1_done, 1); + mbar_init(&s.acc_ready, 2); // warps 0 and 4 drained their columns into acc1 + mbar_init(&s.h_ready, 4); + mbar_init(&s.fc2_pass[0], 1); + mbar_init(&s.fc2_pass[1], 1); + mbar_init(&s.fc2_free, 4); + mbar_init(&s.acc_drained, 4); // warps 4..7 read the last FC2 pass: the next item's FC1 may overwrite columns 0..255 + mbar_init(&s.route_done, 1); + fence_mbar_init(); + } + if (warp == 2) { tmem_alloc(&s.tmem, 512); tmem_relinquish(); tc_fence_before(); } + if (warp == 1 && lane == 0) { prefetch_tmap(&tm13); prefetch_tmap(&tm2); } + __syncthreads(); + tc_fence_after(); + const uint32_t tm = s.tmem; + // ring stage g (global over the CTA's items) of item-relative stage st of expert e + auto issue_stage = [&](int g, int st, int e, uint64_t pol) { + const int b = g % kStages; + mbar_expect_tx(&s.full[b], kStageBytes); + uint8_t* dst = s.ring[b]; + if (st < kFc1Stages) { + // the expert's rows as two boxes of 128 rows, K bytes [128 st, 128 st + 128) + tma_load_2d_w(dst, &tm13, &s.full[b], st * 128, e * kN1, pol); + tma_load_2d_w(dst + 16384, &tm13, &s.full[b], st * 128, e * kN1 + 128, pol); + // scale atoms of row blocks 0 and 1, k-blocks [4 st, 4 st + 4) + const uint8_t* sf = p.w13_sf + (size_t)e * (kN1 * kH / 16); + bulk_load_w(dst + 32768, sf + (size_t)(st * 4) * 512, 2048, &s.full[b], pol); + bulk_load_w(dst + 32768 + 2048, sf + (size_t)(64 + st * 4) * 512, 2048, &s.full[b], pol); + } else { + // W2 tiles 2 f and 2 f + 1 (256 rows x 64 B each) and their scale atoms (2 row blocks x 2 k-blocks) + const int f = st - kFc1Stages; +#pragma unroll + for (int tl = 0; tl < 2; ++tl) { + const int row0 = (2 * f + tl) * 256; + tma_load_2d_w(dst + tl * 16384, &tm2, &s.full[b], 0, e * kH + row0, pol); + bulk_load_w(dst + 32768 + tl * 2048, p.w2_sf + (size_t)e * (kH * 8) + (size_t)(row0 / 128) * 2 * 512, 2048, + &s.full[b], pol); + } + } + }; + auto prefetch_stage = [&](int st, int e) { // the same bytes as issue_stage, to L2 only + if (st < kFc1Stages) { + tma_prefetch_2d(&tm13, st * 128, e * kN1); + tma_prefetch_2d(&tm13, st * 128, e * kN1 + 128); + const uint8_t* sf = p.w13_sf + (size_t)e * (kN1 * kH / 16); + bulk_prefetch_l2(sf + (size_t)(st * 4) * 512, 2048); + bulk_prefetch_l2(sf + (size_t)(64 + st * 4) * 512, 2048); + } else { + const int row0 = (st - kFc1Stages) * 512; + tma_prefetch_2d(&tm2, 0, e * kH + row0); + tma_prefetch_2d(&tm2, 0, e * kH + row0 + 256); + bulk_prefetch_l2(p.w2_sf + (size_t)e * (kH * 8) + (size_t)(row0 / 128) * 2 * 512, 4096); + } + }; + // Rows of item it into row buffer rb, by one warp: its expert's tokens in increasing order, row group item_g + auto set_rows = [&](int rb, int it) { + const int e = s.item_e[it], g0 = s.item_g[it] * kRows; + int before = 0; // the expert's tokens below 32 h +#pragma unroll 1 + for (int h = 0; h < kGroups && 32 * h < M; ++h) { + const int t = lane + 32 * h; + int ks = -1; + if (t < M) { + const uint32_t* w = reinterpret_cast(&s.ids[t][0]); +#pragma unroll + for (int j = 0; j < kTopK / 2; ++j) { + const uint32_t v = w[j]; + ks = (int)(v & 0xFFFFu) == e ? 2 * j : ks; + ks = (int)(v >> 16) == e ? 2 * j + 1 : ks; + } + } + const uint32_t bal = __ballot_sync(0xffffffffu, ks >= 0); + const int r = before + __popc(bal & ((1u << lane) - 1u)) - g0; + if (ks >= 0 && r >= 0 && r < kRows) { + s.rtok[rb][r] = t; + s.rkslot[rb][r] = ks; + } + before += __popc(bal); + if (before >= g0 + kRows) break; + } + if (lane == 0) s.nrows[rb] = min(kRows, before - g0); + }; + // their x rows and scale bytes into xa / xsf, by threads tid = 0 .. nt - 1 + auto copy_rows = [&](int rb, int tid, int nt) { + const int n1 = s.nrows[rb] * 128, n = s.nrows[rb] * 144; // 16-B granules: x rows, then their scale bytes + if (GATHER_ASYNC) { + for (int idx = tid; idx < n; idx += nt) { + const uint8_t* src; + uint8_t* dst; + if (idx < n1) { + const int r = idx >> 7, c = (idx >> 3) & 15, q = idx & 7; + src = p.x_fp4 + (size_t)s.rtok[rb][r] * (kH / 2) + c * 128 + q * 16; + dst = s.xa + c * 2048 + (r / 8) * 1024 + (r % 8) * 128 + ((q ^ (r % 8)) * 16); + } else { + const int r = (idx - n1) >> 4, q = (idx - n1) & 15; + src = p.x_sf + (size_t)s.rtok[rb][r] * (kH / 16) + q * 16; + dst = &s.xsf[r][q * 16]; + } + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(smem_u32(dst)), "l"(src) : "memory"); + } + asm volatile("cp.async.wait_all;" ::: "memory"); + fence_proxy_async_smem(); + return; + } + // GATHER_BATCH loads a thread ahead of their stores: the global latency is paid once a batch + for (int i0 = tid; i0 < n; i0 += GATHER_BATCH * nt) { + uint4 v[GATHER_BATCH]; +#pragma unroll + for (int j = 0; j < GATHER_BATCH; ++j) { + const int idx = i0 + j * nt; + if (idx < n1) { // row r, K chunk c (128 B), 16-B granule q of the chunk + const int r = idx >> 7, c = (idx >> 3) & 15, q = idx & 7; + v[j] = *reinterpret_cast(p.x_fp4 + (size_t)s.rtok[rb][r] * (kH / 2) + c * 128 + q * 16); + } else if (idx < n) { + const int r = (idx - n1) >> 4, q = (idx - n1) & 15; + v[j] = *reinterpret_cast(p.x_sf + (size_t)s.rtok[rb][r] * (kH / 16) + q * 16); + } + } +#pragma unroll + for (int j = 0; j < GATHER_BATCH; ++j) { + const int idx = i0 + j * nt; + if (idx < n1) { + const int r = idx >> 7, c = (idx >> 3) & 15, q = idx & 7; + *reinterpret_cast(s.xa + c * 2048 + (r / 8) * 1024 + (r % 8) * 128 + ((q ^ (r % 8)) * 16)) = v[j]; + } else if (idx < n) { + const int r = (idx - n1) >> 4, q = (idx - n1) & 15; + *reinterpret_cast(&s.xsf[r][q * 16]) = v[j]; + } + } + } + fence_proxy_async_smem(); // the tensor core reads xa through the async proxy + }; + if (warp == 1) { + // expert cta is this CTA's first item whenever it has tokens: its weight stream starts before the routing is read + if (lane == 0) { + stamp(p, 9); + const uint64_t pol = policy_evict_first(); + for (int st = 0; st < SPEC_STAGES; ++st) { + issue_stage(st, st, cta, pol); + if (STAGE_STAMPS) stamp(p, 64 + st); + } + // HBM is otherwise idle until the routing is in: the next stages go to L2 now + for (int st = SPEC_STAGES; st < SPEC_STAGES + SPEC_PF && st < kTotStages; ++st) prefetch_stage(st, cta); + stamp(p, 2); + } + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (lane == 0) stamp(p, 10); + mbar_wait(&s.route_done, 0); + } else { + asm volatile("griddepcontrol.wait;" ::: "memory"); + // the routing as given: the 224 routing threads load every slot's expert at once (16 B a load), then store and + // count them + const int tid = (warp == 0 ? 0 : warp - 1) * 32 + lane, S = M * kTopK, S4 = S / 4; + constexpr int kPer4 = (kMaxTok * kTopK / 4 + 223) / 224; + int4 v[kPer4]; +#pragma unroll + for (int j = 0; j < kPer4; ++j) + v[j] = tid + 224 * j < S4 ? reinterpret_cast(p.ids_in)[tid + 224 * j] : make_int4(0, 0, 0, 0); +#pragma unroll + for (int j = 0; j < kPer4; ++j) { + const int i4 = tid + 224 * j; + if (i4 < S4) { + *reinterpret_cast(&s.ids[0][0] + 4 * i4) = + make_uint2((uint32_t)v[j].x | ((uint32_t)v[j].y << 16), (uint32_t)v[j].z | ((uint32_t)v[j].w << 16)); + atomicAdd(&s.cnt[v[j].x], 1); + atomicAdd(&s.cnt[v[j].y], 1); + atomicAdd(&s.cnt[v[j].z], 1); + atomicAdd(&s.cnt[v[j].w], 1); + } + } + for (int s2 = 4 * S4 + tid; s2 < S; s2 += 224) { + const int e2 = p.ids_in[s2]; + (&s.ids[0][0])[s2] = (uint16_t)e2; + atomicAdd(&s.cnt[e2], 1); + } + named_bar(4, 224); + if (warp == 0) { + // Work items: own experts cta, cta + G, ... that have tokens (row group 0), then this CTA's further row groups: + // the o-th one over all experts (expert order, then group) is CTA G - 1 - o % G's + int ni = 0; + for (int e = cta; e < kE; e += G) { + if (s.cnt[e] > 0) { + if (lane == 0) { + s.item_e[ni] = e; + s.item_g[ni] = 0; + } + ++ni; + } + } + int xe[16], xsum = 0; +#pragma unroll + for (int j = 0; j < 16; ++j) { + const int n = s.cnt[16 * lane + j]; + xe[j] = n > kRows ? (n - 1) / kRows : 0; + xsum += xe[j]; + } + int incl = xsum; +#pragma unroll + for (int d = 1; d < 32; d <<= 1) { + const int up = __shfl_up_sync(0xffffffffu, incl, d); + incl += lane >= d ? up : 0; + } + const int X = __shfl_sync(0xffffffffu, incl, 31); + for (int o = G - 1 - cta; o < X; o += G) { + const int owner = __ffs(__ballot_sync(0xffffffffu, incl > o)) - 1; + int e = 0, grp = 0; + if (lane == owner) { + int left = o - (incl - xsum); +#pragma unroll + for (int j = 0; j < 16; ++j) { + if (left >= 0 && left < xe[j]) { + e = 16 * lane + j; + grp = left + 1; + } + left -= xe[j]; + } + } + e = __shfl_sync(0xffffffffu, e, owner); + grp = __shfl_sync(0xffffffffu, grp, owner); + if (ni >= kMaxItems) __trap(); + if (lane == 0) { + s.item_e[ni] = e; + s.item_g[ni] = grp; + } + ++ni; + } + __syncwarp(); + if (lane == 0) { + s.nitems = ni; + s.spec_ok = s.cnt[cta] > 0 ? 1 : 0; + stamp(p, 1); + mbar_arrive(&s.route_done); + } + } + named_bar(4, 224); + } + const int nitems = s.nitems; + if (nitems == 0) { // no expert with tokens here: let the speculative loads land, then leave + if (warp == 1) { + if (lane == 0) + for (int b = 0; b < SPEC_STAGES; ++b) mbar_wait(&s.full[b], 0); + __syncwarp(); + } + named_bar(7, kThreads); + if (warp == 2) tmem_dealloc(tm, 512); + return; + } + tc_fence_after(); + const bool spec_ok = s.spec_ok != 0; + // ring stage of item 0's stage 0: the speculative fill (stages 0 .. SPEC_STAGES - 1) is item 0's own only if spec_ok + const int base0 = spec_ok ? 0 : SPEC_STAGES; + float ag = 0.f, sc = 0.f, a2 = 0.f; + auto load_scales = [&](int it) { + const int e = s.item_e[it]; + ag = p.g1_alpha[e]; + sc = p.g1_scale_c[e]; + a2 = p.g2_alpha[e]; + }; + // FC1 drain of quarter-0 accumulator columns [64 c0, 64 c0 + 128) (TMEM lanes 0..15 hold the item rows) into acc1 + auto drain_fc1 = [&](int c0) { +#pragma unroll 1 + for (int c = c0; c < c0 + 2; ++c) { + uint32_t r[64]; + tmem_ld_x64(tm + (uint32_t)(c * 64), r); + tmem_ld_wait(); + if (lane < 16) { +#pragma unroll + for (int i = 0; i < 16; ++i) + *reinterpret_cast(s.acc1 + lane * kARow + c * 64 + 4 * i) = + make_uint4(r[4 * i], r[4 * i + 1], r[4 * i + 2], r[4 * i + 3]); + } + } + tc_fence_before(); + __syncwarp(); + }; + // SwiGLU + NVFP4 epilogue (h = 0: the one FC1 pass), tasks (row, 16-wide block, half) = 16 x 8 x 2 = 256, + // task = et + k * nt. + // Accumulator columns of each 32-block: [0, 8) up 2m, [8, 16) gate 2m, [16, 24) up 2m + 1, [24, 32) gate 2m + 1 + auto epi_tasks = [&](int h, int et, int nt) { +#pragma unroll 1 + for (int task = et; task < 256; task += nt) { + const int row = task / 16, blk = (task / 2) % 8, hf = task % 2; + const float4* ar = reinterpret_cast(s.acc1 + row * kARow + 32 * blk + 4 * hf); + const float4 u_even = ar[0], g_even = ar[2], u_odd = ar[4], g_odd = ar[6]; + const float ga[8] = {g_even.x, g_odd.x, g_even.y, g_odd.y, g_even.z, g_odd.z, g_even.w, g_odd.w}; + const float ua[8] = {u_even.x, u_odd.x, u_even.y, u_odd.y, u_even.z, u_odd.z, u_even.w, u_odd.w}; + // trtllm-gen's SwiGLU + NVFP4 epilogue op for op (from its SASS); a = FC1 dequant, c = a x FC2 input scale: + // h' = (u a)(g sigmoid(g a)), sigmoid = rcp(ex2(-(g (a log2e))) + 1), block scale e4m3(max|h'| c / 6), FP4 (h' / scale) c + const float a_l2e = mul_ftz(ag, 1.4426950216293334961f); + float hv[8]; + float amax = 0.f; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const float ex = ex2_approx(-(ga[i] * a_l2e)); + const float sg = rcp_approx(ex + 1.0f); + const float v = fmaf(ua[i], ag, 0.0f) * (ga[i] * sg); + hv[i] = v; + amax = fmaxf(amax, fabsf(v)); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + const uint32_t sf = cvt_e4m3(fmaxf(0.f, fabsf(amax * sc)) * 0.16666667163372039795f); + const float sff = e4m3_to_f32(sf); + const float osc = sff == 0.f ? 0.f : rcp_approx(sff); + float qv[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) qv[i] = mul_ftz(hv[i], osc) * sc; + const uint32_t packed = cvt_e2m1x8(qv); + const int gblk = 8 * h + blk; // 16-wide block over the 128 intermediates; k-blocks 0, 1 + const int chunk = gblk / 2, r8 = row % 8; + const uint32_t off = (uint32_t)((row / 8) * 1024 + r8 * 128 + ((chunk ^ r8) * 16) + (gblk % 2) * 8 + hf * 4); +#pragma unroll + for (int cp = 0; cp < 8; ++cp) *reinterpret_cast(s.ha + cp * 2048 + off) = packed; + if (hf == 0) reinterpret_cast(&s.hsf[row][gblk / 4])[gblk % 4] = (uint8_t)sf; + } + }; + // Rows of item it (its tokens in increasing order, at most 16) and their x rows / scale bytes, by warps 0 and 3; + // row info is double-buffered by item parity (the previous item's FC2 drains still read theirs). + auto gather = [&](int it) { + const int rb = it & 1; + if (warp == 0) set_rows(rb, it); + named_bar(5, 64); + copy_rows(rb, (warp == 0 ? 0 : 32) + lane, 64); + named_bar(5, 64); + if (warp == 0 && lane == 0) { + mbar_arrive(&s.x_full); + if (DIAG_ROUTE && it == 0) stamp(p, 14); + } + }; + + if (warp == 1) { + if (lane == 0) { + // ---------------- TMA producer: every item's weight stream through the ring + const uint64_t pol = policy_evict_first(); + for (int it = 0; it < nitems; ++it) { + const int e = s.item_e[it], base = base0 + it * kTotStages; + const int st0 = (it == 0 && spec_ok) ? SPEC_STAGES : 0; + for (int st = st0; st < kTotStages; ++st) { + const int g = base + st; + if (FC2_PF > 0 && st == kFc1Stages - FC2_PF_LEAD) { + // the FC2 tiles beyond the ring go to L2 now, while the last FC1 stages and h are in progress + for (int j = kStages; j < kStages + FC2_PF; ++j) prefetch_stage(kFc1Stages + j, e); + } + if (PF_SET != 0 && ((st == st0 && (PF_SET & 1)) || (st == kFc1Stages + 4 && (PF_SET & 4)) || + (st == kFc1Stages + 6 && (PF_SET & 8)))) { + // these stages wait for MMAs that are slow to start after a phase change; let HBM fetch them meanwhile + const int n = st == st0 ? PF_N1 : st == kFc1Stages + 4 ? PF_N2 : 4; + for (int k = st; k < st + n && k < kTotStages; ++k) prefetch_stage(k, e); + } + if (g >= kStages) mbar_wait(&s.empty[g % kStages], ((g / kStages) - 1) & 1); + issue_stage(g, st, e, pol); + if (STAGE_STAMPS && it == 0) stamp(p, 64 + st); + } + } + } + __syncwarp(); + } else if (warp == 2) { + if (lane == 0) { + // ---------------- MMA issuer (tensor pipe: MMAs only) + const uint32_t idesc = make_idesc_nvf4(128u, 256u); + if (!spec_ok) { // the speculative fill was another CTA's expert: recycle its ring stages without MMAs + for (int g = 0; g < SPEC_STAGES; ++g) { + mbar_wait(&s.full[g % kStages], (g / kStages) & 1); + mbar_wait(&s.sf_full[g % kStages], (g / kStages) & 1); + mbar_arrive(&s.empty[g % kStages]); + } + } + for (int it = 0; it < nitems; ++it) { + const int base = base0 + it * kTotStages, par = it & 1; + if (it > 0) mbar_wait(&s.acc_drained, (it - 1) & 1); + mbar_wait(&s.x_full, par); + tc_fence_after(); + int n = 0; + for (int st = 0; st < kFc1Stages; ++st) { + const int g = base + st, b = g % kStages, ks = st; + mbar_wait(&s.full[b], (g / kStages) & 1); + mbar_wait(&s.sf_full[b], (g / kStages) & 1); + tc_fence_after(); + if (st == 0 && it == 0) stamp(p, 3); + if (st == 0 && it > 0 && it == nitems - 1) stamp(p, 11); + if (DIAG_ROUTE && st == 0 && it < 4) stamp(p, 24 + it); // DIAG: item it's first FC1 MMA + if (DIAG_ROUTE && st == kFc1Stages - 1 && it < 4) stamp(p, 28 + it); // DIAG: its last FC1 stage landed +#pragma unroll + for (int q = 0; q < 4; ++q, ++n) { + const int kb = ks * 4 + q; // k-blocks in K order: one accumulator, the stock summation order + umma_nvf4(tm, make_sdesc(smem_u32(s.xa + (kb / 4) * 2048) + (kb % 4) * 32, 0, 1024, 2), + make_sdesc(smem_u32(s.ring[b]) + q * 32, 0, 1024, 2), idesc, + tm + kColSfaX + (uint32_t)(kb % kSfaSlots) * 4, tm + kColSfb + (uint32_t)(n % kSfbSlots) * 8, + kb > 0); + } + tc_commit(&s.empty[b]); + } + tc_commit(&s.fc1_done); + mbar_wait(&s.h_ready, par); + tc_fence_after(); + if (it == 0) stamp(p, 6); + for (int st = kFc1Stages; st < kTotStages; ++st) { + const int g = base + st, b = g % kStages, f = st - kFc1Stages; + if (f == 4) { mbar_wait(&s.fc2_free, par); tc_fence_after(); } + mbar_wait(&s.full[b], (g / kStages) & 1); + mbar_wait(&s.sf_full[b], (g / kStages) & 1); + tc_fence_after(); +#pragma unroll + for (int tl = 0; tl < 2; ++tl) { + const int j = 2 * f + tl, jt = j % 8; +#pragma unroll + for (int kb = 0; kb < 2; ++kb) { + const int n2 = 64 + j * 2 + kb; // scale-factor slot, numbered on from FC1's 64 k-blocks + umma_nvf4(tm, make_sdesc(smem_u32(s.ha) + kb * 32, 0, 1024, 2), + make_sdesc(smem_u32(s.ring[b] + tl * 16384) + kb * 32, 0, 512, 4), idesc, + tm + kColSfa + (uint32_t)(jt * 2 + kb) * 4, tm + kColSfb + (uint32_t)(n2 % kSfbSlots) * 8, + (jt | kb) != 0); + } + } + tc_commit(&s.empty[b]); + if (f % 4 == 3) tc_commit(&s.fc2_pass[f / 4]); + } + } + } + __syncwarp(); + } else if (warp == 0 || warp == 3) { + // ---------------- x gather, FC1 drains (warp 0: columns [0, 128); warp 4 drains [128, 256)) and the epilogue + const int et = (warp == 0 ? 0 : 32) + lane; // 0..63 + gather(0); + for (int it = 0; it < nitems; ++it) { + const int par = it & 1; + load_scales(it); + if (warp == 0) { + mbar_wait(&s.fc1_done, par); + tc_fence_after(); + if (lane == 0 && it == 0) stamp(p, 4); + drain_fc1(0); + if (lane == 0) { + mbar_arrive(&s.acc_ready); + if (it == 0 && !DIAG_ROUTE) stamp(p, 13); + } + } + mbar_wait(&s.acc_ready, par); + epi_tasks(0, et, 192); + if (warp == 3 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 14); + fence_proxy_async_smem(); + named_bar(6, 192); // h and its scale words are complete for the TMEM warpgroup + if (it + 1 < nitems) gather(it + 1); // FC1 of item it is done: x and its scale bytes are free + } + } else { + // ---------------- TMEM-side warpgroup (warps 4..7): q = lane quarter (TMEM lanes 32 q .. 32 q + 31) + const int q = warp - 4; + const uint32_t lane_addr = (uint32_t)(q * 32) << 16; + if (!spec_ok) { // the speculative fill is recycled without MMAs: release its scale-factor side + for (int g = 0; g < SPEC_STAGES; ++g) { + mbar_wait(&s.full[g % kStages], (g / kStages) & 1); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[g % kStages]); + } + } + for (int it = 0; it < nitems; ++it) { + const int base = base0 + it * kTotStages, par = it & 1, rb = it & 1; + load_scales(it); + mbar_wait(&s.x_full, par); + const int nrows = s.nrows[rb]; + auto stage_sf2 = [&](int st) { // FC2 tiles 2 f, 2 f + 1: k-blocks 0, 1 of their row blocks (rows 0..127 | 128..255) + const int g = base + st, b = g % kStages; + mbar_wait(&s.full[b], (g / kStages) & 1); + if (STAGE_STAMPS && threadIdx.x == 128 && it == 0) stamp(p, 16 + st); + if (g >= 4) mbar_wait(&s.empty[(g - 4) % kStages], ((g - 4) / kStages) & 1); + const uint8_t* sf = s.ring[b] + 32768; + const int n0 = 64 + (st - kFc1Stages) * 4; + for (int i = 0; i < 4; ++i) { + const int tl = i / 2, kb = i % 2; + const uint4 v0 = *reinterpret_cast(sf + tl * 2048 + kb * 512 + lane * 16); + const uint4 v1 = *reinterpret_cast(sf + tl * 2048 + 1024 + kb * 512 + lane * 16); + const uint32_t r[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; + tmem_st_x8(tm + lane_addr + kColSfb + (uint32_t)((n0 + i) % kSfbSlots) * 8, r); + } + tmem_st_wait(); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[b]); + }; + // FC1 stage st: weight scale factors of k-blocks [4 ks, 4 ks + 4) (8 columns each: row blocks 2 h | 2 h + 1) and + // the x scale factors of the same k-blocks (row m -> lane m % 32, column m / 32 of each 4-column atom, every + // quarter holds a copy; rows >= nrows, including the A rows read past the 16 real ones, get zero scales) + for (int st = 0; st < kFc1Stages; ++st) { + const int g = base + st, b = g % kStages, ks = st; + mbar_wait(&s.full[b], (g / kStages) & 1); + if (STAGE_STAMPS && threadIdx.x == 128 && it == 0) stamp(p, 16 + st); + if (g >= 4) mbar_wait(&s.empty[(g - 4) % kStages], ((g - 4) / kStages) & 1); + const uint8_t* sf = s.ring[b] + 32768; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const uint4 v0 = *reinterpret_cast(sf + i * 512 + lane * 16); + const uint4 v1 = *reinterpret_cast(sf + 2048 + i * 512 + lane * 16); + const uint32_t r[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; + tmem_st_x8(tm + lane_addr + kColSfb + (uint32_t)((st * 4 + i) % kSfbSlots) * 8, r); + } + uint32_t xr[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int kb = ks * 4 + i / 4; + xr[i] = (i % 4 == 0 && lane < nrows) ? *reinterpret_cast(&s.xsf[lane][kb * 4]) : 0u; + } + tmem_st_x16(tm + lane_addr + kColSfaX + (uint32_t)((ks * 4) % kSfaSlots) * 4, xr); + tmem_st_wait(); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.sf_full[b]); + } + // EARLY_SF2: stage the first FC2 scale factors before h (they wait for FC2 data to land, which delays h) + for (int st = kFc1Stages; st < kFc1Stages + EARLY_SF2; ++st) stage_sf2(st); + if (q == 0) { // drain the upper 128 columns + mbar_wait(&s.fc1_done, par); + tc_fence_after(); + drain_fc1(2); + if (lane == 0) mbar_arrive(&s.acc_ready); + } + mbar_wait(&s.acc_ready, par); + epi_tasks(0, 64 + (warp - 4) * 32 + lane, 192); + fence_proxy_async_smem(); + named_bar(6, 192); // h complete + if (threadIdx.x == 128 && it == 0) stamp(p, 5); + // h scale factors: 16 atoms (tile jt, k-block kb) of 4 columns; row m = 32 c + lane is nonzero only in its tile + { + uint32_t r[32]; +#pragma unroll 1 + for (int half = 0; half < 2; ++half) { +#pragma unroll + for (int i = 0; i < 32; ++i) { + const int a = half * 8 + i / 4, col = i % 4; // atom a = 2 jt + kb + const int jt = a / 2, kb = a % 2; + const int row = 32 * col + lane; + r[i] = (row / 16 == jt && row % 16 < nrows) ? s.hsf[row % 16][kb] : 0u; + } + tmem_st_x32(tm + lane_addr + kColSfa + half * 32, r); + } + tmem_st_wait(); + } + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.h_ready); + if (q == 0 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 15); + for (int st = kFc1Stages + EARLY_SF2; st < kTotStages; ++st) { + if (st == kFc1Stages + 4) { + // pass 0 is fully staged; drain it once its MMAs are done, then free the accumulator for pass 1 + mbar_wait(&s.fc2_pass[0], par); + tc_fence_after(); + if (q == 0 && lane == 0 && it == 0) stamp(p, 7); + fc2_drain16<0, 4>(p, s, tm, q, lane, 0, a2, rb); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.fc2_free); + } + stage_sf2(st); + } + mbar_wait(&s.fc2_pass[1], par); + tc_fence_after(); + if (q == 0 && lane == 0 && it == 0) stamp(p, 12); + if (TAIL8 && it + 1 == nitems) fc2_drain16<0, 2>(p, s, tm, q, lane, 1, a2, rb); + else fc2_drain16<0, 4>(p, s, tm, q, lane, 1, a2, rb); + tc_fence_before(); + __syncwarp(); + if (lane == 0) mbar_arrive(&s.acc_drained); + } + } + if (TAIL8 && warp < 4) { // the last item's last pass, upper columns: every warp is free now (TMEM lane quarter = warp) + const int it = nitems - 1; + load_scales(it); + mbar_wait(&s.fc2_pass[1], it & 1); + tc_fence_after(); + fc2_drain16<2, 4>(p, s, tm, warp, lane, 1, a2, it & 1); + } + tc_fence_before(); + named_bar(7, kThreads); // every warp, including the drains, before TMEM is released + if (threadIdx.x == 128) stamp(p, 8); + if (warp == 2) tmem_dealloc(tm, 512); +} + +static_assert(sizeof(Smem) + 1024 == Q35_MKM_SMEM_BYTES, "the launcher's dynamic shared memory size"); + +} // namespace qwen_mkm_moe_wide + + +} // namespace sglang diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_topk.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_topk.cu new file mode 100644 index 00000000..12f1f9f6 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_topk.cu @@ -0,0 +1,259 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// The MoE front behind SGLang's router GEMM for verify steps wider than slot route serves (rows 17+): per token, the +// top-10 of its 512 bf16 router logits with the softmax renormalized over them (expert ids and bf16 weights, in the +// order and with the bits of FlashInfer's trtllm-gen routing, RenormalizeNaive), and the NVFP4 quantization of its +// bf16 row (16-element blocks, e4m3 scales, linear layout, the trtllm runner's quantizer), in one launch. The routed +// trtllm-gen MoE then takes the materialized routing and skips its own top-k kernel. +// +// One CTA per token: warp 0 ranks the logits (mkm_moe_m16.cu's routing, bit-exact with the single-block routing), the +// other 7 warps quantize the row (router_front.cu's quantize_block). Launched as a programmatic dependent; it waits +// before reading anything and releases its dependents at once (they read its outputs after their own wait). +// +// Built with NVRTC. + +#include +#include +#include + +// NVRTC has no +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef int int32_t; +typedef unsigned int uint32_t; +#ifndef INFINITY +#define INFINITY __int_as_float(0x7f800000) +#endif + +namespace q35_route_topk { + +constexpr uint32_t kHidden = 4096, kExperts = 512, kTopK = 10, kThreads = 256, kSfVec = 16; +constexpr uint32_t kBlocksPerRow = kHidden / kSfVec; // 256 + +struct Params { + const __nv_bfloat16* logits; // [M, 512] + const __nv_bfloat16* x; // [M, 4096] + const float* global_scale; // [1] + int32_t* topk_ids; // [M, 10] + uint16_t* topk_weights; // [M, 10] bf16 + uint8_t* x_fp4; // [M, 2048] + uint8_t* x_sf; // [M, 256] +}; + +// ---- the NVFP4 quantizer (router_front.cu) +__device__ __forceinline__ float bf16_lo(uint32_t v) { return __uint_as_float(v << 16); } +__device__ __forceinline__ float bf16_hi(uint32_t v) { return __uint_as_float(v & 0xFFFF0000u); } +__device__ __forceinline__ float rcp_approx_ftz(float a) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(a)); + return y; +} +__device__ __forceinline__ uint32_t cvt_e4m3(float a) { + uint32_t out; + asm("{\n.reg .b16 pair;\n.reg .f32 zero;\nmov.f32 zero, 0f00000000;\n" + "cvt.rn.satfinite.e4m3x2.f32 pair, zero, %1;\ncvt.u32.u16 %0, pair;\n}" + : "=r"(out) + : "f"(a)); + return out & 0xFFu; +} +__device__ __forceinline__ float e4m3_to_f32(uint32_t v) { + float out; + asm("{\n.reg .b16 pair;\n.reg .b32 h2;\n.reg .b16 lo, hi;\ncvt.u16.u32 pair, %1;\n" + "cvt.rn.f16x2.e4m3x2 h2, pair;\nmov.b32 {lo, hi}, h2;\ncvt.f32.f16 %0, lo;\n}" + : "=f"(out) + : "r"(v)); + return out; +} +__device__ __forceinline__ uint32_t cvt_e2m1x8(const float (&v)[8]) { + uint32_t out; + asm("{\n.reg .b8 b0, b1, b2, b3;\n" + "cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\ncvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\ncvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" + "mov.b32 %0, {b0, b1, b2, b3};\n}" + : "=r"(out) + : "f"(v[0]), "f"(v[1]), "f"(v[2]), "f"(v[3]), "f"(v[4]), "f"(v[5]), "f"(v[6]), "f"(v[7])); + return out; +} +__device__ __forceinline__ uint32_t bf16x2_abs_max(uint32_t a, uint32_t b) { + uint32_t out; + asm("{\n.reg .b32 ta, tb;\nand.b32 ta, %1, 0x7FFF7FFF;\nand.b32 tb, %2, 0x7FFF7FFF;\nmax.bf16x2 %0, ta, tb;\n}" + : "=r"(out) + : "r"(a), "r"(b)); + return out; +} +__device__ __forceinline__ uint32_t bf16x2_max(uint32_t a, uint32_t b) { + uint32_t out; + asm("max.bf16x2 %0, %1, %2;" : "=r"(out) : "r"(a), "r"(b)); + return out; +} +__device__ __forceinline__ void quantize_block(const uint4& a, const uint4& b, float gs, uint32_t& sf, uint2& packed) { + const uint32_t w[8] = {a.x, a.y, a.z, a.w, b.x, b.y, b.z, b.w}; + const uint32_t m = bf16x2_max(bf16x2_max(bf16x2_abs_max(w[0], w[1]), bf16x2_abs_max(w[2], w[3])), + bf16x2_max(bf16x2_abs_max(w[4], w[5]), bf16x2_abs_max(w[6], w[7]))); + const float amax = fmaxf(bf16_lo(m), bf16_hi(m)); + const float scale = __fmul_rn(gs, __fmul_rn(amax, rcp_approx_ftz(6.0f))); + sf = cvt_e4m3(scale); + const float sf_f = e4m3_to_f32(sf); + const float out_scale = (sf_f == 0.0f) ? 0.0f : rcp_approx_ftz(__fmul_rn(sf_f, rcp_approx_ftz(gs))); + float lo[8], hi[8]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + lo[2 * i] = __fmul_rn(bf16_lo(w[i]), out_scale); + lo[2 * i + 1] = __fmul_rn(bf16_hi(w[i]), out_scale); + hi[2 * i] = __fmul_rn(bf16_lo(w[4 + i]), out_scale); + hi[2 * i + 1] = __fmul_rn(bf16_hi(w[4 + i]), out_scale); + } + packed.x = cvt_e2m1x8(lo); + packed.y = cvt_e2m1x8(hi); +} + +// ---- the routing (mkm_moe_m16.cu): key = bf16 order above 65535 - expert (higher score first, smaller expert on ties) +__device__ __forceinline__ uint32_t topk_key(uint32_t bits, uint32_t expert) { + const uint32_t twiddled = bits ^ ((bits & 0x8000u) ? 0xFFFFu : 0x8000u); + return (twiddled << 16) | (65535u - expert); +} +__device__ __forceinline__ void topk_unkey(uint32_t key, uint32_t& bits, uint32_t& expert) { + expert = 65535u - (key & 0xFFFFu); + const uint32_t twiddled = key >> 16; + bits = twiddled ^ ((twiddled & 0x8000u) ? 0x8000u : 0xFFFFu); +} +__device__ __forceinline__ uint32_t max16(const uint32_t (&k)[16]) { + const uint32_t a = max(max(max(k[0], k[1]), max(k[2], k[3])), max(max(k[4], k[5]), max(k[6], k[7]))); + const uint32_t b = max(max(max(k[8], k[9]), max(k[10], k[11])), max(max(k[12], k[13]), max(k[14], k[15]))); + return max(a, b); +} +__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_sum_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.add.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_topk_of_lanes(uint32_t v, uint32_t lane) { + uint32_t mine = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(v); + v = v == top ? 0u : v; + mine = lane == k ? top : mine; + } + return mine; +} +// Top-10 keys of 512 (16 per lane), lane k holding the k-th largest (keys unique and nonzero). +__device__ __forceinline__ uint32_t warp_top10_redux(uint32_t (&keys)[16], uint32_t lane, uint32_t* s_cand) { + uint32_t v = max16(keys); + uint32_t bound = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + bound = warp_max_u32(v); + v = v == bound ? 0u : v; + } + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) count += keys[i] >= bound; + const uint32_t total = warp_sum_u32(count); + if (total <= 32) { + uint32_t offset = count; +#pragma unroll + for (uint32_t d = 1; d < 32; d <<= 1) { + const uint32_t up = __shfl_up_sync(0xffffffffu, offset, d); + offset += lane >= d ? up : 0u; + } + offset -= count; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) { + if (keys[i] >= bound) s_cand[offset++] = keys[i]; + } + __syncwarp(); + const uint32_t cand = lane < total ? s_cand[lane] : 0u; + __syncwarp(); + return warp_topk_of_lanes(cand, lane); + } + uint32_t mine = 0; +#pragma unroll 1 + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(max16(keys)); +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) keys[i] = keys[i] == top ? 0u : keys[i]; + mine = lane == k ? top : mine; + } + return mine; +} +// FlashInfer calcSoftmax (bf16 scores) as built with -use_fast_math. +__device__ __forceinline__ float fast_expf(float x) { + float y; + asm("{\n .reg .f32 t;\n mul.ftz.f32 t, %1, 0f3FB8AA3B;\n ex2.approx.ftz.f32 %0, t;\n}" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float fast_div(float a, float b) { + float y; + asm("div.full.ftz.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ uint16_t softmax_topk(uint16_t score_bits, uint32_t lane) { + namespace cg = cooperative_groups; + const auto warp = cg::tiled_partition<32>(cg::this_thread_block()); + const float si = __uint_as_float((uint32_t)score_bits << 16); + float max_score = -INFINITY; + if (lane < kTopK) max_score = si >= max_score ? si : max_score; + max_score = cg::reduce(warp, max_score, cg::greater()); + float sum_score = 0.f, new_score = 0.f; + if (lane < kTopK) { + new_score = fast_expf(__fsub_rn(si, max_score)); + sum_score = __fadd_rn(sum_score, new_score); + } + sum_score = cg::reduce(warp, sum_score, cg::plus()); + uint16_t out = score_bits; + if (lane < kTopK) { + const __nv_bfloat16 b = __float2bfloat16_rn(fast_div(new_score, sum_score)); + out = *reinterpret_cast(&b); + } + return out; +} + +__device__ __forceinline__ void route_topk_quant(const Params& p) { + __shared__ __align__(16) uint32_t s_cand[32]; + const uint32_t t = blockIdx.x, warp = threadIdx.x / 32, lane = threadIdx.x % 32; + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + if (warp == 0) { + const uint4* src = reinterpret_cast(p.logits + static_cast(t) * kExperts + lane * 16); + const uint4 v0 = src[0], v1 = src[1]; + uint32_t keys[16]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t w = i < 4 ? (&v0.x)[i] : (&v1.x)[i - 4]; + keys[2 * i] = topk_key(w & 0xFFFFu, lane * 16 + 2 * i); + keys[2 * i + 1] = topk_key(w >> 16, lane * 16 + 2 * i + 1); + } + const uint32_t key = warp_top10_redux(keys, lane, s_cand); + uint32_t bits, expert; + topk_unkey(key, bits, expert); + const uint16_t wbits = softmax_topk(static_cast(bits), lane); + if (lane < kTopK) { + p.topk_ids[t * kTopK + lane] = static_cast(expert); + p.topk_weights[t * kTopK + lane] = wbits; + } + } else if (p.x_fp4 != nullptr) { // null: the router front already quantized the rows + const float gs = p.global_scale[0]; + const __nv_bfloat16* row = p.x + static_cast(t) * kHidden; + for (uint32_t blk = threadIdx.x - 32; blk < kBlocksPerRow; blk += kThreads - 32) { + const uint4* src = reinterpret_cast(row + blk * kSfVec); + uint32_t sf; + uint2 packed; + quantize_block(src[0], src[1], gs, sf, packed); + *reinterpret_cast(p.x_fp4 + static_cast(t) * (kHidden / 2) + blk * 8) = packed; + p.x_sf[static_cast(t) * kBlocksPerRow + blk] = static_cast(sf); + } + } +} + +} // namespace q35_route_topk + +extern "C" __global__ __launch_bounds__(q35_route_topk::kThreads, 1) void q35_route_topk_quant( + const __grid_constant__ q35_route_topk::Params p) { + q35_route_topk::route_topk_quant(p); +} diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/router_front.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/router_front.cu index 41bdce53..c82c8893 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/router_front.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/router_front.cu @@ -921,7 +921,9 @@ SGL_DEVICE void router_front_kernel(const Params& p) { const uint32_t warp = threadIdx.x / 32; const uint32_t expert0 = blockIdx.x * kExpertsPerCta; - const uint32_t M = p.num_tokens; + // token tile blockIdx.y: rows [t0, t0 + M) of x, logits and the NVFP4 input (one tile without route_gemm) + const uint32_t t0 = blockIdx.y * kMaxM; + const uint32_t M = min(p.num_tokens - t0, kMaxM); const uint32_t tb = blockIdx.x * 8; float* s_part = reinterpret_cast(smem + L::kPart); const uint32_t bar_w = smem_addr(smem + L::kBarW); @@ -964,7 +966,7 @@ SGL_DEVICE void router_front_kernel(const Params& p) { #pragma unroll for (uint32_t row = 0; row < kMaxM; ++row) { if (row < M) { - bulk_load(smem_addr(smem + L::kX + row * kRowBytes), p.x + static_cast(row) * kHidden, kHidden * 2, + bulk_load(smem_addr(smem + L::kX + row * kRowBytes), p.x + static_cast(t0 + row) * kHidden, kHidden * 2, bar_x); } } @@ -986,8 +988,8 @@ SGL_DEVICE void router_front_kernel(const Params& p) { uint32_t sf; uint2 packed; quantize_block(src[0], src[1], p.global_scale[0], sf, packed); - *reinterpret_cast(p.x_fp4 + static_cast(blockIdx.x) * (kHidden / 2) + blk * 8) = packed; - p.x_sf[static_cast(blockIdx.x) * kBlocksPerRow + blk] = static_cast(sf); + *reinterpret_cast(p.x_fp4 + static_cast(t0 + blockIdx.x) * (kHidden / 2) + blk * 8) = packed; + p.x_sf[static_cast(t0 + blockIdx.x) * kBlocksPerRow + blk] = static_cast(sf); } } __syncthreads(); @@ -997,7 +999,7 @@ SGL_DEVICE void router_front_kernel(const Params& p) { const uint32_t e = threadIdx.x % kExpertsPerCta; const float lo = s_part[m * kExpertsPerCta + e]; const float hi = s_part[(16 + m) * kExpertsPerCta + e]; - p.logits[m * kExperts + expert0 + e] = __float2bfloat16_rn(__fadd_rn(lo, hi)); + p.logits[(t0 + m) * kExperts + expert0 + e] = __float2bfloat16_rn(__fadd_rn(lo, hi)); } if (has_sh && threadIdx.x < M * 8) { const uint32_t m = threadIdx.x / 8; diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index 4484c9ea..7ccd089b 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -110,6 +110,8 @@ # bit for bit. A numeric slot's outputs are held to fp64 on the same quantized operands # (test/kernels), at most these multiples of the error of SGLang's own op. NUMERIC_TOLERANCES = { + # The tiled router front stays within one BF16 ULP of SGLang's whole-M GEMM. + "route_gemm": {"route_gemm.router_logits.bf16_ulp_vs_stock": 1.0}, "gdn_push": { "gdn_push.out_proj.rel_l2_vs_stock": 1.05, "gdn_push.out_proj.max_abs_vs_stock": 2.0, diff --git a/test/kernels/test_qwen3_5_moe.py b/test/kernels/test_qwen3_5_moe.py index e4c9cac9..6cf4d05c 100644 --- a/test/kernels/test_qwen3_5_moe.py +++ b/test/kernels/test_qwen3_5_moe.py @@ -9,8 +9,11 @@ SGLang serves Qwen3.5 at TP8 and TP4 (intermediate 128 / 256 a rank; routing from the logits, RenormalizeNaive, deferred finalize): each token's 10 expert rows and their weights, at 1, 4 and 7 rows (the 2-CTA engines) and at TP8 also 8, 14 and 16 rows (the - 16-token engine), also with every token routed to the same experts; graph replay of - both; + 16-token engine), 17, 21, 28 and 32 rows (the 32-token engine), 33, 56 and 64 rows + (the 64-token engine, also on route_topk's routing), 65, 84, 112 and 128 rows (the + 128-token engine, on route_topk's routing) and 129, 168, 224, 256, 257, 336, 448 and 512 + rows (the wide engines, on route_topk's routing), also with every token routed to the + same experts; graph replay of each; - the router front with the shared gate/up stage (slot route_shared, numeric): its logits and NVFP4 input bit for bit the plain launch's, and its gate/up rows against fp64 on SGLang's static FP8 quantization of the rows, at most 1.05 x the mean relative @@ -166,12 +169,12 @@ def _stock(layer, logits, x_fp4, x_sf): return rows[permuted.reshape(-1).long()[: m * TOP_K]], weights -@pytest.mark.parametrize("m", [1, 4, 7, 8, 14, 16]) +@pytest.mark.parametrize("m", [1, 4, 7, 8, 14, 16, 17, 21, 28, 32, 33, 56, 64]) @pytest.mark.parametrize("seed", [0, 1]) @pytest.mark.parametrize("shared_route", [False, True]) def test_mkm_is_the_trtllm_routed_experts(layer, m, seed, shared_route): - if m > mkm.max_tokens(layer["inter"]): - pytest.skip("no 16-token engine at this intermediate") + if m > mkm.max_tokens(layer["inter"], m32=True, m64=True): + pytest.skip("no engine for these rows at this intermediate") logits, x_fp4, x_sf = _inputs(m, seed) if shared_route: # every token routes to the same 10 experts: one slot computes each for all rows @@ -184,10 +187,10 @@ def test_mkm_is_the_trtllm_routed_experts(layer, m, seed, shared_route): assert not differ.any(), f"rows {differ.nonzero().flatten().tolist()} differ" -@pytest.mark.parametrize("m", [7, 14]) +@pytest.mark.parametrize("m", [7, 14, 28, 56]) def test_mkm_replays_from_a_cuda_graph(layer, m): - if m > mkm.max_tokens(layer["inter"]): - pytest.skip("no 16-token engine at this intermediate") + if m > mkm.max_tokens(layer["inter"], m32=True, m64=True): + pytest.skip("no engine for these rows at this intermediate") logits, x_fp4, x_sf = _inputs(m, 2) eager = _fused(layer, logits, x_fp4, x_sf) graph = torch.cuda.CUDAGraph() @@ -198,6 +201,80 @@ def test_mkm_replays_from_a_cuda_graph(layer, m): assert all(torch.equal(a, b) for a, b in zip(eager, captured)) +@pytest.mark.parametrize( + "m", [33, 35, 56, 64, 65, 84, 112, 128, 129, 168, 224, 256, 257, 336, 448, 512] +) +@pytest.mark.parametrize("shared_route", [False, True]) +def test_mkm_on_route_topks_routing_is_the_trtllm_routed_experts( + layer, m, shared_route +): + """Slots mkm_m64 / mkm_m128 / mkm_m256 / mkm_m512: the 64 / 128 / 256 / 512-token engine + on route_topk's routing (ids and weights) gives the runner's rows on the logits, bit for + bit.""" + if m > mkm.max_tokens( + layer["inter"], m32=True, m64=True, m128=True, m256=True, m512=True + ): + pytest.skip("no engine for these rows at this intermediate") + logits, x_fp4, x_sf = _inputs(m, 5) + if shared_route: + logits = logits[:1].expand(m, -1).contiguous() + x = torch.randn(m, H, device="cuda").to(torch.bfloat16) + ids, weights, _, _ = router.route_topk_quant( + logits, x, torch.tensor([X_SCALE], device="cuda"), quantize=False + ) + rows, got_weights = mkm.mkm_moe( + layer["maps"], + logits, + x_fp4, + x_sf, + layer["w13_sf"].view(torch.uint8), + layer["w2_sf"].view(torch.uint8), + layer["g1"], + layer["g1c"], + layer["g2"], + routing=(ids, weights), + ) + want_rows, want_weights = _stock(layer, logits, x_fp4, x_sf) + assert torch.equal(got_weights.view(torch.int16), want_weights.view(torch.int16)) + differ = (rows.view(torch.int16) != want_rows.view(torch.int16)).any(dim=1) + assert not differ.any(), f"rows {differ.nonzero().flatten().tolist()} differ" + + +@pytest.mark.parametrize("m", [56, 112, 224, 448]) +def test_mkm_on_route_topks_routing_replays_from_a_cuda_graph(layer, m): + if m > mkm.max_tokens( + layer["inter"], m32=True, m64=True, m128=True, m256=True, m512=True + ): + pytest.skip("no engine for these rows at this intermediate") + logits, x_fp4, x_sf = _inputs(m, 6) + x = torch.randn(m, H, device="cuda").to(torch.bfloat16) + ids, weights, _, _ = router.route_topk_quant( + logits, x, torch.tensor([X_SCALE], device="cuda"), quantize=False + ) + + def run(): + return mkm.mkm_moe( + layer["maps"], + logits, + x_fp4, + x_sf, + layer["w13_sf"].view(torch.uint8), + layer["w2_sf"].view(torch.uint8), + layer["g1"], + layer["g1c"], + layer["g2"], + routing=(ids, weights), + ) + + eager = run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run() + graph.replay() + torch.cuda.synchronize() + assert all(torch.equal(a, b) for a, b in zip(eager, captured)) + + # ---------------------------------------------------------------- the router front @@ -304,3 +381,119 @@ def test_router_gemm_quant_shared_replays_from_a_cuda_graph(): graph.replay() torch.cuda.synchronize() assert all(torch.equal(a, b) for a, b in zip(eager, captured)) + + +# ---------------------------------------------------------------- slot route_topk + + +def _routed_stock(layer, routing, x_fp4, x_sf): + """``_stock`` on a materialized routing (ids int32, weights bf16 [M, 10]): the + runner's routed path (FlashInfer's routed MoE, as SGLang calls it).""" + from flashinfer.fused_moe import trtllm_fp4_block_scale_routed_moe + from flashinfer.fused_moe.core import ActivationType + + m = x_fp4.shape[0] + rows, weights, permuted = trtllm_fp4_block_scale_routed_moe( + topk_ids=routing, + routing_bias=None, + hidden_states=x_fp4, + hidden_states_scale=x_sf.view(torch.float8_e4m3fn), + gemm1_weights=layer["w13"], + gemm1_weights_scale=layer["w13_sf"], + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=layer["w2"], + gemm2_weights_scale=layer["w2_sf"], + gemm2_bias=None, + output1_scale_scalar=layer["g1c"], + output1_scale_gate_scalar=layer["g1"], + output2_scale_scalar=layer["g2"], + per_token_scale=None, + num_experts=E, + top_k=TOP_K, + n_group=0, + topk_group=0, + intermediate_size=layer["inter"], + local_expert_offset=0, + local_num_experts=E, + routed_scaling_factor=None, + routing_method_type=1, + do_finalize=False, + activation_type=ActivationType.Swiglu, + tune_max_num_tokens=max(8, m), + enable_pdl=True, + )[:3] + return rows[permuted.reshape(-1).long()[: m * TOP_K]], weights + + +def _topk_inputs(m, seed): + """Router logits with ties (a bf16 grid) and a row with -inf / equal scores, the + input rows and the routed experts' input scale.""" + g = torch.Generator(device="cuda").manual_seed(seed) + logits = (torch.randn(m, E, device="cuda", generator=g) * 2).to(torch.bfloat16) + logits[0, ::3] = logits[0, 1] # many ties + if m > 2: + logits[1] = 0.5 # every score equal: the lowest expert ids + x = torch.randn(m, H, device="cuda", generator=g).to(torch.bfloat16) + x[0, :32] = 0 + x[-1, 32:48] *= 1e4 + return logits, x, torch.tensor([X_SCALE], device="cuda") + + +@pytest.mark.parametrize("m", [1, 7, 17, 21, 28, 56, 112, 224, 448, 896]) +def test_route_topk_is_the_trtllm_routing_and_the_runners_quantizer(layer, m): + from flashinfer import fp4_quantize + + for seed in range(2): + logits, x, scale = _topk_inputs(m, seed=1000 * m + seed) + ids, weights, x_fp4, x_sf = router.route_topk_quant(logits, x, scale) + q, q_sf = fp4_quantize(x, scale, 16, False, False) + assert torch.equal(x_fp4, q.view(torch.uint8).reshape(m, -1)) + assert torch.equal(x_sf, q_sf.view(torch.uint8).reshape(m, -1)) + # the routed experts on this routing equal the runner's own routing from the + # logits, row for row, weights included + want_rows, want_w = _stock(layer, logits, x_fp4, x_sf) + got_rows, got_w = _routed_stock(layer, (ids, weights), x_fp4, x_sf) + assert torch.equal(got_w.view(torch.int16), want_w.view(torch.int16)) + assert torch.equal(got_rows.view(torch.int16), want_rows.view(torch.int16)) + + +def test_route_topk_replays_from_a_cuda_graph(): + logits, x, scale = _topk_inputs(56, seed=3) + router.build_topk(x.device.index) + eager = router.route_topk_quant(logits, x, scale) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = router.route_topk_quant(logits, x, scale) + graph.replay() + torch.cuda.synchronize() + assert all(torch.equal(a, b) for a, b in zip(eager, captured)) + + +@pytest.mark.parametrize("m", [17, 21, 28, 35, 49, 56]) +def test_router_front_tiles_are_its_16_row_launches(m): + """Slot route_gemm (numeric past 16 rows): the router front over 16-row token tiles in + one launch equals the 16-row launch on each tile, bit for bit; its NVFP4 input is the + runner's quantizer's; its logits stay within a bf16 ulp of cuBLAS's GEMM; and + route_topk's routing-only launch routes as its quantizing launch.""" + from flashinfer import fp4_quantize + + x, w, scale = _router_inputs(m, seed=900 + m) + router.build(x.device.index, 16) + logits, x_fp4, x_sf = router.router_gemm_quant_tiled(x, w, scale) + for t0 in range(0, m, 16): + tile = router.router_gemm_quant(x[t0 : t0 + 16].contiguous(), w, scale) + for a, b in zip((logits, x_fp4, x_sf), tile): + assert torch.equal(a[t0 : t0 + 16], b) + q, q_sf = fp4_quantize(x, scale, 16, False, False) + assert torch.equal(x_fp4, q.view(torch.uint8).reshape(m, -1)) + assert torch.equal(x_sf, q_sf.view(torch.uint8).reshape(m, -1)) + stock = torch.nn.functional.linear(x, w).float() + ulp = stock.abs().clamp_min(1e-30) * 2.0**-7 + assert ((logits.float() - stock).abs() <= ulp).all() + routed = router.route_topk_quant(logits, x, scale) + only = router.route_topk_quant(logits, x, scale, quantize=False) + assert only[2] is None and only[3] is None + assert torch.equal(routed[0], only[0]) and torch.equal(routed[1], only[1]) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 117fab47..0940eb14 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -117,15 +117,40 @@ def test_verify_sample_is_refused_with_the_switch_that_replaces_it(monkeypatch): def test_expert_parallel_keeps_route_and_boundary_slots(monkeypatch): requested = frozenset( - {"route", "mkm", "mkm_wide", "hop", "gdn_push", "draft_moe", "draft_hop"} + { + "route", + "route_topk", + "mkm", + "mkm_wide", + "mkm_m32", + "mkm_m64", + "mkm_m128", + "mkm_m256", + "mkm_m512", + "hop", + "gdn_push", + "draft_moe", + "draft_hop", + } ) monkeypatch.setattr(runtime, "SLOTS", requested) assert runtime.apply_expert_parallel(1) == frozenset() assert runtime.SLOTS == requested assert runtime.apply_expert_parallel(8) == frozenset( - {"mkm", "mkm_wide", "draft_moe"} + { + "mkm", + "mkm_wide", + "mkm_m32", + "mkm_m64", + "mkm_m128", + "mkm_m256", + "mkm_m512", + "draft_moe", + } + ) + assert runtime.SLOTS == frozenset( + {"route", "route_topk", "hop", "gdn_push", "draft_hop"} ) - assert runtime.SLOTS == frozenset({"route", "hop", "gdn_push", "draft_hop"}) def test_the_runtime_gives_the_executor_the_build_contexts_binding(monkeypatch): diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 6456a1a4..48938560 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -162,6 +162,31 @@ def test_an_sglang_without_the_eagle_forward_view_is_refused(monkeypatch): {"route", "mkm", "mkm_wide"}, "slot 'mkm_wide' serves the steps of more than one request ('multi_request')", ), + ({"mkm_m32"}, "slot 'mkm_m32' takes slot 'route_topk''s rows"), + ( + {"route_topk", "mkm_m32"}, + "slot 'mkm_m32' serves the steps of more than two requests ('multi_request')", + ), + ({"mkm_m64"}, "slot 'mkm_m64' runs on slot 'route_topk''s routing"), + ( + {"route_topk", "mkm_m64"}, + "slot 'mkm_m64' serves the steps of more than four requests ('multi_request')", + ), + ({"mkm_m128"}, "slot 'mkm_m128' runs on slot 'route_topk''s routing"), + ( + {"route_topk", "mkm_m128"}, + "slot 'mkm_m128' serves the steps of more than nine requests ('multi_request')", + ), + ({"mkm_m256"}, "slot 'mkm_m256' runs on slot 'route_topk''s routing"), + ( + {"route_topk", "mkm_m256"}, + "slot 'mkm_m256' serves the steps of more than 18 requests ('multi_request')", + ), + ({"mkm_m512"}, "slot 'mkm_m512' runs on slot 'route_topk''s routing"), + ( + {"route_topk", "mkm_m512"}, + "slot 'mkm_m512' serves the steps of more than 36 requests ('multi_request')", + ), ], ) def test_a_slot_set_that_cannot_run_is_refused_before_the_load( @@ -217,7 +242,17 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): **{**dict.fromkeys(plan.Facts.__struct_fields__, True), "mkm_wide": False} ) assert _on(plan.slots(7, every)) == sorted( - set(plan.Slots.__struct_fields__) - {"route_shared"} + set(plan.Slots.__struct_fields__) + - { + "mkm_m32", + "mkm_m64", + "mkm_m128", + "mkm_m256", + "mkm_m512", + "route_gemm", + "route_shared", + "route_topk", + } ) # Two requests also take every measured low-M launch and the hop's FP8 copy. assert _on(plan.slots(14, every)) == [ @@ -238,20 +273,39 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): ] # The 28/56-row low-M paths stay stock; GDN core and MKA serve through 224. for rows in (28, 56, 112, 224): - assert _on(plan.slots(rows, every)) == [ + expected = [ "gdn_conv", "gdn_core", "gdn_pdl", "gdn_verify", "mka", + "route_topk", ] + if rows <= plan.MKM_M32_MAX_ROWS: + expected.append("mkm_m32") + if plan.MKM_M32_MAX_ROWS < rows <= plan.MKM_M64_MAX_ROWS: + expected.append("mkm_m64") + if plan.MKM_M64_MAX_ROWS < rows <= plan.MKM_M128_MAX_ROWS: + expected.append("mkm_m128") + if plan.MKM_M128_MAX_ROWS < rows <= plan.MKM_M256_MAX_ROWS: + expected.append("mkm_m256") + if rows <= plan.ROUTE_GEMM_MAX_ROWS: + expected.append("route_gemm") + assert _on(plan.slots(rows, every)) == sorted(expected) for rows in (448, 896): - assert _on(plan.slots(rows, every)) == [ + expected = [ "gdn_conv", "gdn_pdl", "gdn_verify", "mka", + "route_topk", ] + if rows <= plan.MKM_M512_MAX_ROWS: + expected.append("mkm_m512") + assert _on(plan.slots(rows, every)) == sorted(expected) + assert _on(plan.slots(28, plan.Facts(route_topk=True))) == ["route_topk"] + assert _on(plan.slots(7, plan.Facts(route_topk=True))) == ["route_topk"] + assert not plan.slots(28, plan.Facts(route_gemm=True)).route_gemm # SGLang offers every multiple of 7 up to MAX_ROWS; every request count is a candidate offered = tuple(range(7, plan.MAX_ROWS + 1, 7)) assert plan.candidates(offered, 7) == offered @@ -284,6 +338,26 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): assert _on(plan.slots(14, wide)) == ["mkm", "route"] assert _on(plan.slots(28, wide)) == [] assert _on(plan.slots(14, plan.Facts(route=True, mkm=True))) == ["route"] + m32 = plan.Facts(route_topk=True, mkm_m32=True) + assert _on(plan.slots(28, m32)) == ["mkm_m32", "route_topk"] + assert _on(plan.slots(14, m32)) == ["route_topk"] + assert _on(plan.slots(56, m32)) == ["route_topk"] + assert _on(plan.slots(28, plan.Facts(mkm_m32=True))) == [] + m64 = plan.Facts(route_topk=True, mkm_m64=True) + assert _on(plan.slots(56, m64)) == ["mkm_m64", "route_topk"] + assert _on(plan.slots(28, m64)) == ["route_topk"] + assert _on(plan.slots(112, m64)) == ["route_topk"] + assert _on(plan.slots(56, plan.Facts(mkm_m64=True))) == [] + m128 = plan.Facts(route_topk=True, mkm_m128=True) + assert _on(plan.slots(112, m128)) == ["mkm_m128", "route_topk"] + assert _on(plan.slots(56, m128)) == ["route_topk"] + assert _on(plan.slots(224, m128)) == ["route_topk"] + assert _on(plan.slots(112, plan.Facts(mkm_m128=True))) == [] + wide = plan.Facts(route_topk=True, mkm_m256=True, mkm_m512=True) + assert _on(plan.slots(224, wide)) == ["mkm_m256", "route_topk"] + assert _on(plan.slots(448, wide)) == ["mkm_m512", "route_topk"] + assert _on(plan.slots(112, wide)) == ["route_topk"] + assert _on(plan.slots(896, wide)) == ["route_topk"] # the router front and the hop's mailbox are built for the widest width they serve assert plan.widest((7, 14, 28), plan.ROUTE_MAX_ROWS) == 14 assert plan.widest((7, 14, 28), plan.HOP_MAX_ROWS) == 14 @@ -296,9 +370,15 @@ def test_the_slot_plans_row_bounds_are_the_kernels(): from sgl_nvidia_mega.qwen3_5.sm103.moe import mkm, router assert plan.ROUTE_MAX_ROWS == router.MAX_TOKENS + assert plan.ROUTE_TOPK_MAX_ROWS == router.TOPK_MAX_ROWS assert plan.ROUTE_SHARED_ROWS == router.SHARED_MAX_M assert plan.MKM_MAX_ROWS == mkm.MAX_TOKENS assert plan.MKM_WIDE_MAX_ROWS == mkm.M16_MAX_TOKENS + assert plan.MKM_M32_MAX_ROWS == mkm.M32_MAX_TOKENS + assert plan.MKM_M64_MAX_ROWS == mkm.M64_MAX_TOKENS + assert plan.MKM_M128_MAX_ROWS == mkm.M128_MAX_TOKENS + assert plan.MKM_M256_MAX_ROWS == mkm.M256_MAX_TOKENS + assert plan.MKM_M512_MAX_ROWS == mkm.M512_MAX_TOKENS assert plan.GDN_PUSH_MAX_ROWS == hop.OPROJ_MAX_M assert plan.MKA_MAX_ROWS == mka.MAX_ROWS from sgl_nvidia_mega.qwen3_5.attention import mka as mka_step @@ -311,6 +391,9 @@ def test_a_numeric_slot_turns_the_claim_into_its_tolerances(): from sglang.kernels.plugin_kernels.nvidia.spec import NumericsKind assert spec.numerics(frozenset({"route", "hop"})).kind is NumericsKind.BITWISE + tiled_route = spec.numerics(frozenset({"route_topk", "route_gemm"})) + assert tiled_route.kind is NumericsKind.TOLERANCE + assert dict(tiled_route.tolerances) == spec.NUMERIC_TOLERANCES["route_gemm"] claim = spec.numerics(frozenset({"gdn_verify", "hop", "gdn_push"})) assert claim.kind is NumericsKind.TOLERANCE assert dict(claim.tolerances) == spec.NUMERIC_TOLERANCES["gdn_push"] diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 7bf31cb7..b44466bc 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -513,7 +513,7 @@ def combine(combine_input, log=log, combined=combined): "c(512,)", "g2(512,)", ), - (), + (("routing", None),), ), ("combine", "main"), ] From 5d9d24275cbefef0f46526848814db5925bb8ca8 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 15:26:40 -0700 Subject: [PATCH 07/62] perf: bake validated megakernel schedules --- sgl_nvidia_mega/qwen3_5/plan.py | 6 +- sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py | 65 ++++++++-- sgl_nvidia_mega/qwen3_5/sm103/moe/router.py | 25 ++-- .../qwen3_5/sm103/moe/src/mkm_moe.cu | 18 ++- .../qwen3_5/sm103/moe/src/mkm_moe_m32.cu | 114 ++++++++++++++---- .../qwen3_5/sm103/moe/src/mkm_moe_m64.cu | 73 ++++++++--- .../qwen3_5/sm103/moe/src/router_front.cu | 40 ++++-- sgl_nvidia_mega/qwen3_5/spec.py | 6 +- test/kernels/test_qwen3_5_moe.py | 12 +- test/unit/test_qwen3_5_spec.py | 30 ++++- 10 files changed, 309 insertions(+), 80 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index cea84f8a..5c54a948 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -71,12 +71,12 @@ MKA_WIDE_MAX_ROWS = MAX_ROWS # slots lowm / lowm_se: the rows up to which each low-M launch measured faster than SGLang's quantize + # cuBLAS path on GB300 (results/trt-1022/l6): the input projections (in_proj_qkvz + in_proj_ba, qkv_proj), -# the output projections (out_proj, o_proj) and the shared expert. The wider 28/ -# 56-row launches regressed the target step, so they stay on SGLang's stock path. +# the output projections (out_proj, o_proj) and the shared expert. The shared expert's +# 21-row launch also measured faster; the 28/56-row launches stay on SGLang's stock path. LOWM_IN_MAX_ROWS = 14 LOWM_OUTPROJ_MAX_ROWS = 14 LOWM_OPROJ_MAX_ROWS = 14 -LOWM_SE_MAX_ROWS = 14 +LOWM_SE_MAX_ROWS = 21 class Slots(msgspec.Struct, frozen=True): diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py b/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py index 5aa31d1e..a85df1e5 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py @@ -41,9 +41,19 @@ class Engine(NamedTuple): defines: tuple = () # extra compile options +# The validated TP8 schedule changes only when the weight and scale-factor loads issue; +# the engine's rows remain bitwise identical. +MKM7_DEFINES = ( + "-DSPEC_STAGES=3", + "-DPACE_D=4", + "-DFC2SF_PRE=0", +) + # per intermediate size a rank holds (TP8, TP4) ENGINES = { - 128: Engine("mkm_moe.cu", "q35_mkm_moe", 211968, "CU_TENSOR_MAP_SWIZZLE_64B"), + 128: Engine( + "mkm_moe.cu", "q35_mkm_moe", 211968, "CU_TENSOR_MAP_SWIZZLE_64B", MKM7_DEFINES + ), 256: Engine( "mkm_moe_tp4.cu", "q35_mkm_moe_tp4", 216064, "CU_TENSOR_MAP_SWIZZLE_128B" ), @@ -116,6 +126,36 @@ class Engine(NamedTuple): M256_MAX_TOKENS: 0, M512_MAX_TOKENS: 0, } +ONEPASS_MAX_ROWS = 42 + + +def engine_options(rows: int, plain: bool = False) -> tuple: + """Validated bitwise schedule defines for a wide engine. + + The 64-token engine keeps its stock schedule past ``ONEPASS_MAX_ROWS``. + """ + if rows == M32_MAX_TOKENS: + return ( + "-DTHR_D=2", + "-DTOPK_NOINLINE=1", + "-DONEPASS=1", + "-DMMA_SFONLY=1", + ) + if rows == M64_MAX_TOKENS: + return ("-DMMA_SFONLY=1",) + (() if plain else ("-DONEPASS=1",)) + return () + + +def reserved_sms(rows: int, m: int) -> int: + """The validated SM reserve for an engine bucket at an actual row count.""" + if rows == M32_MAX_TOKENS and m <= 21: + return 0 + if rows == M64_MAX_TOKENS: + if m <= 35: + return 32 + if m <= 42: + return 16 + return RESERVED_SMS[rows] class Params(ctypes.Structure): @@ -166,8 +206,15 @@ def engine_rows(m: int) -> int: return next(r for r in ENGINE_ROWS if m <= r) +def onepass_plain(rows: int, m: int) -> bool: + """Whether ``m`` rows use the 64-token engine's stock two-pass schedule.""" + return rows == M64_MAX_TOKENS and m > ONEPASS_MAX_ROWS + + @cache -def _kernel(device_index: int, intermediate: int, rows: int) -> jit.Kernel: +def _kernel( + device_index: int, intermediate: int, rows: int, plain: bool = False +) -> jit.Kernel: """The engine of ``engine_rows`` bound ``rows`` at ``intermediate`` per rank.""" engine = { MAX_TOKENS: ENGINES, @@ -188,7 +235,8 @@ def _kernel(device_index: int, intermediate: int, rows: int) -> jit.Kernel: "-default-device", f"-DQ35_MKM_SMEM_BYTES={engine.smem_bytes}", ) - + engine.defines, + + engine.defines + + engine_options(rows, plain), ) return jit.load(cubin, engine.kernel, (THREADS, 1, 1), engine.smem_bytes) @@ -254,6 +302,8 @@ def build( for hi in ENGINE_ROWS: if min_rows <= hi and max_rows > lo: _kernel(device_index, intermediate, hi) + if onepass_plain(hi, max(lo + 1, min(max_rows, hi))): + _kernel(device_index, intermediate, hi, True) lo = hi @@ -363,15 +413,14 @@ def mkm_moe( out=gemm2_out.data_ptr(), num_tokens=m, ) - kernel = _kernel(device.index, intermediate, rows) + kernel = _kernel(device.index, intermediate, rows, onepass_plain(rows, m)) stream = torch.cuda.current_stream().cuda_stream if rows in (M32_MAX_TOKENS, M64_MAX_TOKENS, M128_MAX_TOKENS) + WIDE_ROWS: # min(M * 10, SMs) CTAs, one per SM, each running its share of the distinct experts # (the wide engines: every SM, CTA c running experts c, c + G, ...) - sms = ( - torch.cuda.get_device_properties(device).multi_processor_count - - RESERVED_SMS[rows] - ) + sms = torch.cuda.get_device_properties( + device + ).multi_processor_count - reserved_sms(rows, m) ids_in = 0 if routing is None else routing[0].data_ptr() if rows in WIDE_ROWS and not ( routing[0].is_contiguous() and ids_in % 16 == 0 and sms >= WIDE_MIN_SMS diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py b/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py index d0eb1477..7ee38653 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py @@ -3,8 +3,9 @@ """The router GEMM and the NVFP4 quantization of its input in one launch, for the trtllm MoE runner, which routes from the logits itself (``src/router_front.cu``). -The logits equal SGLang's cuBLAS router GEMM bit for bit, and the quantization equals -the runner's own NVFP4 input quantizer (16-element blocks, e4m3 scales, linear layout). +The logits stay within one BF16 ULP of SGLang's cuBLAS router GEMM, and the +quantization equals the runner's own NVFP4 input quantizer (16-element blocks, e4m3 +scales, linear layout). ``router_gemm_quant_shared`` also computes the shared expert's static-FP8 gate/up projection of up to 8 rows in the same launch (SGLANG_Q35MF_ROUTE_SHARED on SGLang branch nvidia/qwen35-agentx-gb300); it sums in another order than the FP8 linear's GEMM @@ -32,6 +33,9 @@ _SH_PART_BYTES = 4 * 16 * 8 * 4 # fp32 [kShSplits 4][16][8] SHARED_MAX_M, SHARED_MAX_N = 8, CTAS * 8 # 8 shared rows a CTA MODE_QUANT = 1 +# The validated plain-launch schedule uses four K chains, one per warp, summed in +# chain order. The shared-stage launch keeps two chains for its larger footprint. +KSPLIT = 4 def max_m_bucket(m: int) -> int: @@ -39,11 +43,11 @@ def max_m_bucket(m: int) -> int: return 1 if m <= 1 else 8 if m <= 8 else 16 -def shared_bytes(max_m: int, shared: bool = False) -> int: +def shared_bytes(max_m: int, shared: bool = False, ksplit: int = 2) -> int: """GemvSmem::kBytes: 8 weight rows, max_m input rows, the fp32 partials - and two mbarriers; with the shared stage also 8 FP8 weight rows, max_m FP8 input rows, - its partials and a third mbarrier.""" - plain = 8 * _ROW_BYTES + max_m * _ROW_BYTES + _PART_BYTES + (``ksplit`` chains) and two mbarriers; with the shared stage also 8 FP8 weight rows, + max_m FP8 input rows, its partials and a third mbarrier.""" + plain = 8 * _ROW_BYTES + max_m * _ROW_BYTES + _PART_BYTES * ksplit // 2 if not shared: return plain + 16 return plain + (8 + max_m) * _SH_ROW_BYTES + _SH_PART_BYTES + 24 @@ -85,23 +89,24 @@ class Params(ctypes.Structure): @cache -def _cubin(device_index: int, max_m: int) -> bytes: +def _cubin(device_index: int, max_m: int, ksplit: int = 2) -> bytes: major, minor = torch.cuda.get_device_capability(device_index) return jit.compile_source( SOURCE.read_text(), f"q35_router_front_m{max_m}", f"sm_{major}{minor}a", - ("-std=c++17", "-default-device", f"-DQ35_ROUTER_MAX_M={max_m}"), + ("-std=c++17", "-default-device", f"-DQ35_ROUTER_MAX_M={max_m}") + + ((f"-DQ35_ROUTER_KSPLIT={ksplit}",) if ksplit != 2 else ()), ) @cache def _kernel(device_index: int, max_m: int) -> jit.Kernel: return jit.load( - _cubin(device_index, max_m), + _cubin(device_index, max_m, KSPLIT), "q35_router_gemm_quant", (THREADS, 1, 1), - shared_bytes(max_m), + shared_bytes(max_m, ksplit=KSPLIT), ) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe.cu index 83925119..959908cd 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe.cu @@ -212,6 +212,12 @@ constexpr int kARow = 132; // FC1 accumulator row stride in floats: 16 B aligne #endif #ifndef DIAG_NO_MMA // timing diagnostic only: no MMAs, stages are released as soon as they land #define DIAG_NO_MMA 0 +#endif +#ifndef SPEC_STAGES // ring stages warp 1 issues for its token's expert before the routing verdict (the rest after it) +#define SPEC_STAGES 5 +#endif +#ifndef PACE_D // > 0: the producer issues stage st only once stage st - PACE_D has landed (bounded bytes in flight) +#define PACE_D 0 #endif // partial row stride in floats: 16 B aligned, rows land on distinct banks struct Params { @@ -712,7 +718,7 @@ q35_mkm_moe(const __grid_constant__ CUtensorMap tm13, const __grid_constant__ CU if (lane == 0) { stamp(p, 9); const uint64_t pol = policy_evict_first(); - for (int st = 0; st < kStages; ++st) issue_stage(st, (int)espec, pol); + for (int st = 0; st < SPEC_STAGES; ++st) issue_stage(st, (int)espec, pol); stamp(p, 2); } mbar_wait(&s.route_done, 0); @@ -759,7 +765,7 @@ q35_mkm_moe(const __grid_constant__ CUtensorMap tm13, const __grid_constant__ CU if (e < 0) { // duplicate pair (both CTAs agree): let the speculative loads land; nothing was sent over DSMEM if (warp == 1) { if (lane == 0) - for (int b = 0; b < kStages; ++b) mbar_wait(&s.full[b], 0); + for (int b = 0; b < SPEC_STAGES; ++b) mbar_wait(&s.full[b], 0); __syncwarp(); } mbar_wait(&s.x_full, 0); @@ -774,10 +780,11 @@ q35_mkm_moe(const __grid_constant__ CUtensorMap tm13, const __grid_constant__ CU if (warp == 1) { if (lane == 0) { - // ---------------- TMA producer: the rest of the weight stream (stages 0 .. kStages - 1 went out above) + // ---------------- TMA producer: the rest of the weight stream (stages 0 .. SPEC_STAGES - 1 went out above) const uint64_t pol = policy_evict_first(); - for (int st = kStages; st < kTotStages; ++st) { - mbar_wait(&s.empty[st % kStages], ((st / kStages) - 1) & 1); + for (int st = SPEC_STAGES; st < kTotStages; ++st) { + if (st >= kStages) mbar_wait(&s.empty[st % kStages], ((st / kStages) - 1) & 1); + if (PACE_D > 0 && st >= PACE_D) mbar_wait(&s.full[(st - PACE_D) % kStages], ((st - PACE_D) / kStages) & 1); issue_stage(st, e, pol); } } @@ -1015,6 +1022,7 @@ q35_mkm_moe(const __grid_constant__ CUtensorMap tm13, const __grid_constant__ CU } static_assert(sizeof(Smem) + 1024 == Q35_MKM_SMEM_BYTES, "the launcher's dynamic shared memory size"); +static_assert(SPEC_STAGES >= 1 && SPEC_STAGES <= kStages && PACE_D >= 0 && PACE_D < kStages, "stage issue knobs"); } // namespace qwen_mkm_moe diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m32.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m32.cu index 1e066a50..ee770956 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m32.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m32.cu @@ -209,9 +209,12 @@ constexpr int kARow = 260; // FC1 accumulator row stride in floats: 16 B aligne #ifndef SPEC_STAGES // ring stages issued for the speculated expert before the routing verdict (rest: after it) #define SPEC_STAGES 4 #endif -#ifndef STAGE_STAMPS // DIAG: per-stage issue (slots 64 + st) and landing (16 + st, seen by warp 4) stamps +#ifndef STAGE_STAMPS // DIAG: per-stage issue (slots 64 + st), landing (16 + st, seen by warp 4) and MMA issue (40 + st) stamps #define STAGE_STAMPS 0 #endif +#ifndef STAMP_IT // DIAG: the item whose stages STAGE_STAMPS records +#define STAMP_IT 0 +#endif #ifndef DIAG_SPEC // DIAG: warp 1's speculative ranking, stamps 16 = logits loaded, 17 = keys, 18 = top-10 (stride 32) #define DIAG_SPEC 0 #endif @@ -249,12 +252,34 @@ static_assert(DRAIN_SMEM == 0, "x's shared memory holds the next item's rows whi #ifndef DIAG_ROUTE // DIAG: stamps 13 = warp 0's routing rounds done, 15 = their barrier, 14 = item 0's rows gathered #define DIAG_ROUTE 0 #endif +#ifndef MMA_SFONLY // 1: the MMA warp waits only on a stage's sf_full: its arrivers (the TMEM warpgroup) waited on the +#define MMA_SFONLY 0 // stage's full barrier before staging the scale factors, so the data has landed (one wait less per stage) +#endif +#ifndef ONEPASS // 1: an item of at most 8 rows runs FC2 in one pass (16 h copies of 8 rows: 16 tiles in one accumulator) +#define ONEPASS 0 +#endif +static_assert(!(ONEPASS && TAIL8), "ONEPASS drains its one FC2 pass on the TMEM warpgroup"); +#ifndef THR_D // a CTA with one item, in a grid where some CTAs have two, keeps at most THR_D ring stages in flight +#define THR_D 0 // (its HBM share goes to the CTAs on the critical path); 0 = off +#endif #ifndef FC2_PF // FC2 tiles (beyond the first kStages) prefetched to L2 ahead of the ring; 0 = off #define FC2_PF 0 #endif #ifndef FC2_PF_LEAD // issue the FC2 prefetch when the producer reaches FC1 stage kFc1Stages - FC2_PF_LEAD #define FC2_PF_LEAD 4 #endif +#ifndef FC2_PF_FROM // the first FC2 stage (0 .. 7) the FC2 prefetch covers (old: kStages, the stages beyond the ring) +#define FC2_PF_FROM 4 +#endif +#ifndef FC2_PF_IT0 // items from this index on take the FC2 prefetch +#define FC2_PF_IT0 0 +#endif +#ifndef NEXT_PF // the next item's first NEXT_PF stages prefetched to L2 when the producer reaches stage NEXT_PF_AT +#define NEXT_PF 0 +#endif +#ifndef NEXT_PF_AT +#define NEXT_PF_AT 16 +#endif struct Params { const uint16_t* logits; // [M, 512] bf16 router logits (router_gemm_quant output) @@ -568,7 +593,14 @@ __device__ __forceinline__ uint32_t warp_top10_redux(uint32_t (&keys)[16], uint3 return mine; } // The top-10 keys of one logits row (16 bf16 a lane in v0, v1), lane k < 10 holding the k-th largest +#ifndef TOPK_NOINLINE // 1: one out-of-line copy of the row top-10 shared by every routing warp (instruction cache) +#define TOPK_NOINLINE 0 +#endif +#if TOPK_NOINLINE +__device__ __noinline__ uint32_t row_top10(const uint4 v0, const uint4 v1, uint32_t lane, uint32_t* s_cand) { +#else __device__ __forceinline__ uint32_t row_top10(const uint4 v0, const uint4 v1, uint32_t lane, uint32_t* s_cand) { +#endif uint32_t keys[16]; #pragma unroll for (int i = 0; i < 8; ++i) { @@ -666,8 +698,9 @@ __device__ __forceinline__ void reg_fence64(uint32_t (&r)[64]) { // keep uses o // r-th row of the item, row info buffer rb) template __device__ __forceinline__ void fc2_drain16(const Params& p, Smem& s, uint32_t tm, int q, int lane, int pass, float a2, - int rb) { - const int row = q * 32 + lane, tile = pass * 8 + row / 16, r = row % 16; + int rb, bool one = false) { + // one: the item's single pass holds 16 tiles of 8 rows (accumulator row = 8 tile + item row) + const int row = q * 32 + lane, tile = one ? row / 8 : pass * 8 + row / 16, r = one ? row % 8 : row % 16; const bool valid = r < s.nrows[rb]; __nv_bfloat16* orow = valid ? p.out + (size_t)(s.rtok[rb][r] * kTopK + s.rkslot[rb][r]) * kH + tile * 256 : nullptr; const uint32_t base = tm + ((uint32_t)(q * 32) << 16); @@ -994,7 +1027,7 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ } if (lane == 0) { s.nitems = nall; - s.spec_ok = owner ? 1 : 0; + s.spec_ok = (owner ? 1 : 0) | (U > G ? 2 : 0); // bit 1: some CTA runs more than one expert stamp(p, 1); mbar_arrive(&s.route_done); } @@ -1013,7 +1046,8 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ return; } tc_fence_after(); - const bool spec_ok = s.spec_ok != 0; + const bool spec_ok = (s.spec_ok & 1) != 0; + const bool thr = THR_D > 0 && nitems == 1 && (s.spec_ok & 2) != 0; // ring stage of item 0's stage 0: the speculative fill (stages 0 .. SPEC_STAGES - 1) is item 0's own only if spec_ok const int base0 = spec_ok ? 0 : SPEC_STAGES; float ag = 0.f, sc = 0.f, a2 = 0.f; @@ -1043,9 +1077,9 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ // SwiGLU + NVFP4 epilogue (h = 0: the one FC1 pass), tasks (row, 16-wide block, half) = 16 x 8 x 2 = 256, // task = et + k * nt. // Accumulator columns of each 32-block: [0, 8) up 2m, [8, 16) gate 2m, [16, 24) up 2m + 1, [24, 32) gate 2m + 1 - auto epi_tasks = [&](int h, int et, int nt) { + auto epi_tasks = [&](int h, int et, int nt, bool one) { // one: rows 0..7 only, 16 copies of 8 rows #pragma unroll 1 - for (int task = et; task < 256; task += nt) { + for (int task = et; task < (one ? 128 : 256); task += nt) { const int row = task / 16, blk = (task / 2) % 8, hf = task % 2; const float4* ar = reinterpret_cast(s.acc1 + row * kARow + 32 * blk + 4 * hf); const float4 u_even = ar[0], g_even = ar[2], u_odd = ar[4], g_odd = ar[6]; @@ -1075,8 +1109,13 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ const int gblk = 8 * h + blk; // 16-wide block over the 128 intermediates; k-blocks 0, 1 const int chunk = gblk / 2, r8 = row % 8; const uint32_t off = (uint32_t)((row / 8) * 1024 + r8 * 128 + ((chunk ^ r8) * 16) + (gblk % 2) * 8 + hf * 4); + if (one) { #pragma unroll - for (int cp = 0; cp < 8; ++cp) *reinterpret_cast(s.ha + cp * 2048 + off) = packed; + for (int cp = 0; cp < 16; ++cp) *reinterpret_cast(s.ha + cp * 1024 + off) = packed; + } else { +#pragma unroll + for (int cp = 0; cp < 8; ++cp) *reinterpret_cast(s.ha + cp * 2048 + off) = packed; + } if (hf == 0) reinterpret_cast(&s.hsf[row][gblk / 4])[gblk % 4] = (uint8_t)sf; } }; @@ -1103,9 +1142,13 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ const int st0 = (it == 0 && spec_ok) ? SPEC_STAGES : 0; for (int st = st0; st < kTotStages; ++st) { const int g = base + st; - if (FC2_PF > 0 && st == kFc1Stages - FC2_PF_LEAD) { - // the FC2 tiles beyond the ring go to L2 now, while the last FC1 stages and h are in progress - for (int j = kStages; j < kStages + FC2_PF; ++j) prefetch_stage(kFc1Stages + j, e); + if (FC2_PF > 0 && st == kFc1Stages - FC2_PF_LEAD && it >= FC2_PF_IT0) { + // the FC2 tiles go to L2 now, while the last FC1 stages and h are in progress + for (int j = FC2_PF_FROM; j < FC2_PF_FROM + FC2_PF && j < kFc2Stages; ++j) prefetch_stage(kFc1Stages + j, e); + } + if (NEXT_PF > 0 && st == NEXT_PF_AT && it + 1 < nitems) { + // the next item's first stages go to L2 while this one finishes (their TMA loads then start from L2) + for (int k = 0; k < NEXT_PF; ++k) prefetch_stage(k, s.item_e[it + 1]); } if (PF_SET != 0 && ((st == st0 && (PF_SET & 1)) || (st == kFc1Stages + 4 && (PF_SET & 4)) || (st == kFc1Stages + 6 && (PF_SET & 8)))) { @@ -1114,8 +1157,9 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ for (int k = st; k < st + n && k < kTotStages; ++k) prefetch_stage(k, e); } if (g >= kStages) mbar_wait(&s.empty[g % kStages], ((g / kStages) - 1) & 1); + if (thr && g >= THR_D) mbar_wait(&s.full[(g - THR_D) % kStages], ((g - THR_D) / kStages) & 1); issue_stage(g, st, e, pol); - if (STAGE_STAMPS && it == 0) stamp(p, 64 + st); + if (STAGE_STAMPS && it == STAMP_IT) stamp(p, 64 + st); } } } @@ -1136,14 +1180,16 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ if (it > 0) mbar_wait(&s.acc_drained, (it - 1) & 1); mbar_wait(&s.x_full, par); tc_fence_after(); + const bool one = ONEPASS && s.nrows[par] <= 8; // FC2 in one pass of 16 tiles int n = 0; for (int st = 0; st < kFc1Stages; ++st) { const int g = base + st, b = g % kStages, ks = st; - mbar_wait(&s.full[b], (g / kStages) & 1); + if (!MMA_SFONLY) mbar_wait(&s.full[b], (g / kStages) & 1); mbar_wait(&s.sf_full[b], (g / kStages) & 1); tc_fence_after(); if (st == 0 && it == 0) stamp(p, 3); if (st == 0 && it > 0 && it == nitems - 1) stamp(p, 11); + if (STAGE_STAMPS && it == STAMP_IT) stamp(p, 40 + st); #pragma unroll for (int q = 0; q < 4; ++q, ++n) { const int kb = ks * 4 + q; // k-blocks in K order: one accumulator, the stock summation order @@ -1160,13 +1206,14 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ if (it == 0) stamp(p, 6); for (int st = kFc1Stages; st < kTotStages; ++st) { const int g = base + st, b = g % kStages, f = st - kFc1Stages; - if (f == 4) { mbar_wait(&s.fc2_free, par); tc_fence_after(); } - mbar_wait(&s.full[b], (g / kStages) & 1); + if (f == 4 && !one) { mbar_wait(&s.fc2_free, par); tc_fence_after(); } + if (!MMA_SFONLY) mbar_wait(&s.full[b], (g / kStages) & 1); mbar_wait(&s.sf_full[b], (g / kStages) & 1); tc_fence_after(); + if (STAGE_STAMPS && it == STAMP_IT) stamp(p, 40 + st); #pragma unroll for (int tl = 0; tl < 2; ++tl) { - const int j = 2 * f + tl, jt = j % 8; + const int j = 2 * f + tl, jt = one ? j : j % 8; #pragma unroll for (int kb = 0; kb < 2; ++kb) { const int n2 = 64 + j * 2 + kb; // scale-factor slot, numbered on from FC1's 64 k-blocks @@ -1177,7 +1224,11 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ } } tc_commit(&s.empty[b]); - if (f % 4 == 3) tc_commit(&s.fc2_pass[f / 4]); + if (one) { // both pass barriers complete once per item either way + if (f == kFc2Stages - 1) { tc_commit(&s.fc2_pass[0]); tc_commit(&s.fc2_pass[1]); } + } else if (f % 4 == 3) { + tc_commit(&s.fc2_pass[f / 4]); + } } } } @@ -1200,7 +1251,7 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ } } mbar_wait(&s.acc_ready, par); - epi_tasks(0, et, 192); + epi_tasks(0, et, 192, ONEPASS && s.nrows[par] <= 8); if (warp == 3 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 14); fence_proxy_async_smem(); named_bar(6, 192); // h and its scale words are complete for the TMEM warpgroup @@ -1225,7 +1276,7 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ auto stage_sf2 = [&](int st) { // FC2 tiles 2 f, 2 f + 1: k-blocks 0, 1 of their row blocks (rows 0..127 | 128..255) const int g = base + st, b = g % kStages; mbar_wait(&s.full[b], (g / kStages) & 1); - if (STAGE_STAMPS && threadIdx.x == 128 && it == 0) stamp(p, 16 + st); + if (STAGE_STAMPS && threadIdx.x == 128 && it == STAMP_IT) stamp(p, 16 + st); if (g >= 4) mbar_wait(&s.empty[(g - 4) % kStages], ((g - 4) / kStages) & 1); const uint8_t* sf = s.ring[b] + 32768; const int n0 = 64 + (st - kFc1Stages) * 4; @@ -1247,7 +1298,7 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ for (int st = 0; st < kFc1Stages; ++st) { const int g = base + st, b = g % kStages, ks = st; mbar_wait(&s.full[b], (g / kStages) & 1); - if (STAGE_STAMPS && threadIdx.x == 128 && it == 0) stamp(p, 16 + st); + if (STAGE_STAMPS && threadIdx.x == 128 && it == STAMP_IT) stamp(p, 16 + st); if (g >= 4) mbar_wait(&s.empty[(g - 4) % kStages], ((g - 4) / kStages) & 1); const uint8_t* sf = s.ring[b] + 32768; #pragma unroll @@ -1277,22 +1328,25 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ drain_fc1(2); if (lane == 0) mbar_arrive(&s.acc_ready); } + const bool one = ONEPASS && nrows <= 8; // FC2 in one pass of 16 tiles (8 rows each) mbar_wait(&s.acc_ready, par); - epi_tasks(0, 64 + (warp - 4) * 32 + lane, 192); + epi_tasks(0, 64 + (warp - 4) * 32 + lane, 192, one); fence_proxy_async_smem(); named_bar(6, 192); // h complete if (threadIdx.x == 128 && it == 0) stamp(p, 5); // h scale factors: 16 atoms (tile jt, k-block kb) of 4 columns; row m = 32 c + lane is nonzero only in its tile + // (one: 32 atoms, 16 tiles of 8 rows, columns kColSfa .. kColSfa + 127) { uint32_t r[32]; + const int rr = one ? 8 : 16, nh = one ? 4 : 2; #pragma unroll 1 - for (int half = 0; half < 2; ++half) { + for (int half = 0; half < nh; ++half) { #pragma unroll for (int i = 0; i < 32; ++i) { const int a = half * 8 + i / 4, col = i % 4; // atom a = 2 jt + kb const int jt = a / 2, kb = a % 2; const int row = 32 * col + lane; - r[i] = (row / 16 == jt && row % 16 < nrows) ? s.hsf[row % 16][kb] : 0u; + r[i] = (row / rr == jt && row % rr < nrows) ? s.hsf[row % rr][kb] : 0u; } tmem_st_x32(tm + lane_addr + kColSfa + half * 32, r); } @@ -1302,6 +1356,20 @@ q35_mkm_moe_m32(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ __syncwarp(); if (lane == 0) mbar_arrive(&s.h_ready); if (q == 0 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 15); + if (one) { // every FC2 stage's scale factors, then the one pass's drain + for (int st = kFc1Stages + EARLY_SF2; st < kTotStages; ++st) stage_sf2(st); + mbar_wait(&s.fc2_pass[0], par); + tc_fence_after(); + if (q == 0 && lane == 0 && it == 0) stamp(p, 12); + fc2_drain16<0, 4>(p, s, tm, q, lane, 0, a2, rb, true); + tc_fence_before(); + __syncwarp(); + if (lane == 0) { + mbar_arrive(&s.fc2_free); // once per item either way (the MMA warp's parity) + mbar_arrive(&s.acc_drained); + } + continue; + } for (int st = kFc1Stages + EARLY_SF2; st < kTotStages; ++st) { if (st == kFc1Stages + 4) { // pass 0 is fully staged; drain it once its MMAs are done, then free the accumulator for pass 1 diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m64.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m64.cu index 544757e9..9fac5e87 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m64.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_m64.cu @@ -254,6 +254,16 @@ static_assert(DRAIN_SMEM == 0, "x's shared memory holds the next item's rows whi #ifndef SPEC_PF // stages of the speculated expert past the first ring fill prefetched to L2 while the routing runs #define SPEC_PF 4 #endif +#ifndef MMA_SFONLY // 1: the MMA warp waits only on a stage's sf_full: its arrivers (the TMEM warpgroup) waited on the +#define MMA_SFONLY 0 // stage's full barrier before staging the scale factors, so the data has landed (one wait less per stage) +#endif +#ifndef ONEPASS // 1: an item of at most 8 rows runs FC2 in one pass (16 h copies of 8 rows: 16 tiles in one accumulator) +#define ONEPASS 0 +#endif +static_assert(!(ONEPASS && TAIL8), "ONEPASS drains its one FC2 pass on the TMEM warpgroup"); +#ifndef THR_D // a CTA with one item, in a grid where some CTAs have two, keeps at most THR_D ring stages in flight +#define THR_D 0 // (its HBM share goes to the CTAs on the critical path); 0 = off +#endif #ifndef FC2_PF // FC2 tiles (beyond the first kStages) prefetched to L2 ahead of the ring; 0 = off #define FC2_PF 0 #endif @@ -677,8 +687,9 @@ __device__ __forceinline__ void reg_fence64(uint32_t (&r)[64]) { // keep uses o // r-th row of the item, row info buffer rb) template __device__ __forceinline__ void fc2_drain16(const Params& p, Smem& s, uint32_t tm, int q, int lane, int pass, float a2, - int rb) { - const int row = q * 32 + lane, tile = pass * 8 + row / 16, r = row % 16; + int rb, bool one = false) { + // one: the item's single pass holds 16 tiles of 8 rows (accumulator row = 8 tile + item row) + const int row = q * 32 + lane, tile = one ? row / 8 : pass * 8 + row / 16, r = one ? row % 8 : row % 16; const bool valid = r < s.nrows[rb]; __nv_bfloat16* orow = valid ? p.out + (size_t)(s.rtok[rb][r] * kTopK + s.rkslot[rb][r]) * kH + tile * 256 : nullptr; const uint32_t base = tm + ((uint32_t)(q * 32) << 16); @@ -1090,7 +1101,7 @@ q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ const int nall = ni; if (lane == 0) { s.nitems = nall; - s.spec_ok = owner ? 1 : 0; + s.spec_ok = (owner ? 1 : 0) | (O > G ? 2 : 0); // bit 1: some CTA runs more than one unit stamp(p, 1); if (DIAG_ROUTE && p.trace) { // DIAG values (not times): 22 = items, 23 = speculation kept p.trace[blockIdx.x * kTraceStride + 22] = nall; @@ -1113,7 +1124,8 @@ q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ return; } tc_fence_after(); - const bool spec_ok = s.spec_ok != 0; + const bool spec_ok = (s.spec_ok & 1) != 0; + const bool thr = THR_D > 0 && nitems == 1 && (s.spec_ok & 2) != 0; // ring stage of item 0's stage 0: the speculative fill (stages 0 .. SPEC_STAGES - 1) is item 0's own only if spec_ok const int base0 = spec_ok ? 0 : SPEC_STAGES; float ag = 0.f, sc = 0.f, a2 = 0.f; @@ -1143,9 +1155,9 @@ q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ // SwiGLU + NVFP4 epilogue (h = 0: the one FC1 pass), tasks (row, 16-wide block, half) = 16 x 8 x 2 = 256, // task = et + k * nt. // Accumulator columns of each 32-block: [0, 8) up 2m, [8, 16) gate 2m, [16, 24) up 2m + 1, [24, 32) gate 2m + 1 - auto epi_tasks = [&](int h, int et, int nt) { + auto epi_tasks = [&](int h, int et, int nt, bool one) { // one: rows 0..7 only, 16 copies of 8 rows #pragma unroll 1 - for (int task = et; task < 256; task += nt) { + for (int task = et; task < (one ? 128 : 256); task += nt) { const int row = task / 16, blk = (task / 2) % 8, hf = task % 2; const float4* ar = reinterpret_cast(s.acc1 + row * kARow + 32 * blk + 4 * hf); const float4 u_even = ar[0], g_even = ar[2], u_odd = ar[4], g_odd = ar[6]; @@ -1175,8 +1187,13 @@ q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ const int gblk = 8 * h + blk; // 16-wide block over the 128 intermediates; k-blocks 0, 1 const int chunk = gblk / 2, r8 = row % 8; const uint32_t off = (uint32_t)((row / 8) * 1024 + r8 * 128 + ((chunk ^ r8) * 16) + (gblk % 2) * 8 + hf * 4); + if (one) { #pragma unroll - for (int cp = 0; cp < 8; ++cp) *reinterpret_cast(s.ha + cp * 2048 + off) = packed; + for (int cp = 0; cp < 16; ++cp) *reinterpret_cast(s.ha + cp * 1024 + off) = packed; + } else { +#pragma unroll + for (int cp = 0; cp < 8; ++cp) *reinterpret_cast(s.ha + cp * 2048 + off) = packed; + } if (hf == 0) reinterpret_cast(&s.hsf[row][gblk / 4])[gblk % 4] = (uint8_t)sf; } }; @@ -1214,6 +1231,7 @@ q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ for (int k = st; k < st + n && k < kTotStages; ++k) prefetch_stage(k, e); } if (g >= kStages) mbar_wait(&s.empty[g % kStages], ((g / kStages) - 1) & 1); + if (thr && g >= THR_D) mbar_wait(&s.full[(g - THR_D) % kStages], ((g - THR_D) / kStages) & 1); issue_stage(g, st, e, pol); if (STAGE_STAMPS && it == 0) stamp(p, 64 + st); } @@ -1236,10 +1254,11 @@ q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ if (it > 0) mbar_wait(&s.acc_drained, (it - 1) & 1); mbar_wait(&s.x_full, par); tc_fence_after(); + const bool one = ONEPASS && s.nrows[par] <= 8; // FC2 in one pass of 16 tiles int n = 0; for (int st = 0; st < kFc1Stages; ++st) { const int g = base + st, b = g % kStages, ks = st; - mbar_wait(&s.full[b], (g / kStages) & 1); + if (!MMA_SFONLY) mbar_wait(&s.full[b], (g / kStages) & 1); mbar_wait(&s.sf_full[b], (g / kStages) & 1); tc_fence_after(); if (st == 0 && it == 0) stamp(p, 3); @@ -1262,13 +1281,13 @@ q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ if (it == 0) stamp(p, 6); for (int st = kFc1Stages; st < kTotStages; ++st) { const int g = base + st, b = g % kStages, f = st - kFc1Stages; - if (f == 4) { mbar_wait(&s.fc2_free, par); tc_fence_after(); } - mbar_wait(&s.full[b], (g / kStages) & 1); + if (f == 4 && !one) { mbar_wait(&s.fc2_free, par); tc_fence_after(); } + if (!MMA_SFONLY) mbar_wait(&s.full[b], (g / kStages) & 1); mbar_wait(&s.sf_full[b], (g / kStages) & 1); tc_fence_after(); #pragma unroll for (int tl = 0; tl < 2; ++tl) { - const int j = 2 * f + tl, jt = j % 8; + const int j = 2 * f + tl, jt = one ? j : j % 8; #pragma unroll for (int kb = 0; kb < 2; ++kb) { const int n2 = 64 + j * 2 + kb; // scale-factor slot, numbered on from FC1's 64 k-blocks @@ -1279,7 +1298,11 @@ q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ } } tc_commit(&s.empty[b]); - if (f % 4 == 3) tc_commit(&s.fc2_pass[f / 4]); + if (one) { // both pass barriers complete once per item either way + if (f == kFc2Stages - 1) { tc_commit(&s.fc2_pass[0]); tc_commit(&s.fc2_pass[1]); } + } else if (f % 4 == 3) { + tc_commit(&s.fc2_pass[f / 4]); + } } } } @@ -1302,7 +1325,7 @@ q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ } } mbar_wait(&s.acc_ready, par); - epi_tasks(0, et, 192); + epi_tasks(0, et, 192, ONEPASS && s.nrows[par] <= 8); if (warp == 3 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 14); fence_proxy_async_smem(); named_bar(6, 192); // h and its scale words are complete for the TMEM warpgroup @@ -1379,22 +1402,24 @@ q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ drain_fc1(2); if (lane == 0) mbar_arrive(&s.acc_ready); } + const bool one = ONEPASS && nrows <= 8; // FC2 in one pass of 16 tiles (8 rows each) mbar_wait(&s.acc_ready, par); - epi_tasks(0, 64 + (warp - 4) * 32 + lane, 192); + epi_tasks(0, 64 + (warp - 4) * 32 + lane, 192, one); fence_proxy_async_smem(); named_bar(6, 192); // h complete if (threadIdx.x == 128 && it == 0) stamp(p, 5); // h scale factors: 16 atoms (tile jt, k-block kb) of 4 columns; row m = 32 c + lane is nonzero only in its tile - { + { // (one: 32 atoms, 16 tiles of 8 rows, columns kColSfa .. kColSfa + 127) uint32_t r[32]; + const int rr = one ? 8 : 16, nh = one ? 4 : 2; #pragma unroll 1 - for (int half = 0; half < 2; ++half) { + for (int half = 0; half < nh; ++half) { #pragma unroll for (int i = 0; i < 32; ++i) { const int a = half * 8 + i / 4, col = i % 4; // atom a = 2 jt + kb const int jt = a / 2, kb = a % 2; const int row = 32 * col + lane; - r[i] = (row / 16 == jt && row % 16 < nrows) ? s.hsf[row % 16][kb] : 0u; + r[i] = (row / rr == jt && row % rr < nrows) ? s.hsf[row % rr][kb] : 0u; } tmem_st_x32(tm + lane_addr + kColSfa + half * 32, r); } @@ -1404,6 +1429,20 @@ q35_mkm_moe_m64(const __grid_constant__ CUtensorMap tm13, const __grid_constant_ __syncwarp(); if (lane == 0) mbar_arrive(&s.h_ready); if (q == 0 && lane == 0 && it == 0 && !DIAG_ROUTE) stamp(p, 15); + if (one) { // every FC2 stage's scale factors, then the one pass's drain + for (int st = kFc1Stages + EARLY_SF2; st < kTotStages; ++st) stage_sf2(st); + mbar_wait(&s.fc2_pass[0], par); + tc_fence_after(); + if (q == 0 && lane == 0 && it == 0) stamp(p, 12); + fc2_drain16<0, 4>(p, s, tm, q, lane, 0, a2, rb, true); + tc_fence_before(); + __syncwarp(); + if (lane == 0) { + mbar_arrive(&s.fc2_free); // once per item either way (the MMA warp's parity) + mbar_arrive(&s.acc_drained); + } + continue; + } for (int st = kFc1Stages + EARLY_SF2; st < kTotStages; ++st) { if (st == kFc1Stages + 4) { // pass 0 is fully staged; drain it once its MMAs are done, then free the accumulator for pass 1 diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/router_front.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/router_front.cu index c82c8893..23d458fe 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/router_front.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/router_front.cu @@ -685,7 +685,14 @@ SGL_DEVICE void topk_and_tables(const Params& p, uint32_t* s_mask, int16_t* s_ti // Per CTA: kExpertsPerCta experts as the n8 tile, tokens as the m16 tile (rows // >= M repeat token 0); warp 0 runs the chain of the first half, warp 1 the // second, operands from shared memory. -inline constexpr uint32_t kSplitK = kHidden / 2; +// Q35_ROUTER_KSPLIT > 2 (numeric): K split into that many chains, one per warp 0 .. KSPLIT - 1, their fp32 partials +// added in split order (another order than cuBLAS's); the CTAs of the token rows quantize after the chains at 8. +#ifndef Q35_ROUTER_KSPLIT +#define Q35_ROUTER_KSPLIT 2 +#endif +inline constexpr uint32_t kKSplit = Q35_ROUTER_KSPLIT; +static_assert(kKSplit == 2 || kKSplit == 4 || kKSplit == 8, "router K splits"); +inline constexpr uint32_t kSplitK = kHidden / kKSplit; inline constexpr uint32_t kChainSteps = kSplitK / 16; // Shared-memory rows are padded by 16 B so the 8 rows of an ldmatrix hit distinct banks. inline constexpr uint32_t kRowBytes = kHidden * 2 + 16; @@ -700,8 +707,8 @@ template struct GemvSmem { static constexpr uint32_t kW = 0; static constexpr uint32_t kX = kW + kExpertsPerCta * kRowBytes; - static constexpr uint32_t kPart = kX + kMaxM * kRowBytes; // fp32 [2][16][kExpertsPerCta] - static constexpr uint32_t kShW = kPart + 2 * 16 * kExpertsPerCta * 4; + static constexpr uint32_t kPart = kX + kMaxM * kRowBytes; // fp32 [kKSplit][16][kExpertsPerCta] + static constexpr uint32_t kShW = kPart + kKSplit * 16 * kExpertsPerCta * 4; static constexpr uint32_t kXq = kShW + (kShared ? 8 * kShRowBytes : 0); static constexpr uint32_t kShPart = kXq + (kShared ? kMaxM * kShRowBytes : 0); // fp32 [kShSplits][16][8] static constexpr uint32_t kBarW = kShPart + (kShared ? kShSplits * 16 * 8 * 4 : 0); @@ -972,7 +979,7 @@ SGL_DEVICE void router_front_kernel(const Params& p) { } } - if (warp < 2) { + if (warp < kKSplit) { mbar_wait(bar_w, 0); mbar_wait(bar_x, 0); trace_stamp(p, tb + 2); @@ -980,10 +987,10 @@ SGL_DEVICE void router_front_kernel(const Params& p) { } else if (has_sh) { shared_gate_up(p, M, smem, bar_w, bar_x); } else if ((p.mode & kModeQuant) && blockIdx.x < M) { - // CTA b < M quantizes token row b while warps 0 and 1 run the GEMV chains. + // CTA b < M quantizes token row b while warps 0 .. kKSplit - 1 run the GEMV chains. mbar_wait(bar_x, 0); const uint8_t* row = smem + L::kX + blockIdx.x * kRowBytes; - for (uint32_t blk = threadIdx.x - 64; blk < kBlocksPerRow; blk += kThreads - 64) { + for (uint32_t blk = threadIdx.x - 32 * kKSplit; blk < kBlocksPerRow; blk += kThreads - 32 * kKSplit) { const uint4* src = reinterpret_cast(row + blk * kSfVec * 2); uint32_t sf; uint2 packed; @@ -997,9 +1004,22 @@ SGL_DEVICE void router_front_kernel(const Params& p) { if (threadIdx.x < M * kExpertsPerCta) { const uint32_t m = threadIdx.x / kExpertsPerCta; const uint32_t e = threadIdx.x % kExpertsPerCta; - const float lo = s_part[m * kExpertsPerCta + e]; - const float hi = s_part[(16 + m) * kExpertsPerCta + e]; - p.logits[(t0 + m) * kExperts + expert0 + e] = __float2bfloat16_rn(__fadd_rn(lo, hi)); + float acc = s_part[m * kExpertsPerCta + e]; +#pragma unroll + for (uint32_t k = 1; k < kKSplit; ++k) acc = __fadd_rn(acc, s_part[(k * 16 + m) * kExpertsPerCta + e]); + p.logits[(t0 + m) * kExperts + expert0 + e] = __float2bfloat16_rn(acc); + } + if (kKSplit == 8 && !kShared && (p.mode & kModeQuant) && blockIdx.x < M) { + // every warp ran a GEMV chain: CTA b < M quantizes token row b now + const uint8_t* row = smem + L::kX + blockIdx.x * kRowBytes; + for (uint32_t blk = threadIdx.x; blk < kBlocksPerRow; blk += kThreads) { + const uint4* src = reinterpret_cast(row + blk * kSfVec * 2); + uint32_t sf; + uint2 packed; + quantize_block(src[0], src[1], p.global_scale[0], sf, packed); + *reinterpret_cast(p.x_fp4 + static_cast(t0 + blockIdx.x) * (kHidden / 2) + blk * 8) = packed; + p.x_sf[static_cast(t0 + blockIdx.x) * kBlocksPerRow + blk] = static_cast(sf); + } } if (has_sh && threadIdx.x < M * 8) { const uint32_t m = threadIdx.x / 8; @@ -1045,7 +1065,7 @@ extern "C" __global__ __launch_bounds__(sglang::qwen_router_front::kThreads, 1) sglang::qwen_router_front::router_front_kernel(p); } -#if Q35_ROUTER_MAX_M <= 8 +#if Q35_ROUTER_MAX_M <= 8 && Q35_ROUTER_KSPLIT == 2 // The same plus the shared expert's static-FP8 gate/up projection of the rows (sh_* in Params) // in the same launch: CTA c computes its 8 rows while warps 0 and 1 run the router GEMV. extern "C" __global__ __launch_bounds__(sglang::qwen_router_front::kThreads, 1) void q35_router_gemm_quant_shared( diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index 7ccd089b..27b75d17 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -106,10 +106,12 @@ ) # Against SGLang's forward on the same GPU (test/e2e/test_qwen3_5_stock_parity.py): the -# stock step and the bitwise slots (route, mkm, gdn_verify, gdn_pdl, hop) reproduce it -# bit for bit. A numeric slot's outputs are held to fp64 on the same quantized operands +# stock step and the bitwise slots (mkm, gdn_verify, gdn_pdl, hop) reproduce it bit for +# bit. A numeric slot's outputs are held to stock or fp64 on the same quantized operands # (test/kernels), at most these multiples of the error of SGLang's own op. NUMERIC_TOLERANCES = { + # Four router GEMV chains sum in another order than SGLang's cuBLAS GEMM. + "route": {"route.router_logits.bf16_ulp_vs_stock": 1.0}, # The tiled router front stays within one BF16 ULP of SGLang's whole-M GEMM. "route_gemm": {"route_gemm.router_logits.bf16_ulp_vs_stock": 1.0}, "gdn_push": { diff --git a/test/kernels/test_qwen3_5_moe.py b/test/kernels/test_qwen3_5_moe.py index 6cf4d05c..106258c3 100644 --- a/test/kernels/test_qwen3_5_moe.py +++ b/test/kernels/test_qwen3_5_moe.py @@ -298,12 +298,22 @@ def _router_stock(x, w, scale): return logits, fp4.view(torch.uint8).reshape(x.shape[0], -1), sf.view(torch.uint8) +def _bf16_ulps(a, b): + """Maximum BF16 ULP distance using ordered sign-magnitude bits.""" + ia = a.contiguous().view(torch.int16).int() + ib = b.contiguous().view(torch.int16).int() + ia = torch.where(ia < 0, -(ia & 0x7FFF), ia) + ib = torch.where(ib < 0, -(ib & 0x7FFF), ib) + return int((ia - ib).abs().max()) + + @pytest.mark.parametrize("m", list(range(1, router.MAX_TOKENS + 1))) def test_router_gemm_quant_is_the_stock_router_gemm_and_quantizer(m): x, w, scale = _router_inputs(m, seed=m) got = router.router_gemm_quant(x, w, scale) want = _router_stock(x, w, scale) - for name, a, b in zip(("logits", "x_fp4", "x_sf"), got, want): + assert _bf16_ulps(got[0], want[0]) <= 1 + for name, a, b in zip(("x_fp4", "x_sf"), got[1:], want[1:]): assert torch.equal(a.view(torch.uint8), b.reshape(a.shape).view(torch.uint8)), ( name ) diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 48938560..fbebd160 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -271,6 +271,10 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): "mka_push", "route", ] + # Three requests keep only the shared expert's validated low-M launch. + three = plan.slots(21, every) + assert three.lowm_se + assert not (three.lowm_in or three.lowm_outproj or three.lowm_oproj) # The 28/56-row low-M paths stay stock; GDN core and MKA serve through 224. for rows in (28, 56, 112, 224): expected = [ @@ -387,10 +391,34 @@ def test_the_slot_plans_row_bounds_are_the_kernels(): assert mka_step.FUSED_MAX_ROWS <= mka.WIDE_MAX_ROWS +def test_the_validated_moe_schedules_are_shape_closed(): + from sgl_nvidia_mega.qwen3_5.sm103.moe import mkm + + assert mkm.MKM7_DEFINES == ( + "-DSPEC_STAGES=3", + "-DPACE_D=4", + "-DFC2SF_PRE=0", + ) + assert mkm.engine_options(32) == ( + "-DTHR_D=2", + "-DTOPK_NOINLINE=1", + "-DONEPASS=1", + "-DMMA_SFONLY=1", + ) + assert mkm.engine_options(64) == ("-DMMA_SFONLY=1", "-DONEPASS=1") + assert mkm.engine_options(64, plain=True) == ("-DMMA_SFONLY=1",) + assert [mkm.reserved_sms(32, rows) for rows in (21, 22)] == [0, 32] + assert [mkm.reserved_sms(64, rows) for rows in (35, 42, 43)] == [32, 16, 0] + assert not mkm.onepass_plain(64, 42) + assert mkm.onepass_plain(64, 43) + + def test_a_numeric_slot_turns_the_claim_into_its_tolerances(): from sglang.kernels.plugin_kernels.nvidia.spec import NumericsKind - assert spec.numerics(frozenset({"route", "hop"})).kind is NumericsKind.BITWISE + route = spec.numerics(frozenset({"route", "hop"})) + assert route.kind is NumericsKind.TOLERANCE + assert dict(route.tolerances) == spec.NUMERIC_TOLERANCES["route"] tiled_route = spec.numerics(frozenset({"route_topk", "route_gemm"})) assert tiled_route.kind is NumericsKind.TOLERANCE assert dict(tiled_route.tolerances) == spec.NUMERIC_TOLERANCES["route_gemm"] From 75cfba9b3212adafb5eda910f383a18dd0eb7709 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 15:29:33 -0700 Subject: [PATCH 08/62] feat: add opt-in balanced wide MoE scheduling --- sgl_nvidia_mega/qwen3_5/moe/__init__.py | 9 ++- sgl_nvidia_mega/qwen3_5/moe/mkm.py | 12 ++- sgl_nvidia_mega/qwen3_5/moe/route.py | 5 ++ sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 18 ++++- sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py | 47 ++++++++--- .../qwen3_5/sm103/moe/src/mkm_moe_wide.cu | 78 +++++++++++++++++++ test/kernels/test_qwen3_5_moe.py | 33 +++++--- test/unit/test_qwen3_5_runtime.py | 2 + test/unit/test_qwen3_5_spec.py | 41 +++++++++- 9 files changed, 218 insertions(+), 27 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/moe/__init__.py b/sgl_nvidia_mega/qwen3_5/moe/__init__.py index 07dcf050..e3c46dd1 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/moe/__init__.py @@ -124,7 +124,14 @@ def build(model, ctx) -> route.Bound: ) > (plan.MKM_M256_MAX_ROWS) if m32 or m64 or m128 or m256 or m512: bound = route.build_m32( - model, bound, m32=m32, m64=m64, m128=m128, m256=m256, m512=m512 + model, + bound, + m32=m32, + m64=m64, + m128=m128, + m256=m256, + m512=m512, + wide_bal="mkm_wide_bal" in SLOTS, ) return bound diff --git a/sgl_nvidia_mega/qwen3_5/moe/mkm.py b/sgl_nvidia_mega/qwen3_5/moe/mkm.py index e3a1fa0a..85333b29 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/mkm.py +++ b/sgl_nvidia_mega/qwen3_5/moe/mkm.py @@ -161,7 +161,16 @@ def routing_unsupported(mlp) -> str: return "" -def routed(mlp, maps: tuple, identity: torch.Tensor, logits, x_fp4, x_sf, routing=None): +def routed( + mlp, + maps: tuple, + identity: torch.Tensor, + logits, + x_fp4, + x_sf, + routing=None, + wide_bal: bool = False, +): """The deferred-finalize output of ``mlp``'s routed experts from the router front's logits and quantized input (``maps``: the layer's weight TMA maps; ``identity``: 0, 1, ... int32; ``routing``: route_topk's ids and weights, for the 64-token engine).""" @@ -184,6 +193,7 @@ def routed(mlp, maps: tuple, identity: torch.Tensor, logits, x_fp4, x_sf, routin experts.g1_scale_c, experts.g2_alphas, routing=routing, + wide_bal=wide_bal, ) deferred = FlashInferTrtllmDeferredFinalizeOutput( gemm2_out=gemm2_out, diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index bb430398..c4d35b5a 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -131,6 +131,7 @@ class Bound(msgspec.Struct, frozen=True): m128: bool = False # slot mkm_m128: the 128-token engine built (on m32's maps) m256: bool = False # slot mkm_m256: the 256-token engine built (on m32's maps) m512: bool = False # slot mkm_m512: the 512-token engine built (on m32's maps) + wide_bal: bool = False # opt-in balanced work split for mkm_m256 / mkm_m512 reason: str = "" # why a requested slot is not built @@ -201,6 +202,7 @@ def build_m32( m128: bool = False, m256: bool = False, m512: bool = False, + wide_bal: bool = False, ) -> Bound: """Slots mkm_m32 / mkm_m64 / mkm_m128 / mkm_m256 / mkm_m512 on ``bound`` (route_topk built): the 32 / 64 / 128 / 256 / 512-token engines and every MoE layer's weight maps @@ -226,6 +228,7 @@ def build_m32( intermediate, hi, min_rows=mkm_kernel.ENGINE_ROWS[bounds.index(lo) + 1] + 1, + wide_bal=wide_bal, ) return msgspec.structs.replace( bound, @@ -233,6 +236,7 @@ def build_m32( m128=m128, m256=m256, m512=m512, + wide_bal=wide_bal and (m256 or m512), m32={ mlp.layer_id: bound.mkm.get(mlp.layer_id) or mkm_kernel.weight_maps(mlp.experts.w13_weight, mlp.experts.w2_weight) @@ -475,4 +479,5 @@ def routed_given(mlp, x: torch.Tensor, bound: Bound, slot: str): x_fp4, x_sf, routing=(ids, weights), + wide_bal=bound.wide_bal, ) diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index fd17216c..f0f486f6 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -50,10 +50,11 @@ # when requested in the same variable: covering them costs the one-request step about # 0.3% at TP8 OPT_IN_COVERAGE: tuple[str, ...] = ("multi_request",) -# Graph variants of a built slot, requested in the same variable: no width of their own. +# Graph or build variants of a slot, requested in the same variable: no width of their +# own. BAL stays opt-in because its serving evidence covers TP8 only. # skv2_cp: one long request past the measured crossover runs CP-lite in place of SKV2, # in a graph variant of its own (attention/route.py) -OPT_IN_VARIANTS: tuple[str, ...] = ("skv2_cp",) +OPT_IN_VARIANTS: tuple[str, ...] = ("skv2_cp", "mkm_wide_bal") # The MTP draft's EAGLE entries and kernel groups (qwen3_5/draft), requested in the same # variable; the draft spec's build reads them, the target's ignores them DRAFT_ENTRIES: tuple[str, ...] = ("draft_loop", "draft_extend") @@ -83,6 +84,7 @@ "mkm_m128", "mkm_m256", "mkm_m512", + "mkm_wide_bal", "draft_moe", } ) @@ -176,6 +178,18 @@ def slot_reasons() -> list[str]: f"{SLOTS_ENV}: slot '{name}' serves the steps of more than {more} requests " "('multi_request')" ) + if ( + "mkm_wide_bal" in SLOTS + and not { + "mkm_m256", + "mkm_m512", + } + & SLOTS + ): + reasons.append( + f"{SLOTS_ENV}: slot 'mkm_wide_bal' changes the work split of slot " + "'mkm_m256' or 'mkm_m512'" + ) if "route_shared" in SLOTS and "route" not in SLOTS: reasons.append( f"{SLOTS_ENV}: slot 'route_shared' runs in slot 'route''s launch" diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py b/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py index a85df1e5..f62c009c 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py @@ -127,9 +127,10 @@ class Engine(NamedTuple): M512_MAX_TOKENS: 0, } ONEPASS_MAX_ROWS = 42 +WIDE_BAL_MAX_ROWS = 0 # 0 means every row count served by a wide engine -def engine_options(rows: int, plain: bool = False) -> tuple: +def engine_options(rows: int, plain: bool = False, wide_bal: bool = False) -> tuple: """Validated bitwise schedule defines for a wide engine. The 64-token engine keeps its stock schedule past ``ONEPASS_MAX_ROWS``. @@ -143,6 +144,8 @@ def engine_options(rows: int, plain: bool = False) -> tuple: ) if rows == M64_MAX_TOKENS: return ("-DMMA_SFONLY=1",) + (() if plain else ("-DONEPASS=1",)) + if rows in WIDE_ROWS and wide_bal and not plain: + return ("-DBAL_SPLIT=1",) return () @@ -206,14 +209,25 @@ def engine_rows(m: int) -> int: return next(r for r in ENGINE_ROWS if m <= r) -def onepass_plain(rows: int, m: int) -> bool: - """Whether ``m`` rows use the 64-token engine's stock two-pass schedule.""" - return rows == M64_MAX_TOKENS and m > ONEPASS_MAX_ROWS +def plain_schedule(rows: int, m: int, wide_bal: bool = False) -> bool: + """Whether ``m`` rows use an engine's stock schedule. + + The 64-token engine returns to two FC2 passes above its validated bound. An + opt-in balanced wide engine would likewise return to its static split above a + nonzero ``WIDE_BAL_MAX_ROWS``; zero covers every row count. + """ + if rows == M64_MAX_TOKENS: + return m > ONEPASS_MAX_ROWS + return rows in WIDE_ROWS and wide_bal and 0 < WIDE_BAL_MAX_ROWS < m @cache def _kernel( - device_index: int, intermediate: int, rows: int, plain: bool = False + device_index: int, + intermediate: int, + rows: int, + plain: bool = False, + wide_bal: bool = False, ) -> jit.Kernel: """The engine of ``engine_rows`` bound ``rows`` at ``intermediate`` per rank.""" engine = { @@ -236,7 +250,7 @@ def _kernel( f"-DQ35_MKM_SMEM_BYTES={engine.smem_bytes}", ) + engine.defines - + engine_options(rows, plain), + + engine_options(rows, plain, wide_bal), ) return jit.load(cubin, engine.kernel, (THREADS, 1, 1), engine.smem_bytes) @@ -295,15 +309,17 @@ def build( intermediate: int = 128, max_rows: int = MAX_TOKENS, min_rows: int = 1, + wide_bal: bool = False, ) -> None: """Compile the engines that serve ``intermediate`` per rank on ``min_rows`` to ``max_rows`` rows (no compile inside a step).""" lo = 0 for hi in ENGINE_ROWS: if min_rows <= hi and max_rows > lo: - _kernel(device_index, intermediate, hi) - if onepass_plain(hi, max(lo + 1, min(max_rows, hi))): - _kernel(device_index, intermediate, hi, True) + balanced = wide_bal and hi in WIDE_ROWS + _kernel(device_index, intermediate, hi, wide_bal=balanced) + if plain_schedule(hi, max(lo + 1, min(max_rows, hi)), wide_bal=balanced): + _kernel(device_index, intermediate, hi, True, balanced) lo = hi @@ -374,6 +390,7 @@ def mkm_moe( g1_scale_c: torch.Tensor, g2_alpha: torch.Tensor, routing: tuple | None = None, + wide_bal: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """(gemm2_out bf16 [M * 10, 4096], expert_weights bf16 [M, 10]) for logits [M, 512] bf16 (M <= 7, at TP8 M <= 128) and the router front's x_fp4 [M, 2048] / x_sf [M, 256] @@ -381,7 +398,8 @@ def mkm_moe( as uint8 (the engine follows I); g1_alpha, g1_scale_c, g2_alpha [512] fp32. With ``routing`` (route_topk's ids int32 / weights bf16 [M, 10]; 33 to 64 rows, and required at 65 to 512) the engine takes that routing instead of ranking the logits, and returns - those weights.""" + those weights. ``wide_bal`` selects the opt-in balanced unit split of the 256 / + 512-token engines.""" m = logits.shape[0] intermediate = w2_sf.shape[2] * 16 widest = max_tokens( @@ -413,7 +431,14 @@ def mkm_moe( out=gemm2_out.data_ptr(), num_tokens=m, ) - kernel = _kernel(device.index, intermediate, rows, onepass_plain(rows, m)) + balanced = wide_bal and rows in WIDE_ROWS + kernel = _kernel( + device.index, + intermediate, + rows, + plain_schedule(rows, m, wide_bal=balanced), + balanced, + ) stream = torch.cuda.current_stream().cuda_stream if rows in (M32_MAX_TOKENS, M64_MAX_TOKENS, M128_MAX_TOKENS) + WIDE_ROWS: # min(M * 10, SMs) CTAs, one per SM, each running its share of the distinct experts diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_wide.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_wide.cu index de1a7d2a..67751092 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_wide.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/mkm_moe_wide.cu @@ -263,6 +263,9 @@ static_assert(DRAIN_SMEM == 0, "x's shared memory holds the next item's rows whi #ifndef FC2_PF_LEAD // issue the FC2 prefetch when the producer reaches FC1 stage kFc1Stages - FC2_PF_LEAD #define FC2_PF_LEAD 4 #endif +#ifndef BAL_SPLIT // 1: the units past each CTA's own first expert are dealt to balance the CTAs (see the work items) +#define BAL_SPLIT 0 +#endif struct Params { const uint16_t* logits; // [M, 512] bf16 router logits (router_gemm_quant output) @@ -925,6 +928,80 @@ q35_mkm_moe_wide(const __grid_constant__ CUtensorMap tm13, const __grid_constant } named_bar(4, 224); if (warp == 0) { +#if BAL_SPLIT + // Work items, balanced: expert cta's row group 0 first when it has tokens (its stream started before the + // routing was read); every other unit (row group 0 of experts G, G + 1, ... with tokens, then the further row + // groups in expert order) is unit r of one list, dealt first to the CTAs whose expert c has no tokens (in CTA + // order), then round robin from CTA 0: the busiest CTA runs ceil(units / G), not up to 512 / G. + int ni = 0; + if (s.cnt[cta] > 0) { + if (lane == 0) { + s.item_e[0] = cta; + s.item_g[0] = 0; + } + ni = 1; + } + int Z = 0, zrank = 0; // CTAs whose expert c has no tokens, and those before this one + for (int c0 = 0; c0 < G; c0 += 32) { + const uint32_t zm = __ballot_sync(0xffffffffu, c0 + lane < G && s.cnt[c0 + lane] == 0); + Z += __popc(zm); + zrank += __popc(cta >= c0 + 32 ? zm : zm & ((1u << (cta > c0 ? cta - c0 : 0)) - 1u)); + } + int xe[16], xsum = 0, asum = 0; + uint32_t amask = 0; // experts 16 lane + j >= G with tokens (row group 0 in the list) +#pragma unroll + for (int j = 0; j < 16; ++j) { + const int e = 16 * lane + j, n = s.cnt[e]; + if (e >= G && n > 0) { + amask |= 1u << j; + ++asum; + } + xe[j] = n > kRows ? (n - 1) / kRows : 0; + xsum += xe[j]; + } + int ainc = asum, xinc = xsum; +#pragma unroll + for (int d = 1; d < 32; d <<= 1) { + const int ua = __shfl_up_sync(0xffffffffu, ainc, d), ux = __shfl_up_sync(0xffffffffu, xinc, d); + ainc += lane >= d ? ua : 0; + xinc += lane >= d ? ux : 0; + } + const int A2 = __shfl_sync(0xffffffffu, ainc, 31), R = A2 + __shfl_sync(0xffffffffu, xinc, 31); + auto take = [&](int r) { // append unit r of the list (warp-uniform r) + const bool first = r < A2; + const int o = first ? r : r - A2; + const int owner = __ffs(__ballot_sync(0xffffffffu, (first ? ainc : xinc) > o)) - 1; + int e = 0, grp = 0; + if (lane == owner) { + int left = o - (first ? ainc - asum : xinc - xsum); +#pragma unroll + for (int j = 0; j < 16; ++j) { + if (first) { + if ((amask >> j) & 1u) { + if (left == 0) e = 16 * lane + j; + --left; + } + } else { + if (left >= 0 && left < xe[j]) { + e = 16 * lane + j; + grp = left + 1; + } + left -= xe[j]; + } + } + } + e = __shfl_sync(0xffffffffu, e, owner); + grp = __shfl_sync(0xffffffffu, grp, owner); + if (ni >= kMaxItems) __trap(); + if (lane == 0) { + s.item_e[ni] = e; + s.item_g[ni] = grp; + } + ++ni; + }; + if (s.cnt[cta] == 0 && zrank < R) take(zrank); + for (int r = Z + cta; r < R; r += G) take(r); +#else // Work items: own experts cta, cta + G, ... that have tokens (row group 0), then this CTA's further row groups: // the o-th one over all experts (expert order, then group) is CTA G - 1 - o % G's int ni = 0; @@ -974,6 +1051,7 @@ q35_mkm_moe_wide(const __grid_constant__ CUtensorMap tm13, const __grid_constant } ++ni; } +#endif __syncwarp(); if (lane == 0) { s.nitems = ni; diff --git a/test/kernels/test_qwen3_5_moe.py b/test/kernels/test_qwen3_5_moe.py index 106258c3..57bfc97f 100644 --- a/test/kernels/test_qwen3_5_moe.py +++ b/test/kernels/test_qwen3_5_moe.py @@ -201,16 +201,7 @@ def test_mkm_replays_from_a_cuda_graph(layer, m): assert all(torch.equal(a, b) for a, b in zip(eager, captured)) -@pytest.mark.parametrize( - "m", [33, 35, 56, 64, 65, 84, 112, 128, 129, 168, 224, 256, 257, 336, 448, 512] -) -@pytest.mark.parametrize("shared_route", [False, True]) -def test_mkm_on_route_topks_routing_is_the_trtllm_routed_experts( - layer, m, shared_route -): - """Slots mkm_m64 / mkm_m128 / mkm_m256 / mkm_m512: the 64 / 128 / 256 / 512-token engine - on route_topk's routing (ids and weights) gives the runner's rows on the logits, bit for - bit.""" +def _assert_mkm_on_route_topks_routing(layer, m, shared_route, wide_bal=False): if m > mkm.max_tokens( layer["inter"], m32=True, m64=True, m128=True, m256=True, m512=True ): @@ -233,6 +224,7 @@ def test_mkm_on_route_topks_routing_is_the_trtllm_routed_experts( layer["g1c"], layer["g2"], routing=(ids, weights), + wide_bal=wide_bal, ) want_rows, want_weights = _stock(layer, logits, x_fp4, x_sf) assert torch.equal(got_weights.view(torch.int16), want_weights.view(torch.int16)) @@ -240,6 +232,27 @@ def test_mkm_on_route_topks_routing_is_the_trtllm_routed_experts( assert not differ.any(), f"rows {differ.nonzero().flatten().tolist()} differ" +@pytest.mark.parametrize( + "m", [33, 35, 56, 64, 65, 84, 112, 128, 129, 168, 224, 256, 257, 336, 448, 512] +) +@pytest.mark.parametrize("shared_route", [False, True]) +def test_mkm_on_route_topks_routing_is_the_trtllm_routed_experts( + layer, m, shared_route +): + """Slots mkm_m64 / mkm_m128 / mkm_m256 / mkm_m512 give the runner's + materialized-routing rows bit for bit.""" + _assert_mkm_on_route_topks_routing(layer, m, shared_route) + + +@pytest.mark.parametrize("m", [129, 168, 224, 256, 257, 336, 448, 512]) +@pytest.mark.parametrize("shared_route", [False, True]) +def test_wide_engines_balanced_split_is_the_trtllm_routed_experts( + layer, m, shared_route +): + """Opt-in BAL changes only assignment: rows stay bitwise stock.""" + _assert_mkm_on_route_topks_routing(layer, m, shared_route, wide_bal=True) + + @pytest.mark.parametrize("m", [56, 112, 224, 448]) def test_mkm_on_route_topks_routing_replays_from_a_cuda_graph(layer, m): if m > mkm.max_tokens( diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 0940eb14..13a23fd4 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -127,6 +127,7 @@ def test_expert_parallel_keeps_route_and_boundary_slots(monkeypatch): "mkm_m128", "mkm_m256", "mkm_m512", + "mkm_wide_bal", "hop", "gdn_push", "draft_moe", @@ -145,6 +146,7 @@ def test_expert_parallel_keeps_route_and_boundary_slots(monkeypatch): "mkm_m128", "mkm_m256", "mkm_m512", + "mkm_wide_bal", "draft_moe", } ) diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index fbebd160..6d300d14 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -187,6 +187,11 @@ def test_an_sglang_without_the_eagle_forward_view_is_refused(monkeypatch): {"route_topk", "mkm_m512"}, "slot 'mkm_m512' serves the steps of more than 36 requests ('multi_request')", ), + ( + {"mkm_wide_bal"}, + "slot 'mkm_wide_bal' changes the work split of slot 'mkm_m256' or " + "'mkm_m512'", + ), ], ) def test_a_slot_set_that_cannot_run_is_refused_before_the_load( @@ -409,8 +414,40 @@ def test_the_validated_moe_schedules_are_shape_closed(): assert mkm.engine_options(64, plain=True) == ("-DMMA_SFONLY=1",) assert [mkm.reserved_sms(32, rows) for rows in (21, 22)] == [0, 32] assert [mkm.reserved_sms(64, rows) for rows in (35, 42, 43)] == [32, 16, 0] - assert not mkm.onepass_plain(64, 42) - assert mkm.onepass_plain(64, 43) + assert not mkm.plain_schedule(64, 42) + assert mkm.plain_schedule(64, 43) + + +@pytest.mark.parametrize( + "rows", [56, 112, 126, 128, 129, 133, 140, 168, 224, 256, 280, 512] +) +def test_bal_changes_only_the_wide_engine_schedule(rows): + from sgl_nvidia_mega.qwen3_5.sm103.moe import mkm + + engine = mkm.engine_rows(rows) + off = mkm.engine_options(engine) + on = mkm.engine_options(engine, wide_bal=True) + assert on == off + (("-DBAL_SPLIT=1",) if rows > 128 else ()) + assert mkm.plain_schedule(engine, rows, wide_bal=True) == mkm.plain_schedule( + engine, rows + ) + assert mkm.engine_options(engine, plain=True, wide_bal=True) == ( + mkm.engine_options(engine, plain=True) + ) + + +@pytest.mark.parametrize( + "rows,plain", [(129, False), (210, False), (211, True), (280, True)] +) +def test_bal_internal_row_cap_selects_the_static_split(monkeypatch, rows, plain): + from sgl_nvidia_mega.qwen3_5.sm103.moe import mkm + + monkeypatch.setattr(mkm, "WIDE_BAL_MAX_ROWS", 210) + engine = mkm.engine_rows(rows) + assert mkm.plain_schedule(engine, rows, wide_bal=True) is plain + assert ( + "-DBAL_SPLIT=1" in mkm.engine_options(engine, plain, wide_bal=True) + ) is not plain def test_a_numeric_slot_turns_the_claim_into_its_tolerances(): From 5a973ec6ffd64f8066e03ffc7642a5022c6fc258 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 15:32:07 -0700 Subject: [PATCH 09/62] feat: extend multi-request draft exchanges --- sgl_nvidia_mega/qwen3_5/attention/split_kv.py | 12 +- sgl_nvidia_mega/qwen3_5/draft/attention.py | 66 ++++++--- sgl_nvidia_mega/qwen3_5/draft/layer.py | 2 +- sgl_nvidia_mega/qwen3_5/draft/loop.py | 16 ++- sgl_nvidia_mega/qwen3_5/draft/parts.py | 37 +++++ sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 7 + sgl_nvidia_mega/qwen3_5/sm103/draft/embed.py | 132 ++++++++++++++++++ .../qwen3_5/sm103/draft/src/draft_embed.cu | 130 +++++++++++++++++ .../qwen3_5/sm103/draft/src/draft_fc.cu | 38 +++-- .../qwen3_5/sm103/draft/src/draft_top1.cu | 13 +- sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py | 6 +- test/kernels/test_qwen3_5_draft_xchg.py | 95 ++++++++++++- test/unit/test_qwen3_5_runtime.py | 50 +++++-- 13 files changed, 544 insertions(+), 60 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/draft/embed.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_embed.cu diff --git a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py index 5b960996..7aa16cd7 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py +++ b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py @@ -231,11 +231,11 @@ def cp_reason() -> str: return "" -def cp_core(attn, q, kv, meta, forward_batch, q_len: int): - """CP-lite on one request's ``q_len`` FP8 query rows ``q`` of ``attn`` over the pool - ``kv`` at ``meta``'s pages and lengths (this step's rows written): [rows, heads * - head_dim] bf16, or None where CP-lite declines (decided by shapes alone, so every - rank of the KV group agrees).""" +def cp_core(attn, q, kv, meta, forward_batch, q_len: int, min_bs: int | None = 1): + """CP-lite on requests of ``q_len`` FP8 query rows over the pool ``kv``: + [rows, heads * head_dim] bf16, or None where CP-lite declines. ``min_bs`` overrides + CP-lite's configured minimum for a host-selected one-request graph; ``None`` keeps + the configured gate for a multi-request graph.""" from sglang.kernels.ops.attention.utils import canonicalize_stride from sglang.srt.layers.attention.trtllm_mha_backend import ( DEFAULT_WORKSPACE_SIZE_MB, @@ -270,7 +270,7 @@ def cp_core(attn, q, kv, meta, forward_batch, q_len: int): workspace, attn, forward_batch, - min_bs=1, + min_bs=min_bs, ) return None if o is None else o.view(-1, attn.tp_q_head_num * attn.head_dim) diff --git a/sgl_nvidia_mega/qwen3_5/draft/attention.py b/sgl_nvidia_mega/qwen3_5/draft/attention.py index 165c23ac..02d2cef6 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/attention.py +++ b/sgl_nvidia_mega/qwen3_5/draft/attention.py @@ -316,11 +316,7 @@ def decode(self, q8, meta, kv, rows_per_request: int): attn = self.layer.attn bmm1_scale = 1.0 * kv.k_scale[0] * attn.scaling bmm2_scale = kv.v_scale[0] - if ( - self.cp_b1 is not None - and q8.shape[0] == rows_per_request - and self.cp_b1.takes(rows_per_request) - ): + if self.cp_b1 is not None and self.cp_b1.takes(rows_per_request, q8.shape[0]): o = self.cp_b1.run(self, q8, meta, kv, rows_per_request) if o is not None: return o @@ -423,30 +419,37 @@ class DraftCP: """Slots draft_q1_cp / draft_skv2_cp: in the draft loop's ``Q1_CP`` and the draft extend's ``SPLIT_KV_CP`` graph variants (one request past its measured crossover), the decode of the request's rows as CP-lite, which then serves one request - (SGLang's q35_cp_lite through ``attention.split_kv.cp_core``'s call).""" + (SGLang's q35_cp_lite through ``attention.split_kv.cp_core``'s call). Slot + draft_multi also preserves CP-lite on multi-request graphs where SGLang built it; + CP-lite's own shape gate decides whether it runs.""" - def __init__(self, split_kv, route, steps: bool, extend: bool): + def __init__(self, split_kv, route, steps: bool, extend: bool, multi: bool = False): self.split_kv, self.route = split_kv, route - self.steps, self.extend = steps, extend + self.steps, self.extend, self.multi = steps, extend, multi self._logged: set = set() @classmethod def build(cls, parts): steps = "draft_q1_cp" in parts.slots extend = "draft_skv2_cp" in parts.slots - if not (steps or extend): + multi = "draft_multi" in parts.slots + if not (steps or extend or multi): return None from sgl_nvidia_mega.qwen3_5.attention import route, split_kv why = split_kv.cp_reason() if why: + if multi and not (steps or extend): + return None raise RuntimeError( f"NVIDIA megakernel: slots draft_q1_cp / draft_skv2_cp cannot run: {why}" ) - return cls(split_kv, route, steps, extend) + return cls(split_kv, route, steps, extend, multi) - def takes(self, rows_per_request: int) -> bool: - """The graph being captured is this decode's CP variant.""" + def takes(self, rows_per_request: int, rows: int) -> bool: + """Whether this graph lets CP-lite try the decode.""" + if self.multi and rows > rows_per_request: + return True label = _q1_label() if rows_per_request == 1: return self.steps and label == Q1_CP @@ -455,21 +458,38 @@ def takes(self, rows_per_request: int) -> bool: def run(self, attention, q8, meta, kv, rows_per_request: int): layer = attention.layer attn = layer.attn - o = self.split_kv.cp_core(attn, q8, kv, meta, attention, rows_per_request) + multi = q8.shape[0] > rows_per_request + o = self.split_kv.cp_core( + attn, + q8, + kv, + meta, + attention, + rows_per_request, + min_bs=None if multi else 1, + ) if o is not None and torch.cuda.is_current_stream_capturing(): key = (rows_per_request, q8.shape[0]) if key not in self._logged: self._logged.add(key) - logger.info( - "NVIDIA megakernel: Qwen3.5 MTP draft slot %s runs the %s's %d " - "rows as CP-lite (from %d tokens of KV)", - "draft_q1_cp" if rows_per_request == 1 else "draft_skv2_cp", - "step" if rows_per_request == 1 else "extend", - q8.shape[0], - self.route.cp_min_length( - "draft" if rows_per_request == 1 else "extend" - ), - ) + if multi: + logger.info( + "NVIDIA megakernel: Qwen3.5 MTP draft slot draft_multi keeps " + "CP-lite for %d requests x %d rows", + q8.shape[0] // rows_per_request, + rows_per_request, + ) + else: + logger.info( + "NVIDIA megakernel: Qwen3.5 MTP draft slot %s runs the %s's " + "%d rows as CP-lite (from %d tokens of KV)", + "draft_q1_cp" if rows_per_request == 1 else "draft_skv2_cp", + "step" if rows_per_request == 1 else "extend", + q8.shape[0], + self.route.cp_min_length( + "draft" if rows_per_request == 1 else "extend" + ), + ) return o diff --git a/sgl_nvidia_mega/qwen3_5/draft/layer.py b/sgl_nvidia_mega/qwen3_5/draft/layer.py index 33dc81ea..06f32904 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/layer.py +++ b/sgl_nvidia_mega/qwen3_5/draft/layer.py @@ -57,7 +57,7 @@ def fc_input(parts, token_ids, hidden): """The fc output [rows, hidden] of the tokens ``token_ids`` [rows] and the hidden rows ``hidden`` [rows, hidden] they follow.""" lm = parts.lm - embeds = lm.pre_fc_norm_embedding(parts.model.embed_tokens(token_ids)) + embeds = lm.pre_fc_norm_embedding(parts.embed(token_ids)) hidden = lm.pre_fc_norm_hidden(hidden) return lm.fc(torch.cat([embeds, hidden], dim=-1)) diff --git a/sgl_nvidia_mega/qwen3_5/draft/loop.py b/sgl_nvidia_mega/qwen3_5/draft/loop.py index 4defa7bd..9f39ca70 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/loop.py +++ b/sgl_nvidia_mega/qwen3_5/draft/loop.py @@ -100,6 +100,12 @@ def step(self, io) -> None: steps = io.num_steps - 1 tokens = io.draft_tokens_out tokens[:, :1].copy_(io.topk_index) + # request-major slots: every step's rows of several requests made dense in one copy, + # for the front and the FP8 KV write + n, k = io.num_requests, io.num_steps + self._locs = ( + io.out_cache_loc[: n * k].view(n, k).t().contiguous() if n > 1 else None + ) positions, mrope = io.positions, io.mrope_positions if parts.mkd is not None and io.num_requests == SERVED_BS: self._mkd_steps(io, steps) @@ -126,15 +132,13 @@ def _fc(self, token, hidden): mkd = self.parts.mkd if mkd is None or token.shape[0] > FC_SPLIT_MAX_TOKENS: return layer.fc_input(self.parts, token, hidden) - return mkd.split_fc(self.parts.model.embed_tokens(token), hidden) + return mkd.split_fc(self.parts.embed(token), hidden) def _model_step(self, io, i: int, x): """Draft step i's decoder layer on its fc output ``x``: the final-norm rows.""" - # request-major slots: step i's of several requests are a strided view, which the - # front and the FP8 KV write read as dense rows - cache_loc = io.out_cache_loc[i :: io.num_steps] - if io.num_requests > 1: - cache_loc = cache_loc.contiguous() + cache_loc = ( + self._locs[i] if self._locs is not None else io.out_cache_loc[i : i + 1] + ) return layer.layer( self.parts, x, diff --git a/sgl_nvidia_mega/qwen3_5/draft/parts.py b/sgl_nvidia_mega/qwen3_5/draft/parts.py index f9e59ddf..ab7b1515 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/parts.py +++ b/sgl_nvidia_mega/qwen3_5/draft/parts.py @@ -27,6 +27,9 @@ "srt/layers/vocab_parallel_embedding.py": ( "shard_indices", "org_vocab_start_index", + "org_vocab_end_index", + "added_vocab_start_index", + "added_vocab_end_index", ), "srt/models/qwen3_5.py": ( "layers", @@ -67,6 +70,8 @@ def __init__(self, draft, layer): # sm103.draft.top1.Top1Exchange: the multi-request steps' sharded top-1 (slot # draft_multi) self.top1x: Optional[Any] = None + # sm103.draft.embed.EmbedExchange: those steps' embedding rows (slot draft_embed) + self.embedx: Optional[Any] = None # (entry, bs) pairs whose capture logged the engagement line self.engaged: set = set() @@ -91,6 +96,8 @@ def build(cls, draft) -> Parts: parts.mkd = parts._build_mkd() if "draft_multi" in slots: parts.top1x = parts._build_top1x() + if "draft_embed" in slots: + parts.embedx = parts._build_embedx() logger.info( "NVIDIA megakernel: Qwen3.5 MTP draft parts built (slots %s; MoE %s, hop " "%s, exchange %s)", @@ -205,6 +212,36 @@ def _build_top1x(self): device=self.device, ) + def _build_embedx(self): + from sgl_nvidia_mega.qwen3_5.sm103.draft.embed import EmbedExchange + + b = self.binding + embed = self.model.embed_tokens + shard = embed.shard_indices + if shard.added_vocab_end_index != shard.added_vocab_start_index: + raise RuntimeError( + "NVIDIA megakernel: slot draft_embed cannot run: the embedding has " + "added vocabulary" + ) + return EmbedExchange( + world=b.tp_size, + table=embed.weight.data, + vocab_start=int(shard.org_vocab_start_index), + vocab_local=int(shard.org_vocab_end_index - shard.org_vocab_start_index), + group=b.tp_device_group, + cpu_group=b.tp_cpu_group, + device=self.device, + ) + + def embed(self, token_ids: torch.Tensor) -> torch.Tensor: + """The embedding rows of ``token_ids``: slot draft_embed's exchange where it is + built and the rows fit it, else SGLang's vocabulary-parallel embedding.""" + from sgl_nvidia_mega.qwen3_5.sm103.draft.embed import MAX_ROWS + + if self.embedx is not None and token_ids.shape[0] <= MAX_ROWS: + return self.embedx(token_ids) + return self.model.embed_tokens(token_ids) + def note_engaged(self, entry: str, bs: int) -> None: """Log once per (entry, bs) captured: the engagement line every rank prints.""" if torch.cuda.is_current_stream_capturing() and (entry, bs) not in self.engaged: diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index f0f486f6..cbb5b3c9 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -68,6 +68,9 @@ # draft_loop and draft_extend also serve TP8's draft graphs of 2 to 128 requests, # and those of two or more take the sharded top-1 (with multi_request) "draft_multi", + # Their embedding rows are multicast by the rank holding the token, in place of + # the vocabulary-parallel embedding's all-reduce. + "draft_embed", # One long request's steps / extend use CP-lite in graph variants of their own. "draft_q1_cp", "draft_skv2_cp", @@ -225,6 +228,10 @@ def slot_reasons() -> list[str]: f"{SLOTS_ENV}: slot 'draft_multi' serves the draft steps of two or more " "requests ('multi_request')" ) + if "draft_embed" in SLOTS and "draft_multi" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'draft_embed' runs in slot 'draft_multi''s steps" + ) if "draft_skv2" in SLOTS and "draft_extend" not in SLOTS: reasons.append( f"{SLOTS_ENV}: slot 'draft_skv2' runs inside entry 'draft_extend'" diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/embed.py b/sgl_nvidia_mega/qwen3_5/sm103/draft/embed.py new file mode 100644 index 00000000..5bb23a81 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/embed.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""The embedding rows of a multi-row MTP draft step (``src/draft_embed.cu``): the rank whose +vocabulary shard holds a token multicasts its table row, in one launch, in place of the +vocabulary-parallel embedding's masked lookup and all-reduce (the same rows: -0.0 as +0.0, +as that all-reduce adds the other ranks' zero rows).""" + +from __future__ import annotations + +import ctypes +from functools import cache + +import torch + +from sgl_nvidia_mega.common import jit +from sgl_nvidia_mega.qwen3_5.sm103.draft.top1 import _launch_2d +from sgl_nvidia_mega.qwen3_5.sm103.draft.xchg import SRC, WORLDS, symmetric_words + +__all__ = ["EmbedExchange"] + +# draft_embed::kSlots / kMaxRows / kThreads * kWords +SLOTS = 3 +MAX_ROWS = 128 +CHUNK_WORDS = 256 * 4 + + +class DraftEmbedParams(ctypes.Structure): + """sglang::DraftEmbedParams.""" + + _fields_ = [ + (name, ctypes.c_void_p) + for name in ("tokens", "table", "out", "epoch", "mc", "uc") + ] + [ + ("tokens_stride", ctypes.c_int64), + ("m", ctypes.c_int32), + ("hidden", ctypes.c_int32), + ("vocab_start", ctypes.c_int32), + ("vocab_local", ctypes.c_int32), + ("rank", ctypes.c_int32), + ("world", ctypes.c_int32), + ] + + +@cache +def _kernel(device_index: int) -> jit.Kernel: + major, minor = torch.cuda.get_device_capability(device_index) + cubin = jit.compile_source( + (SRC / "draft_embed.cu").read_text(), + "q35_draft_embed", + f"sm_{major}{minor}a", + ("-std=c++17", "-default-device"), + headers={"mkd_common.cuh": (SRC / "mkd_common.cuh").read_text()}, + ) + return jit.load(cubin, "draft_embed_xchg", (256, 1, 1), 0) + + +class EmbedExchange: + """The exchange buffer and epoch of one rank. Construct on every rank of ``group`` (a + collective), before any capture.""" + + def __init__( + self, + *, + world: int, + table: torch.Tensor, + vocab_start: int, + vocab_local: int, + group, + cpu_group, + device: torch.device, + ): + import torch.distributed as dist + + if world not in WORLDS: + raise ValueError(f"the embedding exchange serves TP {WORLDS}, not {world}") + if ( + table.dtype != torch.bfloat16 + or table.dim() != 2 + or not table.is_contiguous() + ): + raise ValueError( + f"embedding exchange: table {table.dtype} {tuple(table.shape)}" + ) + if table.shape[1] % 2 or table.shape[0] < vocab_local: + raise ValueError(f"embedding exchange: table {tuple(table.shape)}") + _kernel(device.index) + self.table, self.hidden = table, table.shape[1] + self.rank, self.world = dist.get_rank(group), world + self.vocab_start, self.vocab_local, self.device = ( + vocab_start, + vocab_local, + device, + ) + self._buf, self._handle, self.mc = symmetric_words( + SLOTS * MAX_ROWS * self.hidden // 2 + world, group, device + ) + self.uc = self._buf.data_ptr() + # [epoch from 1, CTA arrivals] + self.epoch = torch.tensor([1, 0], dtype=torch.int32, device=device) + torch.cuda.synchronize(device) + dist.barrier(group=cpu_group) + + def __call__(self, token_ids: torch.Tensor) -> torch.Tensor: + """The embedding rows [m, hidden] bf16 of ``token_ids`` [m] (int64, any stride).""" + m = int(token_ids.shape[0]) + if ( + not 1 <= m <= MAX_ROWS + or token_ids.dim() != 1 + or token_ids.dtype != torch.int64 + ): + raise ValueError( + f"embedding exchange: tokens {token_ids.dtype} {tuple(token_ids.shape)}" + ) + out = torch.empty(m, self.hidden, dtype=torch.bfloat16, device=self.device) + p = DraftEmbedParams( + tokens=token_ids.data_ptr(), + table=self.table.data_ptr(), + out=out.data_ptr(), + epoch=self.epoch.data_ptr(), + mc=self.mc, + uc=self.uc, + tokens_stride=token_ids.stride(0), + m=m, + hidden=self.hidden, + vocab_start=self.vocab_start, + vocab_local=self.vocab_local, + rank=self.rank, + world=self.world, + ) + chunks = (self.hidden // 2 + CHUNK_WORDS - 1) // CHUNK_WORDS + _launch_2d(_kernel(self.device.index), (chunks, m), p, self.device) + return out diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_embed.cu b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_embed.cu new file mode 100644 index 00000000..fe06868f --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_embed.cu @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// The embedding rows of a multi-row MTP draft step's tokens without the vocabulary-parallel +// embedding's all-reduce (draft_embed_xchg): the rank whose vocabulary shard holds a row's token +// writes that row of its embedding table to every rank through the multicast mapping +// (mkd_xchg's transport: epoch-tagged 64-bit words, each carrying two bf16, three slots rotating +// by epoch, bounded waits that trap); every rank then reads the row from its own copy. +// +// The rows equal SGLang's: there every other rank contributes a zero row to the all-reduce, so +// the sum is the owner's row with -0.0 read as +0.0, which is what is written here. +// +// A rank may own every row of a call and so wait on no one; an arrival word per rank (written at +// entry, awaited by one CTA before the launch completes) keeps every rank within one call of the +// others, so a slot is never rewritten before every rank has read it. +// +// Grid (chunks, m): chunk c of row r moves 8-byte words [c * kThreads * kWords, ...) of the row. +// Built with NVRTC. + +#include "mkd_common.cuh" + +namespace sglang { + +namespace draft_embed { + +inline constexpr int kSlots = 3; +inline constexpr int kMaxRows = 128; +inline constexpr int kThreads = 256; +inline constexpr int kWords = 4; // 8-byte words a thread moves +inline constexpr uint32_t kSpinLimit = 1u << 27; + +// -0.0 as +0.0 in each bf16 half (zero rows added to it by an all-reduce) +SGL_DEVICE uint32_t plus_zero(uint32_t v) { + if ((v & 0xFFFFu) == 0x8000u) v &= 0xFFFF0000u; + if ((v >> 16) == 0x8000u) v &= 0x0000FFFFu; + return v; +} + +SGL_DEVICE uint32_t poll(const uint64_t* src, uint32_t tag) { + uint64_t w; + uint32_t spins = 0; + do { + asm volatile("ld.relaxed.sys.global.u64 %0, [%1];" : "=l"(w) : "l"(src) : "memory"); + if (++spins > kSpinLimit) __trap(); + } while (static_cast(w >> 32) != tag); + return static_cast(w); +} + +} // namespace draft_embed + +struct DraftEmbedParams { + const int64_t* tokens; // the token of row r at tokens[r * tokens_stride], written by an earlier grid + const uint32_t* table; // this rank's embedding shard, [vocab_local, hidden] bf16 as bf16 pairs + uint32_t* out; // [m, hidden] bf16 as bf16 pairs + uint32_t* epoch; // [2]: call counter shared by all ranks (from 1), CTA arrivals + uint64_t* mc; // multicast base of [kSlots, kMaxRows, hidden / 2] words + const uint64_t* uc; // this rank's own copy of it + int64_t tokens_stride; + int32_t m; + int32_t hidden; + int32_t vocab_start; // global id of this rank's table row 0 + int32_t vocab_local; // the ids this rank holds: [vocab_start, vocab_start + vocab_local) + int32_t rank; + int32_t world; +}; + +} // namespace sglang + +using sglang::DraftEmbedParams; + +extern "C" __global__ void __launch_bounds__(256) draft_embed_xchg(const DraftEmbedParams __grid_constant__ p) { + using namespace sglang::draft_embed; + sglang::device::PDLWaitPrimary(); + const int row = static_cast(blockIdx.y); + const int words = p.hidden / 2; + // every CTA reads the call's tag before it arrives; the last arrival advances it + const uint32_t tag = *reinterpret_cast(p.epoch); + const int64_t base = (static_cast(tag % kSlots) * kMaxRows + row) * words; + const int64_t local = p.tokens[row * p.tokens_stride] - p.vocab_start; + const int w0 = static_cast(blockIdx.x) * kThreads * kWords + static_cast(threadIdx.x); + // the arrival words follow the kSlots * kMaxRows rows + const int64_t arrive = static_cast(kSlots) * kMaxRows * words; + const bool barrier = blockIdx.x == 0 && blockIdx.y == 0; + if (barrier && threadIdx.x == 0) { + asm volatile("multimem.st.relaxed.sys.global.b64 [%0], %1;" ::"l"(p.mc + arrive + p.rank), + "l"(static_cast(tag)) : "memory"); + } + if (local >= 0 && local < p.vocab_local) { + const uint32_t* src = p.table + local * words; +#pragma unroll + for (int j = 0; j < kWords; ++j) { + const int w = w0 + j * kThreads; + if (w < words) { + const uint64_t word = static_cast(plus_zero(src[w])) | (static_cast(tag) << 32); + asm volatile("multimem.st.relaxed.sys.global.b64 [%0], %1;" ::"l"(p.mc + base + w), "l"(word) : "memory"); + } + } + } + // the thread's words in flight together, then each re-polled only while its tag is stale + uint64_t v[kWords]; +#pragma unroll + for (int j = 0; j < kWords; ++j) { + const int w = w0 + j * kThreads; + if (w < words) asm volatile("ld.relaxed.sys.global.u64 %0, [%1];" : "=l"(v[j]) : "l"(p.uc + base + w) : "memory"); + } +#pragma unroll + for (int j = 0; j < kWords; ++j) { + const int w = w0 + j * kThreads; + if (w < words) + p.out[static_cast(row) * words + w] = + static_cast(v[j] >> 32) == tag ? static_cast(v[j]) : poll(p.uc + base + w, tag); + } + if (barrier && static_cast(threadIdx.x) < p.world) { + // every rank has entered this call (tags only grow: a rank one call ahead also passes) + uint64_t w; + uint32_t spins = 0; + do { + asm volatile("ld.relaxed.sys.global.u64 %0, [%1];" : "=l"(w) : "l"(p.uc + arrive + threadIdx.x) : "memory"); + if (++spins > kSpinLimit) __trap(); + } while (static_cast(w) - tag >= 0x80000000u); + } + __syncthreads(); + if (threadIdx.x == 0) { + __threadfence(); + if (atomicAdd(p.epoch + 1, 1u) == gridDim.x * gridDim.y - 1) { + p.epoch[1] = 0u; + p.epoch[0] = tag + 1; + } + } +} diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_fc.cu b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_fc.cu index ce142c10..defb8e3c 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_fc.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_fc.cu @@ -166,6 +166,7 @@ namespace fc_gather { inline constexpr int kSlots = 3; inline constexpr uint32_t kSpinLimit = 1u << 27; inline constexpr int kMaxTokens = 16; // the fc split's widest batch (2.2 _FC_SPLIT_MAX_TOKENS) +inline constexpr int kBatch = 8; // polled words in flight per thread SGL_DEVICE void mc_store(uint64_t* addr, uint32_t payload, uint32_t tag) { const uint64_t w = static_cast(payload) | (static_cast(tag) << 32); @@ -206,16 +207,33 @@ SGL_DEVICE void fc_gather_kernel(const FcGatherParams& p) { for (int64_t i = first; i < n; i += stride) { fc_gather::mc_store(dst + i, __float_as_uint(p.partial[i]), tag); } - for (int64_t i = first; i < kWorld * n; i += stride) { - const int64_t r = i / n, j = i % n; - const uint64_t* src = p.uc + (slot * kWorld + r) * kPacket + j; - uint64_t w; - uint32_t spins = 0; - do { - asm volatile("ld.relaxed.sys.global.u64 %0, [%1];" : "=l"(w) : "l"(src) : "memory"); - if (++spins > fc_gather::kSpinLimit) __trap(); - } while (static_cast(w >> 32) != tag); - p.out[(j / kH) * kWorld * kH + r * kH + j % kH] = __uint_as_float(static_cast(w)); + // kBatch independent loads in flight per thread, then each word re-polled only while its tag is stale: + // the words that already landed cost one round trip together, not one each + const int64_t total = kWorld * n; + for (int64_t base = first; base < total; base += stride * fc_gather::kBatch) { + uint64_t w[fc_gather::kBatch]; +#pragma unroll + for (int u = 0; u < fc_gather::kBatch; ++u) { + const int64_t i = base + u * stride; + if (i < total) { + const uint64_t* src = p.uc + (slot * kWorld + i / n) * kPacket + i % n; + asm volatile("ld.relaxed.sys.global.u64 %0, [%1];" : "=l"(w[u]) : "l"(src) : "memory"); + } + } +#pragma unroll + for (int u = 0; u < fc_gather::kBatch; ++u) { + const int64_t i = base + u * stride; + if (i < total) { + const int64_t r = i / n, j = i % n; + const uint64_t* src = p.uc + (slot * kWorld + r) * kPacket + j; + uint32_t spins = 0; + while (static_cast(w[u] >> 32) != tag) { + if (++spins > fc_gather::kSpinLimit) __trap(); + asm volatile("ld.relaxed.sys.global.u64 %0, [%1];" : "=l"(w[u]) : "l"(src) : "memory"); + } + p.out[(j / kH) * kWorld * kH + r * kH + j % kH] = __uint_as_float(static_cast(w[u])); + } + } } __syncthreads(); if (threadIdx.x == 0) { diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_top1.cu b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_top1.cu index 4ef844fb..7bfc074a 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_top1.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_top1.cu @@ -157,12 +157,21 @@ SGL_DEVICE void draft_top1_xchg_kernel(const DraftTop1Params& p) { mc_store(dst + 1, static_cast(k >> 32), tag); } if (t < p.m) { + // every rank's two words of this row in flight together, then each re-polled only while its tag + // is stale: the keys that already landed cost one round trip, not one each + uint64_t v[2 * kWorld]; +#pragma unroll + for (int k = 0; k < 2 * kWorld; ++k) { + const uint64_t* w = p.uc + ((slot * kWorld + k / 2) * kMaxRows + t) * 2 + k % 2; + asm volatile("ld.relaxed.sys.global.u64 %0, [%1];" : "=l"(v[k]) : "l"(w) : "memory"); + } uint64_t win = 0; #pragma unroll for (int r = 0; r < kWorld; ++r) { const uint64_t* w = p.uc + ((slot * kWorld + r) * kMaxRows + t) * 2; - const uint32_t lo = poll(w, tag); - const uint32_t hi = poll(w + 1, tag); + const uint32_t lo = static_cast(v[2 * r] >> 32) == tag ? static_cast(v[2 * r]) : poll(w, tag); + const uint32_t hi = + static_cast(v[2 * r + 1] >> 32) == tag ? static_cast(v[2 * r + 1]) : poll(w + 1, tag); win = max64(win, (static_cast(hi) << 32) | lo); } const int64_t id = static_cast(0xFFFFFFFFu - static_cast(win)); diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py b/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py index 144c5108..a9eebfb4 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py @@ -52,7 +52,7 @@ _MAX_SMEM_BYTES = 216 * 1024 # fc_gather: the split fc's widest batch (2.2 _FC_SPLIT_MAX_TOKENS) and its grid GATHER_MAX_TOKENS = 16 -GATHER_CTAS, GATHER_THREADS = 32, 512 +GATHER_THREADS = 512 def _order_key(value: float) -> int: @@ -714,6 +714,7 @@ def __init__( # [epoch, CTA arrivals]: the epoch tags the call's words, from 1 self.epoch_done = torch.tensor([1, 0], dtype=torch.int32, device=device) sms = _sm_count(device.index) + self._ctas = sms self._kernels = _draft_fc_kernels( device.index, world, draft_fc_partial_rows(HIDDEN, 2 * HIDDEN // world, sms) ) @@ -738,5 +739,6 @@ def __call__(self, partial: torch.Tensor) -> torch.Tensor: rank=self.rank, m=m, ) - _launch(self._kernels["fc_gather"], GATHER_CTAS, p, self.device, pdl=False) + # one CTA an SM: with the batched polls every thread waits about one round trip + _launch(self._kernels["fc_gather"], self._ctas, p, self.device, pdl=False) return out diff --git a/test/kernels/test_qwen3_5_draft_xchg.py b/test/kernels/test_qwen3_5_draft_xchg.py index 63a799e4..6014026d 100644 --- a/test/kernels/test_qwen3_5_draft_xchg.py +++ b/test/kernels/test_qwen3_5_draft_xchg.py @@ -14,7 +14,10 @@ graph with fresh inputs, equal to its eager run; - the sharded top-1 of slot draft_multi (``sm103/draft/top1``) picks, for 1-16 rows, the argmax of the all-gathered logits under both NaN rules (ties to the lowest id, -0.0 as - +0.0), writes the token column, 1.0 and positions + 1, and replays from a CUDA graph. + +0.0), writes the token column, 1.0 and positions + 1, and replays from a CUDA graph; +- the embedding exchange of slot draft_embed (``sm103/draft/embed``) gives, for 1-128 rows, + the rows of SGLang's vocabulary-parallel embedding (masked lookup + all-reduce) bit for + bit, -0.0 entries included, and replays from a CUDA graph. The bitwise comparison with the kernels of SGLang branch nvidia/qwen35-agentx-gb300 needs that branch installed as ``sglang`` and is not part of this test. @@ -429,5 +432,95 @@ def run(): dist.barrier(group=cpu_group) +V_EMBED = 512 # embedding rows a rank holds here + + +@pytest.fixture(scope="module") +def embedx(ranks): + from sgl_nvidia_mega.qwen3_5.sm103.draft.embed import EmbedExchange + + rank, world, group, cpu_group = ranks + full = _randn(world * V_EMBED, H, g=_gen(9100)) # the same table on every rank + full[3, :64] = -0.0 + full[V_EMBED, 1::2] = -0.0 # rank 1's first row + table = full[rank * V_EMBED : (rank + 1) * V_EMBED].contiguous() + ex = EmbedExchange( + world=world, + table=table, + vocab_start=rank * V_EMBED, + vocab_local=V_EMBED, + group=group, + cpu_group=cpu_group, + device=torch.device("cuda", torch.cuda.current_device()), + ) + return ex, table + + +def _embed_want(table, ids, rank, group): + """SGLang's VocabParallelEmbedding at TP > 1: out-of-shard ids read row 0 and are + zeroed, then the all-reduce.""" + start = rank * V_EMBED + mask = (ids < start) | (ids >= start + V_EMBED) + rows = torch.nn.functional.embedding((ids - start) * ~mask, table) + rows.masked_fill_(mask.unsqueeze(-1), 0) + dist.all_reduce(rows, group=group) + return rows + + +def _embed_ids(world, m, seed): + ids = torch.randint(0, world * V_EMBED, (m,), device="cuda", generator=_gen(seed)) + edges = [3, V_EMBED, V_EMBED - 1, world * V_EMBED - 1, 0] + ids[: min(m, len(edges))] = torch.tensor(edges[:m], device="cuda") + return ids + + +@pytest.mark.parametrize("m", [1, 2, 7, 16, 28, 112, 128]) +def test_embedding_exchange_is_the_vocab_parallel_embedding(ranks, embedx, m): + rank, world, group, cpu_group = ranks + ex, table = embedx + ids = _embed_ids(world, m, 9200 + m) + tokens = torch.zeros(m, NUM_STEPS, dtype=torch.int64, device="cuda") + tokens[:, 2] = ids # a strided column, as the draft tokens' next column + got = ex(tokens[:, 2]) + want = _embed_want(table, ids, rank, group) + torch.cuda.synchronize() + assert torch.equal(got.view(torch.int16), want.view(torch.int16)) + dist.barrier(group=cpu_group) + + +def test_embedding_exchange_replays_from_a_cuda_graph(ranks, embedx): + rank, world, group, cpu_group = ranks + ex, table = embedx + m = 16 + ids = torch.zeros(m, dtype=torch.int64, device="cuda") + outs = [] + + def run(): + outs.clear() + for _ in range(NUM_STEPS - 1): + outs.append(ex(ids)) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + dist.barrier(group=cpu_group) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + torch.cuda.synchronize() + dist.barrier(group=cpu_group) + for seed in (1, 2, 3): + ids.copy_(_embed_ids(world, m, 9300 + seed)) + graph.replay() + torch.cuda.synchronize() + want = _embed_want(table, ids, rank, group) + assert all( + torch.equal(o.view(torch.int16), want.view(torch.int16)) for o in outs + ) + dist.barrier(group=cpu_group) + + if __name__ == "__main__": multigpu_pytest_main(__name__, __file__, num_gpus=(4,)) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 13a23fd4..4ffd3573 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -333,9 +333,17 @@ def one(n): ] -@pytest.mark.parametrize("rows, label", [(1, "q1_cp"), (7, "split_kv_cp")]) -def test_the_draft_decode_takes_cp_lite_for_one_request_in_its_cp_variant( - monkeypatch, rows, label +@pytest.mark.parametrize( + "rows,label,requests,min_bs", + [ + (1, "q1_cp", 1, 1), + (7, "split_kv_cp", 1, 1), + (1, "q1_stock", 4, None), + (7, "split_kv_stock", 4, None), + ], +) +def test_the_draft_decode_keeps_cp_lite_for_cp_variants_and_multi_request( + monkeypatch, rows, label, requests, min_bs ): from sgl_nvidia_mega.qwen3_5.attention import route, split_kv from sgl_nvidia_mega.qwen3_5.draft import attention as draft_attention @@ -370,19 +378,21 @@ def attend( attn.heads, attn.head_dim, attn.page_size = heads, dim, page attn.q1 = attn.skv2 = None attn._noted = set() - attn.cp_b1 = draft_attention.DraftCP(split_kv, route, steps=True, extend=True) + attn.cp_b1 = draft_attention.DraftCP( + split_kv, route, steps=True, extend=True, multi=requests > 1 + ) cache = torch.zeros(4 * page, 1, dim, dtype=torch.float8_e4m3fn) kv = SimpleNamespace( layer_ids=(5,), k=(cache,), v=(cache,), k_scale=(2.0,), v_scale=(3.0,) ) meta = SimpleNamespace( - page_table=torch.zeros(1, 4, dtype=torch.int32), - cache_seqlens=torch.full((1,), 100, dtype=torch.int32), + page_table=torch.zeros(requests, 4, dtype=torch.int32), + cache_seqlens=torch.full((requests,), 100, dtype=torch.int32), ) - q8 = torch.zeros(rows, heads * dim, dtype=torch.float8_e4m3fn) + q8 = torch.zeros(requests * rows, heads * dim, dtype=torch.float8_e4m3fn) out = attn.decode(q8, meta, kv, rows_per_request=rows) - assert out.shape == (rows, heads * dim) and bool((out == 1).all()) - assert calls == [((rows, heads, dim), rows, 1.0, 3.0, 1)] + assert out.shape == (requests * rows, heads * dim) and bool((out == 1).all()) + assert calls == [((requests * rows, heads, dim), rows, 1.0, 3.0, min_bs)] @pytest.mark.parametrize( @@ -399,6 +409,10 @@ def attend( "slot 'draft_multi' serves the draft steps of two or more requests " "('multi_request')", ), + ( + {"draft_loop", "multi_request", "draft_embed"}, + "slot 'draft_embed' runs in slot 'draft_multi''s steps", + ), ], ) def test_the_draft_slots_name_what_they_need(monkeypatch, slots, want): @@ -416,6 +430,9 @@ def __init__(self, slots=frozenset(), tp_size=8): def note_engaged(self, entry, bs): self.engaged.append((entry, bs)) + def embed(self, token_ids): + return self.model.embed_tokens(token_ids) + def _stub_model_step(monkeypatch, calls): """The draft's model step as a log: fc_input -> x, layer -> hidden, top1 -> token.""" @@ -571,6 +588,21 @@ def top1_into( assert io.positions.tolist() == [10, 20, 30, 40] +def test_the_draft_embeds_through_the_exchange_where_it_fits(): + from sgl_nvidia_mega.qwen3_5.draft.parts import Parts + from sgl_nvidia_mega.qwen3_5.sm103.draft.embed import MAX_ROWS + + calls = [] + parts = Parts.__new__(Parts) + parts.model = SimpleNamespace(embed_tokens=lambda t: calls.append("sglang") or t) + parts.embedx = None + parts.embed(torch.zeros(4, dtype=torch.int64)) + parts.embedx = lambda t: calls.append("exchange") or t + parts.embed(torch.zeros(MAX_ROWS, dtype=torch.int64)) + parts.embed(torch.zeros(MAX_ROWS + 1, dtype=torch.int64)) + assert calls == ["sglang", "exchange", "sglang"] + + @pytest.mark.parametrize("requests, split", [(16, True), (32, False)]) def test_the_split_fc_takes_the_draft_steps_up_to_its_widest_batch( monkeypatch, requests, split From 2dc92c2cf5792af81494929fd8ad906edd13c2cf Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 15:42:42 -0700 Subject: [PATCH 10/62] feat: fuse validated draft boundaries --- sgl_nvidia_mega/qwen3_5/draft/extend.py | 6 + sgl_nvidia_mega/qwen3_5/draft/layer.py | 92 ++++++++- sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py | 140 +++++++++++++ sgl_nvidia_mega/qwen3_5/draft/parts.py | 18 +- sgl_nvidia_mega/qwen3_5/draft/xchg.py | 88 ++++++++- sgl_nvidia_mega/qwen3_5/moe/__init__.py | 8 +- sgl_nvidia_mega/qwen3_5/moe/mkm.py | 2 + sgl_nvidia_mega/qwen3_5/moe/route.py | 24 ++- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 6 + .../qwen3_5/sm103/draft/src/draft_fc_sum.cu | 163 ++++++++++++++++ .../qwen3_5/sm103/draft/src/mkd_fc.cu | 60 +++++- .../qwen3_5/sm103/draft/src/mkd_xchg.cu | 14 +- sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py | 94 ++++++++- sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py | 13 +- test/kernels/test_qwen3_5_draft_xchg.py | 184 +++++++++++++++++- test/kernels/test_qwen3_5_moe.py | 18 +- test/unit/test_qwen3_5_runtime.py | 115 +++++++++++ 17 files changed, 1012 insertions(+), 33 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_fc_sum.cu diff --git a/sgl_nvidia_mega/qwen3_5/draft/extend.py b/sgl_nvidia_mega/qwen3_5/draft/extend.py index d13669ff..33265b7b 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/extend.py +++ b/sgl_nvidia_mega/qwen3_5/draft/extend.py @@ -91,6 +91,12 @@ def step(self, io) -> None: # copies at the end of the graph selected = torch.index_select(hidden, 0, io.select_index, out=io.hidden_out) if mkd is not None: + local = mkd.extend_head(selected) + if local is not None: + mkd.exchange_extend( + local, selected, io.topk_index_out, io.topk_p_out, mkd_head=True + ) + return lm = parts.lm local = lm.logits_processor._compute_lm_head(selected, lm.lm_head, None) mkd.exchange_extend(local, selected, io.topk_index_out, io.topk_p_out) diff --git a/sgl_nvidia_mega/qwen3_5/draft/layer.py b/sgl_nvidia_mega/qwen3_5/draft/layer.py index 06f32904..ff305147 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/layer.py +++ b/sgl_nvidia_mega/qwen3_5/draft/layer.py @@ -15,15 +15,19 @@ and their arg-max (``torch.argmax``'s choice, as SGLang's top-1). Stock: the ops SGLang's draft runs for these rows, in its order. Slot ``draft_moe`` -replaces the MoE, slot ``draft_hop`` the two all-reduce points and ``o_proj``. +replaces the MoE, slot ``draft_hop`` the two all-reduce points and ``o_proj``; slot +``draft_fin`` (numeric) hands a deferred MoE finalize to the hop's final norm launch. """ from __future__ import annotations +import logging from functools import partial import torch +logger = logging.getLogger(__name__) + # The SGLang names this module reads (test/unit/test_contract.py pins them) SGLANG_READS = { "srt/models/qwen3_5_mtp.py": ( @@ -43,6 +47,17 @@ "tensor_model_parallel_all_reduce", "tensor_model_parallel_all_gather", ), + "srt/models/qwen2_moe.py": ("supports_deferred_finalize",), + "srt/layers/layer_boundary/fusions/moe_finalize.py": ( + "MoeDeferredFinalize", + "routed_output", + "expert_weights", + "permuted_indices", + "gated_shared_output", + "m", + "complete", + ), + "srt/layers/layernorm.py": ('"gemma_weight"',), } # The SGLang functions whose body this module mirrors, by the digest of their code @@ -108,9 +123,18 @@ def layer(parts, x, positions, cache_loc, meta, kv, *, rows_per_request: int): else None ), ) + # Compatible NVFP4 draft experts reuse the target MoE engines from two rows. + mkm = ( + parts.moe_mkm + if parts.moe_mkm is not None and parts.moe_mkm.takes(rows, rows_per_request) + else None + ) moe = ( parts.moe - if rows_per_request == 1 and parts.moe is not None and parts.moe.ready(rows) + if mkm is None + and rows_per_request == 1 + and parts.moe is not None + and parts.moe.ready(rows) else None ) if hop is not None: @@ -130,7 +154,42 @@ def layer(parts, x, positions, cache_loc, meta, kv, *, rows_per_request: int): moe_in, residual = lay.post_attention_layernorm( tensor_model_parallel_all_reduce(attn_out), residual ) - moe_out = moe(moe_in) if moe is not None else lay.mlp(moe_in) + fin_hop = hop if hop is not None else extend_hop + fin = fin_hop is not None and moe is None and "draft_fin" in parts.slots + if fin and not lay.mlp.supports_deferred_finalize: + _note_fin( + rows, "SGLang's: the experts do not defer their finalize", hop is None + ) + fin = False + if fin: + owed = ( + mkm(lay, moe_in) + if mkm is not None + else lay.mlp(moe_in, defer_finalize=True) + ) + if _hop_takes(fin_hop.hop, owed): + _note_fin(rows, "in the final norm's hop (slot draft_fin)", hop is None) + out, _ = fin_hop.hop.moe_finalize_all_reduce_rms_norm( + routed_output=owed.routed_output, + expert_weights=owed.expert_weights, + permuted_indices=owed.permuted_indices, + gated_shared_output=owed.gated_shared_output, + residual=residual, + gamma=parts.model.norm.gemma_weight, + ) + return out + _note_fin( + rows, "SGLang's: the hop does not read this deferred output", hop is None + ) + moe_out = owed.complete() + if hop is None and mkm is None: + # SGLang's deferred finish already applies its configured reduction. + out, _ = parts.model.norm(moe_out, residual) + return out + elif mkm is not None: + moe_out = mkm(lay, moe_in).complete() + else: + moe_out = moe(moe_in) if moe is not None else lay.mlp(moe_in) if hop is not None: out, _ = hop.final_norm( moe_out, @@ -145,6 +204,33 @@ def layer(parts, x, positions, cache_loc, meta, kv, *, rows_per_request: int): return out +_FIN_NOTED: set[tuple[bool, int]] = set() + + +def _note_fin(rows: int, where: str, extend: bool = False) -> None: + """Log once per captured row count where slot ``draft_fin`` finalized the MoE.""" + if torch.cuda.is_current_stream_capturing() and (extend, rows) not in _FIN_NOTED: + _FIN_NOTED.add((extend, rows)) + logger.info( + "NVIDIA megakernel: Qwen3.5 MTP draft %s of %d rows: MoE finalize %s", + "extend" if extend else "step", + rows, + where, + ) + + +def _hop_takes(hop, owed) -> bool: + """Whether the hop's finalize reads this deferred MoE output as SGLang wrote it.""" + shared = owed.gated_shared_output + return ( + hop.supports(owed.m) + and owed.routed_output.dtype == torch.bfloat16 + and owed.expert_weights.dtype == torch.bfloat16 + and owed.permuted_indices.dtype == torch.int32 + and (shared is None or shared.dtype == torch.bfloat16) + ) + + def top1(parts, hidden): """SGLang's draft top-1 [rows] int64 of the final-norm rows ``hidden``.""" from sglang.srt.distributed import tensor_model_parallel_all_gather diff --git a/sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py b/sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py new file mode 100644 index 00000000..e03ff8a1 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""The MTP draft's compatible NVFP4 experts on the target MoE engines. + +Slot ``draft_moe`` keeps its one-row BF16 launch. From two through 64 rows, a draft +whose experts already have the target NVFP4 layout reuses the router front and the +7/16/32/64-token engines. Incompatible weights simply keep the normal draft path. +""" + +from __future__ import annotations + +import logging +from types import SimpleNamespace +from typing import Optional + +import torch + +logger = logging.getLogger(__name__) + +# The SGLang names this module reads (test/unit/test_contract.py pins them) +SGLANG_READS = { + "srt/models/qwen2_moe.py": ("experts", "layer_id"), + "srt/layers/quantization/modelopt_quant.py": ('"w13_weight"', '"w2_weight"'), +} + +MIN_ROWS = 2 +MAX_ROWS = 64 +# Past 12 rows the 32-token engine leaves room for the BF16 shared expert; its rows +# equal the 16-token engine bit for bit. +M16_MAX_ROWS = 12 + +# the MoE slots moe.step reads, all off but the ones a width takes +_SLOT_NAMES = ( + "route", + "mkm", + "route_topk", + "route_gemm", + "route_shared", + "mkm_m32", + "mkm_m64", + "mkm_m128", + "mkm_m256", + "mkm_m512", + "lowm_se", + "hop_fp8", +) + + +class DraftMkm: + """The draft layer's bound for moe.step and the per-width slots it takes.""" + + def __init__(self, bound, wide: bool): + self.bound = bound + self.wide = wide + self._noted: set = set() + + @classmethod + def build(cls, mlp) -> tuple[Optional[DraftMkm], str]: + from sgl_nvidia_mega.qwen3_5.moe import mkm, route + from sgl_nvidia_mega.qwen3_5.sm103.moe import mkm as mkm_kernel + from sgl_nvidia_mega.qwen3_5.sm103.moe import router + + w13 = getattr(getattr(mlp, "experts", None), "w13_weight", None) + if not isinstance(w13, torch.Tensor) or w13.dtype != torch.uint8: + return None, "routed experts are not in the target NVFP4 layout" + reason = route.unsupported(mlp) or mkm.unsupported(mlp) + if reason: + return None, reason + device = torch.cuda.current_device() + intermediate = mlp.experts.w2_weight.shape[-1] * 2 + reason = mkm_kernel.m32_unsupported(device, intermediate) + if reason: + return None, f"32 / 64-token engines: {reason}" + # compiled kernels are cached per shape: what the target built is reused, the rest added + router.build(device, router.MAX_TOKENS) + router.build_topk(device) + mkm_kernel.build(device, intermediate, MAX_ROWS) + wide = intermediate in mkm_kernel.WIDE_ENGINES + maps = mkm_kernel.weight_maps(mlp.experts.w13_weight, mlp.experts.w2_weight) + bound = route.Bound( + route=True, + mkm={mlp.layer_id: maps}, + mkm_wide=wide, + # the hop's finalize takes registered identity rows without a map load + identity=torch.arange( + mkm_kernel.M128_MAX_TOKENS * mkm_kernel.TOP_K, + dtype=torch.int32, + device="cuda", + ), + topk=True, + m32={mlp.layer_id: maps}, + m64=True, + m16_max_rows=M16_MAX_ROWS, + ) + return cls(bound, wide), "" + + def takes(self, rows: int, rows_per_request: int) -> bool: + """Whether a draft step or extend of ``rows`` takes the target engines.""" + if rows_per_request == 1 and rows < MIN_ROWS: + return False + return 1 <= rows <= MAX_ROWS + + def state(self, rows: int): + """moe.step's step state for ``rows`` rows: the MoE slots of that width, this bound.""" + from sgl_nvidia_mega.qwen3_5 import plan + + route = rows <= plan.ROUTE_MAX_ROWS + on = dict.fromkeys(_SLOT_NAMES, False) + on.update( + route=route, + mkm=route + and rows <= (plan.MKM_WIDE_MAX_ROWS if self.wide else plan.MKM_MAX_ROWS), + route_topk=not route, + mkm_m32=not route and rows <= plan.MKM_M32_MAX_ROWS, + mkm_m64=plan.MKM_M32_MAX_ROWS < rows <= plan.MKM_M64_MAX_ROWS, + ) + if rows not in self._noted and torch.cuda.is_current_stream_capturing(): + self._noted.add(rows) + logger.info( + "NVIDIA megakernel: Qwen3.5 MTP draft MoE of %d rows on the " + "target path: %s", + rows, + [k for k, v in on.items() if v], + ) + return SimpleNamespace( + slots=SimpleNamespace(**on), bound={"moe": self.bound}, scratch={} + ) + + def __call__(self, layer, x: torch.Tensor): + """Return SGLang's deferred-finalize handoff for ``x``.""" + from sgl_nvidia_mega.qwen3_5 import moe + + return moe.step( + layer, x, self.state(int(x.shape[0])), defer=True, reduce=_unreduced + ) + + +def _unreduced(rows: torch.Tensor) -> torch.Tensor: + """The draft layer owns the tensor-parallel reduction after the MoE.""" + return rows diff --git a/sgl_nvidia_mega/qwen3_5/draft/parts.py b/sgl_nvidia_mega/qwen3_5/draft/parts.py index ab7b1515..6ed2f353 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/parts.py +++ b/sgl_nvidia_mega/qwen3_5/draft/parts.py @@ -74,6 +74,8 @@ def __init__(self, draft, layer): self.embedx: Optional[Any] = None # (entry, bs) pairs whose capture logged the engagement line self.engaged: set = set() + # draft.moe_mkm.DraftMkm: compatible quantized draft experts on the target path + self.moe_mkm: Optional[Any] = None @classmethod def build(cls, draft) -> Parts: @@ -98,11 +100,25 @@ def build(cls, draft) -> Parts: parts.top1x = parts._build_top1x() if "draft_embed" in slots: parts.embedx = parts._build_embedx() + if "draft_moe" in slots: + from sgl_nvidia_mega.qwen3_5.draft import moe_mkm + + parts.moe_mkm, why = moe_mkm.DraftMkm.build(parts.layer.mlp) + if parts.moe_mkm is None: + logger.info( + "NVIDIA megakernel: draft MoE target path unavailable (%s)", why + ) logger.info( "NVIDIA megakernel: Qwen3.5 MTP draft parts built (slots %s; MoE %s, hop " "%s, exchange %s)", sorted(slots), - "single launch" if parts.moe is not None else "SGLang's", + ( + "single launch at one row, target path from two rows" + if parts.moe_mkm is not None + else "single launch" + if parts.moe is not None + else "SGLang's" + ), "draft hop" if parts.hop is not None else "SGLang's", "MKD" if parts.mkd is not None else "SGLang's", ) diff --git a/sgl_nvidia_mega/qwen3_5/draft/xchg.py b/sgl_nvidia_mega/qwen3_5/draft/xchg.py index 52430dc1..69460f62 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/xchg.py +++ b/sgl_nvidia_mega/qwen3_5/draft/xchg.py @@ -33,7 +33,6 @@ from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels logger = logging.getLogger(__name__) - # Memo rows built per launch at construction _MEMO_BLOCK = 4096 # CTAs of mkd_draft_head per SM (T1, TP8 shard: 1 -> 39.3 us, 2 -> 40.5 us, 3 -> 41.6 us) @@ -44,7 +43,7 @@ HEAD_SMEM_PAD = 120 * 1024 # SMs a hidden rank leaves to the side-stream hidden partial: the early-released head # would fill every SM, and the partial (which the exchange waits for) would start only as -# head CTAs retire; the head then runs on the other SMs +# head CTAs retire; the four-row partial ends before the head at this measured split PARTIAL_SMS = 24 # Draft and draft-extend widths that take the row-split fc (2.2 _FC_SPLIT_MAX_TOKENS) FC_SPLIT_MAX_TOKENS = kernels.GATHER_MAX_TOKENS @@ -171,6 +170,13 @@ def __init__( self._fc_gather = kernels.FcGather( world=world, rank=rank, group=device_group, cpu_group=cpu_group, device=dev ) + self._fc_allsum = kernels.FcAllsum( + world=world, + rank=rank, + group=device_group, + cpu_group=cpu_group, + device=dev, + ) self.is_hidden_rank = use_hidden self.slice_start, self.width = start, width self.n_embed = hidden // width @@ -234,6 +240,7 @@ def __init__( self.step = 0 self.last_step = False self.partial_pending = False + self._extend_head_noted = False # the vocabulary exchange's per-step outputs (DraftVocabExchange) self.num_rows = 1 self.topk_index = [ @@ -245,8 +252,8 @@ def __init__( for _ in range(num_steps) ] logger.info( - "MKD draft boundary active (tp=%d, %s candidates, %.1f MB exchange buffer); " - "sources %s.", + "MKD draft boundary active (tp=%d, %s candidates, %.1f MB exchange " + "buffer; one-launch split-fc all-sum); sources %s.", world, "memo" if memo else "embedding", words * 8 / 2**20, @@ -564,16 +571,75 @@ def select_extend( raise RuntimeError("MKD draft boundary: the extend exchange serves one row") return self.exchange_extend(local_logits, hidden_states) + def extend_head(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: + """Run a one-request extend head as an MKD step. + + On a hidden rank the next draft's partial runs on the side stream; the exchange + waits on its flag and reads the head's per-CTA maxima. Return ``None`` when the + boundary cannot use this shape. + """ + if self.head_part is None or not self.extend_fc or hidden_states.shape[0] != 1: + return None + x = hidden_states.to(torch.bfloat16).contiguous() + partial_ctas = ( + max(0, min(self.partial_sms, self.head_part.numel() - 1)) + if self.is_hidden_rank + else 0 + ) + if self.is_hidden_rank: + norm = self.norm_hidden + self.side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.side_stream): + kernels.mkd_hidden_partial( + x, + self._fc_shard, + norm.weight.data, + norm.variance_epsilon, + self.slice_start, + self.own_partial, + self.epoch_done, + self.partial_sync, + self.world, + ctas=partial_ctas, + ) + self.partial_pending = True + logits = torch.empty( + 1, self.lm_head_weight.shape[0], dtype=torch.bfloat16, device=x.device + ) + self.head_part_cur = self.head_part[: self.head_part.numel() - partial_ctas] + kernels.mkd_draft_head( + x, + self.lm_head_weight, + logits, + self.head_part_cur, + reverse=False, + smem_pad=self.head_pad, + ) + self.head_used = True + if torch.cuda.is_current_stream_capturing() and not self._extend_head_noted: + self._extend_head_noted = True + logger.info( + "NVIDIA megakernel: MKD draft extend uses the BF16 step head on " + "%d CTAs, %s", + self.head_part_cur.numel(), + f"hidden partial on {partial_ctas} SMs of the side stream" + if self.is_hidden_rank + else "no partial (embedding rank)", + ) + return logits + def exchange_extend( self, local_logits: torch.Tensor, hidden_states: Optional[torch.Tensor], topk_index: Optional[torch.Tensor] = None, topk_p: Optional[torch.Tensor] = None, + mkd_head: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Draft-extend token (torch.argmax NaN rule) through the same one-launch exchange; with extend_fc it also writes the next draft's first fc into fc_buf. The token and - its 1.0 land in ``topk_index`` / ``topk_p`` ([1, 1]) when given.""" + its 1.0 land in ``topk_index`` / ``topk_p`` ([1, 1]) when given. ``mkd_head``: the + logits come from ``extend_head`` (its per-CTA maxima and side-stream partial).""" if topk_index is None: topk_index = torch.empty( (1, 1), dtype=torch.int64, device=local_logits.device @@ -585,7 +651,7 @@ def exchange_extend( payload = ( self.extend_fc and hidden_states is not None and hidden_states.shape[0] == 1 ) - if payload and self.is_hidden_rank: + if payload and self.is_hidden_rank and not mkd_head: norm = self.norm_hidden kernels.draft_fc_partial( hidden_states.contiguous(), @@ -617,8 +683,15 @@ def exchange_extend( kernels.NAN_KEY_TORCH_ARGMAX, fc_out=self.fc_buf if payload else None, token_only=not payload, + partial_flag=self.partial_sync[1:] + if mkd_head and payload and self.partial_pending + else None, scan_part=self.scan_part, + head_part=self.head_part_cur if mkd_head and self.head_used else None, ) + if mkd_head: + self.head_used = False + self._join_side_streams() self.fc_buf_key = self.input_seq if payload else None return topk_p, topk_index @@ -681,6 +754,5 @@ def split_fc( return self._gather_partials(partial) def _gather_partials(self, partial: torch.Tensor) -> torch.Tensor: - gathered = self._fc_gather(partial) out = partial.new_empty(partial.shape, dtype=torch.bfloat16) - return kernels.ordered_partial_sum(gathered, self.world, out) + return self._fc_allsum(partial, out) diff --git a/sgl_nvidia_mega/qwen3_5/moe/__init__.py b/sgl_nvidia_mega/qwen3_5/moe/__init__.py index e3c46dd1..b87ac113 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/moe/__init__.py @@ -136,9 +136,11 @@ def build(model, ctx) -> route.Bound: return bound -def step(layer, x: torch.Tensor, st, defer: bool): +def step(layer, x: torch.Tensor, st, defer: bool, reduce=None): """``layer``'s MoE on the normed input ``x`` [rows, hidden]: the deferred-finalize - handoff with ``defer``, else the unreduced rows [rows, hidden].""" + handoff with ``defer``, else the unreduced rows [rows, hidden]. ``reduce`` overrides + the handoff's completion reduction; the draft passes identity because its layer owns + the tensor-parallel all-reduce.""" from sglang.kernels.ops.elementwise.elementwise import fused_gate_sigmoid_mul_add from sglang.srt.layers.layer_boundary.fusions.moe_finalize import ( MoeDeferredFinalize, @@ -171,7 +173,7 @@ def step(layer, x: torch.Tensor, st, defer: bool): routed, gated_shared_output=shared, m=rows, - reduce=sum_post_experts_output, + reduce=sum_post_experts_output if reduce is None else reduce, ) if shared is not None: if use_fused_gate: diff --git a/sgl_nvidia_mega/qwen3_5/moe/mkm.py b/sgl_nvidia_mega/qwen3_5/moe/mkm.py index 85333b29..c3a07a28 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/mkm.py +++ b/sgl_nvidia_mega/qwen3_5/moe/mkm.py @@ -170,6 +170,7 @@ def routed( x_sf, routing=None, wide_bal: bool = False, + m16_max_rows: int = 16, ): """The deferred-finalize output of ``mlp``'s routed experts from the router front's logits and quantized input (``maps``: the layer's weight TMA maps; ``identity``: @@ -194,6 +195,7 @@ def routed( experts.g2_alphas, routing=routing, wide_bal=wide_bal, + m16_max_rows=m16_max_rows, ) deferred = FlashInferTrtllmDeferredFinalizeOutput( gemm2_out=gemm2_out, diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index c4d35b5a..bfb73f64 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -132,6 +132,9 @@ class Bound(msgspec.Struct, frozen=True): m256: bool = False # slot mkm_m256: the 256-token engine built (on m32's maps) m512: bool = False # slot mkm_m512: the 512-token engine built (on m32's maps) wide_bal: bool = False # opt-in balanced work split for mkm_m256 / mkm_m512 + # A caller-specific upper bound for the 16-token engine. The draft uses 12 so + # its alternate-stream shared expert can overlap rows 13-16 with the 32-token engine. + m16_max_rows: int = 16 reason: str = "" # why a requested slot is not built @@ -383,7 +386,15 @@ def _routed_experts(mlp, x, fused, logits, x_fp4, x_sf): experts = mlp.experts if fused is not None: maps = fused.mkm[mlp.layer_id] - return mkm.routed(mlp, maps, fused.identity, logits, x_fp4, x_sf) + return mkm.routed( + mlp, + maps, + fused.identity, + logits, + x_fp4, + x_sf, + m16_max_rows=fused.m16_max_rows, + ) topk_output = mlp.topk(x, logits) if not TopKOutputChecker.format_is_bypassed(topk_output): raise RuntimeError("the router front requires logits-based bypassed TopK") @@ -451,7 +462,15 @@ def routed_m32(mlp, x: torch.Tensor, bound: Bound, gemm: bool = False): else: logits, _ = mlp.gate(x) _, _, x_fp4, x_sf = router.route_topk_quant(logits, x, scale) - return mkm.routed(mlp, bound.m32[mlp.layer_id], bound.identity, logits, x_fp4, x_sf) + return mkm.routed( + mlp, + bound.m32[mlp.layer_id], + bound.identity, + logits, + x_fp4, + x_sf, + m16_max_rows=bound.m16_max_rows, + ) def routed_given(mlp, x: torch.Tensor, bound: Bound, slot: str): @@ -480,4 +499,5 @@ def routed_given(mlp, x: torch.Tensor, bound: Bound, slot: str): x_sf, routing=(ids, weights), wide_bal=bound.wide_bal, + m16_max_rows=bound.m16_max_rows, ) diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index cbb5b3c9..0471698e 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -63,6 +63,8 @@ "draft_q1", "draft_xchg", "draft_hop", + # numeric: hand deferred MoE rows to the draft hop's fused final norm + "draft_fin", "draft_front", "draft_skv2", # draft_loop and draft_extend also serve TP8's draft graphs of 2 to 128 requests, @@ -223,6 +225,10 @@ def slot_reasons() -> list[str]: reasons.append(f"{SLOTS_ENV}: slot {name!r} runs inside entry 'draft_loop'") if "draft_hop" in SLOTS and "hop" not in SLOTS: reasons.append(f"{SLOTS_ENV}: slot 'draft_hop' runs on slot 'hop''s mailbox") + if "draft_fin" in SLOTS and "draft_hop" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'draft_fin' runs in slot 'draft_hop''s final norm" + ) if "draft_multi" in SLOTS and "multi_request" not in SLOTS: reasons.append( f"{SLOTS_ENV}: slot 'draft_multi' serves the draft steps of two or more " diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_fc_sum.cu b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_fc_sum.cu new file mode 100644 index 00000000..1d63b51e --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/draft_fc_sum.cu @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// The row-split FC's one-launch cross-rank sum: every rank's fp32 partial [m, 4096] is pushed through the +// multicast mapping and summed on every rank in rank order from rank 0, rounded once to bf16, in +// one launch: ordered_partial_sum's arithmetic on the bytes fc_gather moves, without the +// [m, world * 4096] intermediate or a second launch. Lamport transport: fp32 words, all ones (a +// NaN no partial carries) as "not arrived", three stages rotated by a control block, every +// consumer re-arming the words it read. Thread t of CTA c owns columns [4j, 4j + 4) of row +// c / 8, j = (c % 8) * 128 + t. +// +// Built with NVRTC for MKD_WORLD (the TP size). + +#include "mkd_common.cuh" + +#ifndef MKD_WORLD +#define MKD_WORLD 8 +#endif + +namespace sglang { + +namespace fc_allsum { + +inline constexpr int kStages = 3; +inline constexpr int kMaxTokens = 16; // the split fc's widest batch, as fc_gather +inline constexpr int kThreads = 128; +inline constexpr int kCols = 4; // fp32 columns a thread owns: one 16-byte word group +inline constexpr uint32_t kNotArrived = 0xFFFFFFFFu; +inline constexpr uint32_t kSpinLimit = 1u << 27; + +// Per-process control block in device memory, zero at build. +struct Ctrl { + uint32_t stage; // stage of the next call + uint32_t arrive; // CTA tickets of the running call +}; + +SGL_DEVICE uint32_t ld_volatile_u32(const uint32_t* p) { + uint32_t v; + asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +SGL_DEVICE uint4 ld_volatile_v4(const void* p) { + uint4 v; + asm volatile("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p) + : "memory"); + return v; +} + +SGL_DEVICE void st_v4(void* p, uint4 v) { + asm volatile("st.global.v4.u32 [%0], {%1,%2,%3,%4};" ::"l"(p), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w) : "memory"); +} + +SGL_DEVICE void mc_st_v4(void* mc, uint4 v) { + asm volatile("multimem.st.relaxed.sys.global.v4.f32 [%0], {%1,%2,%3,%4};" ::"l"(mc), "f"(__uint_as_float(v.x)), + "f"(__uint_as_float(v.y)), "f"(__uint_as_float(v.z)), "f"(__uint_as_float(v.w)) + : "memory"); +} + +// A partial word equal to the "not arrived" pattern (a NaN) is sent as another NaN. +SGL_DEVICE uint32_t sendable(uint32_t w) { + return w == kNotArrived ? 0x7FFFFFFFu : w; +} + +SGL_DEVICE uint32_t pending_mask(uint4 v) { + return (v.x == kNotArrived) | (v.y == kNotArrived) | (v.z == kNotArrived) | (v.w == kNotArrived); +} + +} // namespace fc_allsum + +struct FcAllsumParams { + const float* partial; // [m, kH] this rank's fp32 partial, written by an earlier grid + bf16_t* out; // [m, kH] the partials summed in rank order, rounded once + float* mc; // multicast base of the mailbox [kStages, world, kMaxTokens, kH] fp32 + float* uc; // this rank's copy of it + fc_allsum::Ctrl* ctrl; // this rank's control block + int32_t rank; + int32_t m; +}; + +/** + * \brief One call of the split fc's cross-rank sum: push this rank's partial, wait for every + * rank's, sum them as ordered_partial_sum does (acc = rank 0, then + rank 1, ... in fp32 round to + * nearest) and round once to bf16. \tparam kH hidden size, \tparam kWorld TP size. The grid is + * m * (kH / 512) CTAs of 128 threads, all resident (no CTA waits on another CTA of its rank). + */ +template +SGL_DEVICE void fc_allsum_kernel(const FcAllsumParams& p) { + using namespace fc_allsum; + constexpr int kCtasPerRow = kH / (kCols * kThreads); + static_assert(kCtasPerRow * kCols * kThreads == kH, "a row is one word group per thread of its CTAs"); + __shared__ uint32_t s_stage; + + device::PDLWaitPrimary(); + uint32_t ticket = 0; + if (threadIdx.x == 0) { + s_stage = ld_volatile_u32(&p.ctrl->stage); + // taken after the stage is read: the last ticket's advance cannot reach this call + ticket = atomicAdd(&p.ctrl->arrive, 1u); + } + __syncthreads(); + const int64_t stage = s_stage; + const int row = static_cast(blockIdx.x) / kCtasPerRow; + const int col = ((static_cast(blockIdx.x) % kCtasPerRow) * kThreads + static_cast(threadIdx.x)) * kCols; + const auto word = [&](int src) { return ((stage * kWorld + src) * kMaxTokens + row) * kH + col; }; + + uint4 own = *reinterpret_cast(p.partial + static_cast(row) * kH + col); + own = make_uint4(sendable(own.x), sendable(own.y), sendable(own.z), sendable(own.w)); + mc_st_v4(p.mc + word(p.rank), own); + + uint4 v[kWorld]; +#pragma unroll + for (int r = 0; r < kWorld; ++r) v[r] = ld_volatile_v4(p.uc + word(r)); + uint32_t pending = 0; +#pragma unroll + for (int r = 0; r < kWorld; ++r) pending |= pending_mask(v[r]) << r; + uint32_t spins = 0; + while (pending) { + if (++spins > kSpinLimit) __trap(); + uint32_t still = 0; +#pragma unroll + for (int r = 0; r < kWorld; ++r) { + if (pending & (1u << r)) { + v[r] = ld_volatile_v4(p.uc + word(r)); + still |= pending_mask(v[r]) << r; + } + } + pending = still; + } + + float acc[kCols] = {__uint_as_float(v[0].x), __uint_as_float(v[0].y), __uint_as_float(v[0].z), + __uint_as_float(v[0].w)}; +#pragma unroll + for (int r = 1; r < kWorld; ++r) { + acc[0] = __fadd_rn(acc[0], __uint_as_float(v[r].x)); + acc[1] = __fadd_rn(acc[1], __uint_as_float(v[r].y)); + acc[2] = __fadd_rn(acc[2], __uint_as_float(v[r].z)); + acc[3] = __fadd_rn(acc[3], __uint_as_float(v[r].w)); + } + __nv_bfloat162 lo = __floats2bfloat162_rn(acc[0], acc[1]); + __nv_bfloat162 hi = __floats2bfloat162_rn(acc[2], acc[3]); + uint2 packed = make_uint2(*reinterpret_cast(&lo), *reinterpret_cast(&hi)); + *reinterpret_cast(p.out + static_cast(row) * kH + col) = packed; + + const uint4 rearm = make_uint4(kNotArrived, kNotArrived, kNotArrived, kNotArrived); +#pragma unroll + for (int r = 0; r < kWorld; ++r) st_v4(p.uc + word(r), rearm); + + if (threadIdx.x == 0 && ticket == gridDim.x - 1) { + p.ctrl->arrive = 0u; + p.ctrl->stage = static_cast((stage + 1) % kStages); + } +} + +} // namespace sglang + +using sglang::FcAllsumParams; + +extern "C" __global__ void __launch_bounds__(128) fc_allsum(const FcAllsumParams __grid_constant__ p) { + sglang::fc_allsum_kernel<4096, MKD_WORLD>(p); +} diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_fc.cu b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_fc.cu index f3a4c8cd..6ee3d5ce 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_fc.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_fc.cu @@ -88,6 +88,24 @@ SGL_DEVICE void gemma_norm_row( __syncthreads(); } +/// \brief The dot of slice_dot from weights already in registers (the same order). +template +SGL_DEVICE float slice_dot_regs(const device::AlignedVector* wv, const bf16_t* s_x) { + using namespace device; + using vec_t = AlignedVector; + constexpr int kDotVecs = kW / (8 * kWarpThreads); + const int lane = static_cast(get_lane_id()); + float acc = 0.0f; +#pragma unroll + for (int k = 0; k < kDotVecs; ++k) { + vec_t xv; + xv.load(s_x, k * kWarpThreads + lane); +#pragma unroll + for (int j = 0; j < 8; ++j) acc = fmaf(cast(wv[k][j]), cast(xv[j]), acc); + } + return warp::reduce_sum(acc); +} + /// \brief One slice partial: warp dot of a weight row segment [kW] with s_x[kW]. template SGL_DEVICE float slice_dot(const bf16_t* w_row, const bf16_t* s_x) { @@ -272,19 +290,48 @@ struct MkdHiddenPartialParams { * the step-boundary exchange waits on the flag instead of a stream join, so it keeps its * programmatic early launch behind the lm_head. */ -template +template SGL_DEVICE void mkd_hidden_partial_kernel(const MkdHiddenPartialParams& p) { using namespace device; + using vec_t = AlignedVector; constexpr int kW = 2 * kH / kWorld; constexpr int kWarps = 8; + constexpr int kDotVecs = kW / (8 * kWarpThreads); __shared__ alignas(16) bf16_t s_x[kH]; __shared__ float s_red[kWarps]; const int warp = static_cast(threadIdx.x / kWarpThreads); const int lane = static_cast(get_lane_id()); mkd_fc::gemma_norm_row(p.h, p.norm_w, p.eps, s_x, s_red); - for (int j = static_cast(blockIdx.x) * kWarps + warp; j < kH; j += static_cast(gridDim.x) * kWarps) { - const float acc = mkd_fc::slice_dot(p.w + static_cast(j) * kW, s_x + p.col); - if (lane == 0) p.out[j] = acc; + const int stride = static_cast(gridDim.x) * kWarps; + if constexpr (kRowsPerIter == 1) { + for (int j = static_cast(blockIdx.x) * kWarps + warp; j < kH; j += stride) { + const float acc = mkd_fc::slice_dot(p.w + static_cast(j) * kW, s_x + p.col); + if (lane == 0) p.out[j] = acc; + } + } else { + // The warp's rows are the same as above (j, j + stride, ...); the weight loads of kRowsPerIter of + // them are issued together (the partial runs beside the draft head, whose stream saturates HBM, so + // one row's load latency at a time left it the slowest input of the exchange). Each row's dot is + // slice_dot's arithmetic. + for (int j0 = static_cast(blockIdx.x) * kWarps + warp; j0 < kH; j0 += stride * kRowsPerIter) { + vec_t wv[kRowsPerIter][kDotVecs]; +#pragma unroll + for (int r = 0; r < kRowsPerIter; ++r) { + const int j = j0 + r * stride; + if (j < kH) { +#pragma unroll + for (int k = 0; k < kDotVecs; ++k) wv[r][k].load(p.w + static_cast(j) * kW, k * kWarpThreads + lane); + } + } +#pragma unroll + for (int r = 0; r < kRowsPerIter; ++r) { + const int j = j0 + r * stride; + if (j < kH) { + const float acc = mkd_fc::slice_dot_regs(wv[r], s_x + p.col); + if (lane == 0) p.out[j] = acc; + } + } + } } __syncthreads(); if (threadIdx.x == 0) { @@ -321,3 +368,8 @@ extern "C" __global__ void __launch_bounds__(256) mkd_fc_memo(const MkdFcMemoPar extern "C" __global__ void __launch_bounds__(256, 1) mkd_hidden_partial(const MkdHiddenPartialParams __grid_constant__ p) { sglang::mkd_hidden_partial_kernel<4096, MKD_WORLD>(p); } + +// Four rows' weight loads in flight per warp (same values). +extern "C" __global__ void __launch_bounds__(256, 1) mkd_hidden_partial_r4(const MkdHiddenPartialParams __grid_constant__ p) { + sglang::mkd_hidden_partial_kernel<4096, MKD_WORLD, 4>(p); +} diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_xchg.cu b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_xchg.cu index 001cce50..d94ea7c3 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_xchg.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_xchg.cu @@ -119,7 +119,7 @@ struct MkdXchgParams { * [c * H / kChunks, ...) of the candidate and partial rows. Every CTA redoes the shard * argmax (a few tens of KB from L2) so no CTA waits for another before sending. */ -template +template SGL_DEVICE void mkd_xchg_kernel(const MkdXchgParams& p) { using namespace device; using vec_t = AlignedVector; @@ -319,7 +319,12 @@ SGL_DEVICE void mkd_xchg_kernel(const MkdXchgParams& p) { } if (p.rank >= kEmbed) { const float* part = p.own_partial + c * kCols; - for (int w = static_cast(threadIdx.x); w < kPartWords; w += kThreads) { + // With kSplitSend the partial goes out from the threads after the candidate, so + // a hidden rank's two loads are not serialized in the same threads. + constexpr int kPartFirst = kSplitSend && kCandWords + kPartWords <= kThreads ? kCandWords : 0; + constexpr int kPartStride = kPartFirst > 0 ? kThreads - kPartFirst : kThreads; + for (int w = static_cast(threadIdx.x) - kPartFirst; w < kPartWords; w += kPartStride) { + if (w < 0) break; if (p.partial_flag != nullptr) { // The partial comes from a side-stream kernel that is not a stream predecessor. uint32_t f, n = 0; @@ -480,3 +485,8 @@ extern "C" __global__ void __launch_bounds__(1024) mkd_xchg_memo(const MkdXchgPa extern "C" __global__ void __launch_bounds__(1024) mkd_xchg_emb(const MkdXchgParams __grid_constant__ p) { sglang::mkd_xchg_kernel<4096, MKD_WORLD, false, 16, 1024, true>(p); } + +// Memo exchange with the hidden partial sent by its own threads. +extern "C" __global__ void __launch_bounds__(1024) mkd_xchg_memo_split(const MkdXchgParams __grid_constant__ p) { + sglang::mkd_xchg_kernel<4096, MKD_WORLD, true, 16, 1024, true, true>(p); +} diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py b/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py index a9eebfb4..fb089f59 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py @@ -53,6 +53,10 @@ # fc_gather: the split fc's widest batch (2.2 _FC_SPLIT_MAX_TOKENS) and its grid GATHER_MAX_TOKENS = 16 GATHER_THREADS = 512 +# fc_allsum (draft_fc_sum.cu): 128 threads a CTA, 4 fp32 columns a thread, so 8 CTAs a row +ALLSUM_THREADS = 128 +ALLSUM_CTAS_PER_ROW = HIDDEN // (4 * ALLSUM_THREADS) +ALLSUM_STAGES = 3 def _order_key(value: float) -> int: @@ -64,6 +68,9 @@ def _order_key(value: float) -> int: # -1e30; torch.argmax ranks NaN above every value (2.2 draft_vocab_argmax.py) NAN_KEY_SPLIT_ARGMAX = _order_key(-1e30) NAN_KEY_TORCH_ARGMAX = 0xFFFFFFFF +# Hidden ranks send the memo row and their side-stream partial from disjoint thread +# ranges. Tests may set this internal policy false for a bitwise lockstep. +SPLIT_MEMO_SEND = True class MkdXchgParams(ctypes.Structure): @@ -160,6 +167,14 @@ class FcGatherParams(ctypes.Structure): ] + [("rank", ctypes.c_int32), ("m", ctypes.c_int32)] +class FcAllsumParams(ctypes.Structure): + """sglang::FcAllsumParams.""" + + _fields_ = [ + (name, ctypes.c_void_p) for name in ("partial", "out", "mc", "uc", "ctrl") + ] + [("rank", ctypes.c_int32), ("m", ctypes.c_int32)] + + def _sm_count(device_index: int) -> int: return torch.cuda.get_device_properties(device_index).multi_processor_count @@ -180,7 +195,7 @@ def _xchg_kernels(device_index: int, world: int) -> dict[str, jit.Kernel]: cubin = _compile(device_index, "mkd_xchg", (f"MKD_WORLD={world}",)) return { name: jit.load(cubin, name, (XCHG_THREADS, 1, 1), 0) - for name in ("mkd_xchg_memo", "mkd_xchg_emb") + for name in ("mkd_xchg_memo", "mkd_xchg_emb", "mkd_xchg_memo_split") } @@ -198,6 +213,9 @@ def _fc_kernels(device_index: int, world: int) -> dict[str, jit.Kernel]: cubin, "mkd_fc_memo", (256, 1, 1), MEMO_BLOCK_ROWS * HIDDEN * 2 ), "mkd_hidden_partial": jit.load(cubin, "mkd_hidden_partial", (256, 1, 1), 0), + "mkd_hidden_partial_r4": jit.load( + cubin, "mkd_hidden_partial_r4", (256, 1, 1), 0 + ), } @@ -226,6 +244,13 @@ def _draft_fc_kernels(device_index: int, world: int, max_rows: int) -> dict: return kernels +@cache +def _fc_sum_kernel(device_index: int, world: int) -> jit.Kernel: + """The split FC's one-launch gather and rank-order sum.""" + cubin = _compile(device_index, "draft_fc_sum", (f"MKD_WORLD={world}",)) + return jit.load(cubin, "fc_allsum", (ALLSUM_THREADS, 1, 1), 0) + + def build(device_index: int, world: int) -> None: """Compile every kernel of the draft step boundary at TP ``world`` (no compile inside a capture).""" @@ -423,7 +448,11 @@ def mkd_xchg( ) kernels = _xchg_kernels(device.index, world) _launch( - kernels["mkd_xchg_memo" if memo else "mkd_xchg_emb"], + kernels[ + ("mkd_xchg_memo_split" if SPLIT_MEMO_SEND else "mkd_xchg_memo") + if memo + else "mkd_xchg_emb" + ], CHUNKS, p, device, @@ -520,12 +549,13 @@ def mkd_hidden_partial( sync: torch.Tensor, world: int, ctas: int = 0, + rows_per_iter: int = 4, ) -> None: """Hidden-slice fc partial [1, H] fp32 of ``h`` [1, H] (draft_fc_partial arithmetic over ``shard`` [H, W] at hidden-half column ``col``); the last CTA then sets ``sync[1]`` to the exchange epoch ``epoch_done[0]`` (``sync[0]`` counts CTAs, zero - between calls). ``ctas`` > 0 launches that many CTAs instead of one per SM (same - values: each row is one warp's).""" + between calls). ``ctas`` > 0 launches that many CTAs instead of one per SM. The + production four-row load schedule is bitwise equal to ``rows_per_iter=1``.""" device = out.device width = 2 * HIDDEN // world _check(h, "h", (1, HIDDEN), torch.bfloat16, aligned=True) @@ -549,7 +579,10 @@ def mkd_hidden_partial( ) kernels = _fc_kernels(device.index, world) grid = ctas if ctas > 0 else _sm_count(device.index) - _launch(kernels["mkd_hidden_partial"], grid, p, device, pdl=False) + if rows_per_iter not in (1, 4): + raise ValueError("mkd_hidden_partial: rows_per_iter must be 1 or 4") + name = "mkd_hidden_partial_r4" if rows_per_iter == 4 else "mkd_hidden_partial" + _launch(kernels[name], grid, p, device, pdl=False) # ---------------------------------------------------------------- draft lm_head @@ -742,3 +775,54 @@ def __call__(self, partial: torch.Tensor) -> torch.Tensor: # one CTA an SM: with the batched polls every thread waits about one round trip _launch(self._kernels["fc_gather"], self._ctas, p, self.device, pdl=False) return out + + +class FcAllsum: + """The split fc's gather and rank-order sum in one launch (draft_fc_sum.cu fc_allsum): + [m, H] fp32 partials on every rank -> [m, H] bf16, summed in rank order from rank 0 and + rounded once, as ``FcGather`` + ``ordered_partial_sum``; m <= GATHER_MAX_TOKENS; + graph-capturable. Its own mailbox and control block (not FcGather's epochs). Build on + every rank of ``group`` (a collective).""" + + def __init__( + self, + *, + world: int, + rank: int, + group, + cpu_group, + device: torch.device, + ): + import torch.distributed as dist + + if world not in WORLDS: + raise ValueError(f"the fc all-sum serves TP {WORLDS}, not {world}") + self.world, self.rank, self.device = world, rank, device + words = ALLSUM_STAGES * world * GATHER_MAX_TOKENS * HIDDEN // 2 + self._buf, self._handle, self.mc = symmetric_words(words, group, device) + # every fp32 word "not arrived" (all ones) + self._buf.fill_(-1) + self.uc = self._buf.data_ptr() + # [stage, arrive] (fc_allsum::Ctrl) + self.ctrl = torch.zeros(2, dtype=torch.int32, device=device) + self._kernel = _fc_sum_kernel(device.index, world) + torch.cuda.synchronize(device) + dist.barrier(group=cpu_group) + + def __call__(self, partial: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + m = int(partial.shape[0]) + if not 1 <= m <= GATHER_MAX_TOKENS: + raise ValueError(f"fc all-sum: {m} rows (1-{GATHER_MAX_TOKENS})") + _check(partial, "partial", (m, HIDDEN), torch.float32, aligned=True) + _check(out, "out", (m, HIDDEN), torch.bfloat16, aligned=True) + p = FcAllsumParams( + partial=partial.data_ptr(), + out=out.data_ptr(), + mc=self.mc, + uc=self.uc, + ctrl=self.ctrl.data_ptr(), + rank=self.rank, + m=m, + ) + _launch(self._kernel, m * ALLSUM_CTAS_PER_ROW, p, self.device, pdl=True) + return out diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py b/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py index f62c009c..ed169457 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/mkm.py @@ -204,8 +204,14 @@ class ParamsM32(Params): ROUTED_ENGINES = (M64_MAX_TOKENS, M128_MAX_TOKENS, M256_MAX_TOKENS, M512_MAX_TOKENS) -def engine_rows(m: int) -> int: - """The bound of the engine that serves ``m`` rows.""" +def engine_rows(m: int, m16_max_rows: int = M16_MAX_TOKENS) -> int: + """The bound of the engine that serves ``m`` rows. + + A caller may hand the tail of the 16-token range to the 32-token engine. Both + consume the same router logits and NVFP4 rows and produce identical values. + """ + if m16_max_rows < m <= M16_MAX_TOKENS: + return M32_MAX_TOKENS return next(r for r in ENGINE_ROWS if m <= r) @@ -391,6 +397,7 @@ def mkm_moe( g2_alpha: torch.Tensor, routing: tuple | None = None, wide_bal: bool = False, + m16_max_rows: int = M16_MAX_TOKENS, ) -> tuple[torch.Tensor, torch.Tensor]: """(gemm2_out bf16 [M * 10, 4096], expert_weights bf16 [M, 10]) for logits [M, 512] bf16 (M <= 7, at TP8 M <= 128) and the router front's x_fp4 [M, 2048] / x_sf [M, 256] @@ -407,7 +414,7 @@ def mkm_moe( ) if not 1 <= m <= widest: raise ValueError(f"mkm: {m} rows (1-{widest})") - rows = engine_rows(m) + rows = engine_rows(m, m16_max_rows) if routing is not None and rows not in ROUTED_ENGINES: raise ValueError( f"mkm: a given routing needs a 64-512-token engine, not {m} rows" diff --git a/test/kernels/test_qwen3_5_draft_xchg.py b/test/kernels/test_qwen3_5_draft_xchg.py index 6014026d..400432c7 100644 --- a/test/kernels/test_qwen3_5_draft_xchg.py +++ b/test/kernels/test_qwen3_5_draft_xchg.py @@ -10,6 +10,8 @@ rounded once to bf16, bit for bit; - the memo rows equal the split fc's partials summed in rank order (checked at construction), and the split fc's gather and rank-order sum are exact; +- the one-launch all-sum equals the gather and rank-order sum bit for bit, also from a + CUDA graph; - a whole draft round (five steps, the extend exchange with extend fc) replays from a CUDA graph with fresh inputs, equal to its eager run; - the sharded top-1 of slot draft_multi (``sm103/draft/top1``) picks, for 1-16 rows, the @@ -118,7 +120,14 @@ def _want_token(local_logits, group, nan_as): return int(torch.argmax(full)) # the first maximum: the lowest id -def test_step_exchange_picks_the_global_argmax(ranks, boundary): +@pytest.mark.parametrize("split_send", [False, True]) +def test_step_exchange_picks_the_global_argmax( + ranks, boundary, split_send, monkeypatch +): + """The split-send and serial-send schedules produce the same boundary outputs.""" + from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels + + monkeypatch.setattr(kernels, "SPLIT_MEMO_SEND", split_send) rank, world, group, _ = ranks mkd, _ = boundary n_embed = world // 2 @@ -212,6 +221,64 @@ def test_extend_exchange_nan_rules(ranks, boundary, nan_key): assert int(topk_index) == want +def test_extend_mkd_head_is_the_scan_path(ranks, boundary): + """The one-request MKD head, side-stream partial and maxima exchange equal the + shard GEMV, logits scan and serial partial, eager and under graph replay.""" + from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels + + rank, world, group, cpu_group = ranks + mkd, _ = boundary + + def scan_path(x): + logits = torch.empty(1, VOCAB, dtype=torch.bfloat16, device="cuda") + part = torch.zeros(mkd.head_part.numel(), dtype=torch.int32, device="cuda") + kernels.mkd_draft_head(x, mkd.lm_head_weight, logits, part) + mkd.note_draft_input() + _, ti = mkd.exchange_extend(logits, x) + return [logits, ti.clone(), mkd.fc_buf.clone()] + + def head_path(x): + logits = mkd.extend_head(x) + assert logits is not None + mkd.note_draft_input() + _, ti = mkd.exchange_extend(logits, x, mkd_head=True) + return [logits, ti.clone(), mkd.fc_buf.clone()] + + def same(got, want): + assert torch.equal(got[0].view(torch.int16), want[0].view(torch.int16)) + assert int(got[1]) == int(want[1]) == _want_token(want[0], group, -1e30) + assert torch.equal(got[2].view(torch.int16), want[2].view(torch.int16)) + + with torch.inference_mode(): + rows = [_randn(1, H, g=_gen(7300 + s)) for s in range(4)] + for x in rows: + want = scan_path(x) + got = head_path(x) + torch.cuda.synchronize() + same(got, want) + dist.barrier(group=cpu_group) + static = rows[0].clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + head_path(static) + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + dist.barrier(group=cpu_group) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = head_path(static) + torch.cuda.synchronize() + dist.barrier(group=cpu_group) + for x in rows: + want = scan_path(x) + static.copy_(x) + graph.replay() + torch.cuda.synchronize() + same(captured, want) + dist.barrier(group=cpu_group) + + @pytest.mark.parametrize("m", [1, 2, 7]) def test_split_fc_gathers_and_sums_in_rank_order(ranks, boundary, m): from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels @@ -241,6 +308,121 @@ def test_split_fc_gathers_and_sums_in_rank_order(ranks, boundary, m): torch.testing.assert_close(got.float(), ref.float(), atol=3e-2, rtol=3e-2) +@pytest.fixture(scope="module") +def allsum(ranks): + from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels + + rank, world, group, cpu_group = ranks + return kernels.FcAllsum( + world=world, + rank=rank, + group=group, + cpu_group=cpu_group, + device=torch.device("cuda", torch.cuda.current_device()), + ) + + +def _fp32_partials(m, rank, seed): + """Rank ``rank``'s fp32 partials, with signed zeros in the first columns of every rank.""" + p = torch.randn(m, H, device="cuda", generator=_gen(seed + 31 * rank)) * 4.0 + p[:, :8] = -0.0 + return p + + +def _gather_ordered_sum(mkd, partial, world): + from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels + + out = torch.empty(partial.shape, dtype=torch.bfloat16, device="cuda") + return kernels.ordered_partial_sum(mkd._fc_gather(partial), world, out) + + +@pytest.mark.parametrize("m", [1, 2, 7, 16]) +def test_fc_allsum_is_the_gather_and_ordered_sum(ranks, boundary, allsum, m): + """One launch gives FcGather + ordered_partial_sum's rows bit for bit, -0.0 + sums included, over calls that visit every stage twice.""" + from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels + + rank, world, group, _ = ranks + mkd, _ = boundary + with torch.inference_mode(): + for call in range(2 * kernels.ALLSUM_STAGES): + partial = _fp32_partials(m, rank, 1000 * m + call) + got = allsum( + partial, torch.empty(m, H, dtype=torch.bfloat16, device="cuda") + ) + want = _gather_ordered_sum(mkd, partial, world) + torch.cuda.synchronize() + assert torch.equal(got.view(torch.int16), want.view(torch.int16)), call + + +def test_fc_allsum_replays_from_a_cuda_graph(ranks, boundary, allsum): + rank, world, group, cpu_group = ranks + mkd, _ = boundary + m = NUM_STEPS + 1 + with torch.inference_mode(): + partial = _fp32_partials(m, rank, 5000) + out = torch.empty(m, H, dtype=torch.bfloat16, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + allsum(partial, out) + torch.cuda.current_stream().wait_stream(stream) + dist.barrier(group=cpu_group) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + allsum(partial, out) + torch.cuda.synchronize() + dist.barrier(group=cpu_group) + for call in range(4): + partial.copy_(_fp32_partials(m, rank, 5100 + call)) + graph.replay() + want = _gather_ordered_sum(mkd, partial, world) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), want.view(torch.int16)), call + dist.barrier(group=cpu_group) + + +@pytest.mark.parametrize("ctas", [8, 24, 0]) +def test_hidden_partial_rows4_is_the_one_row_kernel(ranks, boundary, ctas): + """The four-row load schedule equals the one-row kernel at every hidden slice.""" + from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels + + rank, world, _, _ = ranks + mkd, weights = boundary + width = mkd.width + with torch.inference_mode(): + h = _randn(1, H, g=_gen(7100 + rank + ctas)) + shard = _randn(H, width, g=_gen(7200 + rank), scale=0.02) + got = {} + for rows_per_iter in (1, 4): + for col in range(0, H, width): + out = torch.full( + (1, H), float("nan"), dtype=torch.float32, device="cuda" + ) + epoch_done = torch.tensor([11, 0], dtype=torch.int32, device="cuda") + sync = torch.zeros(2, dtype=torch.int32, device="cuda") + kernels.mkd_hidden_partial( + h, + shard, + weights["norm_hid"], + EPS, + col, + out, + epoch_done, + sync, + world, + ctas=ctas, + rows_per_iter=rows_per_iter, + ) + torch.cuda.synchronize() + assert sync.tolist() == [0, 11], (rows_per_iter, col) + got[(rows_per_iter, col)] = out + for col in range(0, H, width): + assert torch.equal( + got[(4, col)].view(torch.int32), got[(1, col)].view(torch.int32) + ), col + + def test_round_replays_from_a_cuda_graph(ranks, boundary): rank, world, group, cpu_group = ranks mkd, _ = boundary diff --git a/test/kernels/test_qwen3_5_moe.py b/test/kernels/test_qwen3_5_moe.py index 57bfc97f..fb9a817d 100644 --- a/test/kernels/test_qwen3_5_moe.py +++ b/test/kernels/test_qwen3_5_moe.py @@ -112,7 +112,7 @@ def _inputs(m, seed): ) -def _fused(layer, logits, x_fp4, x_sf): +def _fused(layer, logits, x_fp4, x_sf, **kwargs): return mkm.mkm_moe( layer["maps"], logits, @@ -123,6 +123,7 @@ def _fused(layer, logits, x_fp4, x_sf): layer["g1"], layer["g1c"], layer["g2"], + **kwargs, ) @@ -187,6 +188,21 @@ def test_mkm_is_the_trtllm_routed_experts(layer, m, seed, shared_route): assert not differ.any(), f"rows {differ.nonzero().flatten().tolist()} differ" +@pytest.mark.parametrize("m", [13, 14, 16]) +@pytest.mark.parametrize("seed", [0, 1]) +def test_draft_32_token_engine_is_the_trtllm_routed_experts(layer, m, seed): + """The draft's rows past 12 use the 32-token engine with identical values.""" + if layer["inter"] not in mkm.M32_ENGINES: + pytest.skip("no 32-token engine at this intermediate") + logits, x_fp4, x_sf = _inputs(m, seed) + assert mkm.engine_rows(m, m16_max_rows=12) == mkm.M32_MAX_TOKENS + rows, weights = _fused(layer, logits, x_fp4, x_sf, m16_max_rows=12) + want_rows, want_weights = _stock(layer, logits, x_fp4, x_sf) + assert torch.equal(weights.view(torch.int16), want_weights.view(torch.int16)) + differ = (rows.view(torch.int16) != want_rows.view(torch.int16)).any(dim=1) + assert not differ.any(), f"rows {differ.nonzero().flatten().tolist()} differ" + + @pytest.mark.parametrize("m", [7, 14, 28, 56]) def test_mkm_replays_from_a_cuda_graph(layer, m): if m > mkm.max_tokens(layer["inter"], m32=True, m64=True): diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 4ffd3573..28e016fd 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -400,6 +400,10 @@ def attend( [ ({"draft_moe"}, "slot 'draft_moe' runs inside entry 'draft_loop'"), ({"draft_loop", "draft_hop"}, "slot 'draft_hop' runs on slot 'hop''s mailbox"), + ( + {"draft_loop", "draft_fin"}, + "slot 'draft_fin' runs in slot 'draft_hop''s final norm", + ), ( {"draft_loop", "draft_skv2"}, "slot 'draft_skv2' runs inside entry 'draft_extend'", @@ -420,6 +424,117 @@ def test_the_draft_slots_name_what_they_need(monkeypatch, slots, want): assert [r.split(": ", 1)[1] for r in runtime.slot_reasons()] == [want] +def test_the_draft_target_path_uses_its_validated_engine_ranges(): + from sgl_nvidia_mega.qwen3_5.draft import moe_mkm + from sgl_nvidia_mega.qwen3_5.sm103.moe import mkm + + draft = moe_mkm.DraftMkm(None, wide=True) + assert [ + draft.takes(rows, per_request) + for rows, per_request in ((1, 1), (2, 1), (7, 7), (64, 1), (65, 1)) + ] == [ + False, + True, + True, + True, + False, + ] + assert [mkm.engine_rows(rows, m16_max_rows=12) for rows in (7, 12, 13, 16, 17)] == [ + mkm.MAX_TOKENS, + mkm.M16_MAX_TOKENS, + mkm.M32_MAX_TOKENS, + mkm.M32_MAX_TOKENS, + mkm.M32_MAX_TOKENS, + ] + + +@pytest.mark.parametrize( + "fused, path, taken, want", + [ + (False, "sglang", True, ["mlp", "all_reduce", "norm"]), + (False, "mkm", True, ["mkm", "complete", "all_reduce", "norm"]), + (True, "sglang", True, ["mlp deferred", "fin"]), + (True, "mkm", True, ["mkm", "fin"]), + (True, "sglang", False, ["mlp deferred", "complete", "norm"]), + (True, "mkm", False, ["mkm", "complete", "all_reduce", "norm"]), + ], +) +def test_the_draft_extends_moe_output_is_reduced_once( + monkeypatch, fused, path, taken, want +): + from sgl_nvidia_mega.qwen3_5.draft import layer + + rows, log = 14, [] + x = torch.ones(rows, 4) + routed = torch.bfloat16 if taken else torch.float32 + + def owed(reduced): + def complete(): + log.append("complete") + return 8 * x if reduced else x + + return SimpleNamespace( + m=rows, + routed_output=torch.zeros(1, dtype=routed), + expert_weights=torch.zeros(1, dtype=torch.bfloat16), + permuted_indices=torch.zeros(1, dtype=torch.int32), + gated_shared_output=torch.zeros(1, dtype=torch.bfloat16), + complete=complete, + ) + + def mlp(h, defer_finalize=False): + log.append("mlp deferred" if defer_finalize else "mlp") + return owed(reduced=True) if defer_finalize else h + + def fin(*, residual, gamma, **_): + log.append("fin") + return residual + gamma, residual + + def norm(h, residual): + log.append("norm") + return h + residual, residual + + def mkm(lay, h): + log.append("mkm") + return owed(reduced=False) + + mlp.supports_deferred_finalize = True + mkm.takes = lambda r, rows_per_request: True + norm.gemma_weight = 100.0 + extend_hop = SimpleNamespace( + hop=SimpleNamespace( + supports=lambda m: m <= rows, + moe_finalize_all_reduce_rms_norm=fin, + ), + ffn_input=lambda h, residual, norm, moe_single_launch: (h, residual), + ) + parts = SimpleNamespace( + layer=SimpleNamespace( + input_layernorm=lambda h: h, + post_attention_layernorm=None, + mlp=mlp, + ), + attention=lambda h, *args, rows_per_request, o_proj: h, + hop=extend_hop, + moe=None, + moe_mkm=mkm if path == "mkm" else None, + slots=frozenset({"draft_hop", "draft_fin"} if fused else {"draft_hop"}), + model=SimpleNamespace(norm=norm), + ) + + def all_reduce(h): + log.append("all_reduce") + return 8 * h + + monkeypatch.setattr( + "sglang.srt.distributed.tensor_model_parallel_all_reduce", all_reduce + ) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + out = layer.layer(parts, x, None, None, None, None, rows_per_request=7) + assert log == want + assert torch.equal(out, x + 100.0 if "fin" in log else 8 * x + x) + + class _DraftParts: """Parts with no slot built: the stock model step.""" From 3689c9fb934d441a3516d5f1e719456d55d1421a Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 16:03:55 -0700 Subject: [PATCH 11/62] perf: tune validated small decode paths --- sgl_nvidia_mega/qwen3_5/attention/__init__.py | 10 +- sgl_nvidia_mega/qwen3_5/attention/mka.py | 86 +- sgl_nvidia_mega/qwen3_5/boundary/hop.py | 34 +- sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py | 7 +- sgl_nvidia_mega/qwen3_5/gdn/__init__.py | 55 + sgl_nvidia_mega/qwen3_5/gdn/conv.py | 1 + sgl_nvidia_mega/qwen3_5/gdn/front.py | 90 +- sgl_nvidia_mega/qwen3_5/gdn/verify.py | 45 +- sgl_nvidia_mega/qwen3_5/lowm.py | 6 +- sgl_nvidia_mega/qwen3_5/moe/__init__.py | 4 + sgl_nvidia_mega/qwen3_5/moe/route.py | 36 +- sgl_nvidia_mega/qwen3_5/plan.py | 54 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 33 + sgl_nvidia_mega/qwen3_5/runtime/executor.py | 5 + .../qwen3_5/sm103/attention/mka.py | 33 + .../qwen3_5/sm103/attention/skv2.py | 1 + .../qwen3_5/sm103/attention/src/mka.cu | 170 ++ .../qwen3_5/sm103/draft/q1_kernel.py | 1 + sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py | 42 +- sgl_nvidia_mega/qwen3_5/sm103/gdn/front.py | 27 +- sgl_nvidia_mega/qwen3_5/sm103/gdn/src/core.cu | 210 ++- .../qwen3_5/sm103/gdn/src/front.cu | 61 +- .../qwen3_5/sm103/gdn/trt_front.py | 124 ++ sgl_nvidia_mega/qwen3_5/sm103/hop/__init__.py | 217 ++- sgl_nvidia_mega/qwen3_5/sm103/hop/src/hop.cu | 25 +- sgl_nvidia_mega/qwen3_5/sm103/hop/src/rlin.cu | 447 +++++ .../qwen3_5/sm103/trt_lowm/__init__.py | 318 ++++ .../qwen3_5/sm103/trt_lowm/lowm_linear.py | 291 ++++ .../qwen3_5/sm103/trt_lowm/lowm_mma_gemm.py | 1518 +++++++++++++++++ .../qwen3_5/sm103/trt_lowm/utils.py | 483 ++++++ sgl_nvidia_mega/qwen3_5/spec.py | 12 + test/kernels/test_qwen3_5_gdn.py | 457 ++++- test/kernels/test_qwen3_5_hop.py | 62 + test/unit/test_qwen3_5_spec.py | 52 + test/unit/test_qwen3_5_steps.py | 243 ++- 35 files changed, 5175 insertions(+), 85 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/gdn/trt_front.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/hop/src/rlin.cu create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/__init__.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/lowm_linear.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/lowm_mma_gemm.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/utils.py diff --git a/sgl_nvidia_mega/qwen3_5/attention/__init__.py b/sgl_nvidia_mega/qwen3_5/attention/__init__.py index 0abe5d7a..e4254cb0 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/attention/__init__.py @@ -47,7 +47,7 @@ "srt/models/qwen3_5.py:Qwen3_5AttentionDecoderLayer.self_attention": "db0f4d98e0fb1abf", } -ATTENTION_SLOTS = ("skv2", "skv2_cp", "mka", "mka_push") +ATTENTION_SLOTS = ("skv2", "skv2_cp", "mka", "mka_push", "route_lin") class Bound(msgspec.Struct, frozen=True): @@ -57,6 +57,7 @@ class Bound(msgspec.Struct, frozen=True): mka: bool = False mka_wide: bool = False # mka's wide entries serve the steps of several requests mka_push: bool = False + route_lin: bool = False # router-before-all-reduce packed for every layer reason: str = "" # why a requested slot is not built @@ -93,14 +94,17 @@ def build(model, ctx) -> Bound: if reason: return Bound(reason=f"slot skv2_cp: {reason}") if "mka" in SLOTS: - for _, layer in layers: - mka.pack(layer) + for i, layer in layers: + reason = mka.pack(layer, route_lin="route_lin" in SLOTS) + if reason: + return Bound(reason=f"slot route_lin: layer {i}: {reason}") mka_on = "mka" in SLOTS return Bound( skv2="skv2" in SLOTS, mka=mka_on, mka_wide=mka_on and all(layer.num_heads == 4 for _, layer in layers), mka_push="mka_push" in SLOTS, + route_lin="route_lin" in SLOTS, ) diff --git a/sgl_nvidia_mega/qwen3_5/attention/mka.py b/sgl_nvidia_mega/qwen3_5/attention/mka.py index 4c0bc94c..08e3c353 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/mka.py +++ b/sgl_nvidia_mega/qwen3_5/attention/mka.py @@ -65,6 +65,7 @@ ), "srt/layers/dp_attention.py": ("is_allocation_symmetric",), "srt/runtime_context.py": ("get_parallel",), + "srt/layers/layer_boundary/fusions/norm.py": ("fused_norm_gamma",), "srt/model_executor/forward_batch_info.py": ("out_cache_loc",), "srt/kernel_plugins/nvidia/eagle_binding.py|srt/megakernels/eagle_binding.py": ( "kv", @@ -130,8 +131,11 @@ def pool_reasons(kv) -> list[str]: return [] -def pack(layer) -> None: - """Pack ``layer``'s qkv (+ gate) and o_proj weights for the kernels (at build).""" +def pack(layer, route_lin: bool = False) -> str: + """Pack ``layer``'s qkv (+ gate) and o_proj weights at build. + + Returns why its router-before-all-reduce weights cannot be packed, when requested. + """ from sgl_nvidia_mega.qwen3_5.sm103.attention import mka as kernels kernels.build(layer.qkv_proj.weight.device) @@ -139,6 +143,52 @@ def pack(layer) -> None: kernels.pack_qkv(layer.qkv_proj.weight.t(), layer.num_heads), kernels.pack_o(layer.o_proj.weight.t()), ) + return _pack_route_lin(layer) if route_lin else "" + + +# Slot route_lin (numeric): layer id -> router-before-all-reduce weights on this TP8 rank. +RL_WEIGHTS: dict = {} + + +def _pack_route_lin(layer) -> str: + import torch + import torch.distributed as dist + from sglang.srt.layers.layer_boundary.fusions.norm import fused_norm_gamma + from sglang.srt.runtime_context import get_parallel + + from sgl_nvidia_mega.qwen3_5.sm103.hop import RL_WORLD, pack_router_lin + + group = get_parallel().tp_group + mlp = getattr(layer, "mlp", None) + if group.world_size != RL_WORLD: + return f"TP {group.world_size} (serves TP {RL_WORLD})" + if layer.num_heads != 4: + return f"{layer.num_heads} query heads a rank (serves 4)" + boundary = getattr(layer, "ffn_boundary", None) + norm = getattr(getattr(boundary, "plan", None), "norm", None) + gamma = fused_norm_gamma(norm) if norm is not None else None + layer_id = getattr(layer, "layer_id", None) + if ( + mlp is None + or getattr(mlp, "gate", None) is None + or gamma is None + or layer_id != layer.attn.layer_id + or getattr(mlp, "layer_id", layer_id) != layer_id + ): + return "no router or fusible FFN-input norm" + try: + with torch.no_grad(): + RL_WEIGHTS[layer.attn.layer_id] = pack_router_lin( + mlp.gate.weight, + gamma, + layer.o_proj.weight.t(), + layer.o_proj.input_scale, + layer.o_proj.weight_scale, + dist.get_rank(group.device_group), + ) + except ValueError as error: + return str(error) + return "" def step(layer, x, st): @@ -201,6 +251,23 @@ def step(layer, x, st): ): out = torch.empty((m, layer.hidden_size), dtype=torch.bfloat16, device=x.device) hop = _push_hop(layer, st, m) if st.slots.mka_push else None + h0 = _route_lin_h0(hop, layer, m, st) + if h0 is not None: + # Slot route_lin: the router partial rides the o_proj push. + kernels.epilogue_rl( + attn=o, + gate=gate, + w_packed=w_o, + sx=layer.o_proj.input_scale, + sw=layer.o_proj.weight_scale, + out=out, + site=hop.site, + rl_args=hop.rl_push_args(RL_WEIGHTS[attn.layer_id], h0), + ) + hop.mark_pushed_rl(out) + # the router launch does not wait on the epilogue: keep its output's memory out of reuse for the whole step + st.scratch.setdefault("rlin.keep", []).append(out) + return out kernels.epilogue( attn=o, gate=gate, @@ -216,6 +283,21 @@ def step(layer, x, st): return out +def _route_lin_h0(hop, layer, rows: int, st): + """The layer's residual input when its router partial can ride the o_proj push (as gdn/verify.py's); else None.""" + if ( + not st.slots.route_lin + or hop is None + or hop.rl is None + or layer.attn.layer_id not in RL_WEIGHTS + ): + return None + lid, h0 = st.scratch.get("rlin.h0", (None, None)) + if lid != layer.layer_id or h0 is None or tuple(h0.shape) != (rows, 4096): + return None + return h0 + + def _front_step(layer, x, st): """The front path of a wide step: SGLang's qkv GEMM, the front with the FP8 KV write, the attention, the gate with o_proj's input quantization, SGLang's o_proj GEMM.""" diff --git a/sgl_nvidia_mega/qwen3_5/boundary/hop.py b/sgl_nvidia_mega/qwen3_5/boundary/hop.py index 9b1c3458..b12a5251 100644 --- a/sgl_nvidia_mega/qwen3_5/boundary/hop.py +++ b/sgl_nvidia_mega/qwen3_5/boundary/hop.py @@ -174,6 +174,7 @@ def build(model, ctx) -> Bound: group=group.device_group, device=torch.device("cuda", torch.cuda.current_device()), deadline_s=ctx.deadline_s, + route_lin="route_lin" in runtime.SLOTS, ) global _ACTIVE _ACTIVE = hop @@ -318,6 +319,9 @@ def layer_input(hop, layer, hidden, st): norm, residual, *fp8 = _finalize( hop, "layer input", owed, residual, fused_norm_gamma(plan.norm), scale ) + if st.slots.route_lin: + # Slot route_lin: the layer input read by this rank's router partial. + st.scratch["rlin.h0"] = (layer.layer_id, residual) stream.write(residual) if fp8: lowm.stash_fp8(st, norm, fp8[0]) @@ -369,12 +373,30 @@ def moe_input(hop, layer, hidden, st): _engage("MoE input", rows) scale = lowm.moe_input_fp8_scale(layer, st) - norm, residual, *fp8 = hop.all_reduce_residual_rms_norm( - local_contribution=part, - residual=residual, - gamma=fused_norm_gamma(plan.norm), - **({} if scale is None else {"fp8_scale": scale}), - ) + if hop._pushed_rl: + # Slot route_lin: the mixer pushed its router partial too; this hop also gives + # the MoE its logits and NVFP4 input (moe/route.py routed_lin). + lid, h0 = st.scratch.get("rlin.h0", (None, None)) + if lid != layer.layer_id or h0 is None or h0.data_ptr() != residual.data_ptr(): + raise RuntimeError( + f"route-lin: layer {layer.layer_id}'s residual is not the one its router partials read" + ) + _engage("MoE input + router (route-lin)", rows) + (norm, residual, *fp8), rl = hop.all_reduce_residual_rms_norm_rl( + local_contribution=part, + residual=residual, + gamma=fused_norm_gamma(plan.norm), + global_scale=layer.mlp.experts.w13_input_scale_quant.reshape(1), + **({} if scale is None else {"fp8_scale": scale}), + ) + st.scratch["rlin.moe"] = (layer.layer_id, *rl) + else: + norm, residual, *fp8 = hop.all_reduce_residual_rms_norm( + local_contribution=part, + residual=residual, + gamma=fused_norm_gamma(plan.norm), + **({} if scale is None else {"fp8_scale": scale}), + ) stream.write(residual) if fp8: lowm.stash_fp8(st, norm, fp8[0]) diff --git a/sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py b/sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py index e03ff8a1..b0080813 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py +++ b/sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py @@ -81,12 +81,7 @@ def build(cls, mlp) -> tuple[Optional[DraftMkm], str]: route=True, mkm={mlp.layer_id: maps}, mkm_wide=wide, - # the hop's finalize takes registered identity rows without a map load - identity=torch.arange( - mkm_kernel.M128_MAX_TOKENS * mkm_kernel.TOP_K, - dtype=torch.int32, - device="cuda", - ), + identity=route.identity_rows(mkm_kernel.M128_MAX_TOKENS * mkm_kernel.TOP_K), topk=True, m32={mlp.layer_id: maps}, m64=True, diff --git a/sgl_nvidia_mega/qwen3_5/gdn/__init__.py b/sgl_nvidia_mega/qwen3_5/gdn/__init__.py index d04c767a..6a0e95f5 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/__init__.py @@ -48,6 +48,8 @@ ), "srt/layers/attention/hybrid_linear_attn_backend.py": ("linear_attn_backend",), "srt/model_executor/forward_context.py": ("get_attn_backend",), + "srt/layers/layer_boundary/fusions/norm.py": ("fused_norm_gamma",), + "srt/runtime_context.py": ("get_parallel",), } # The SGLang functions whose body this module mirrors, by the digest of their code @@ -70,6 +72,9 @@ class Bound(msgspec.Struct, frozen=True): packed: dict = msgspec.field(default_factory=dict) # slot gdn_front: layer id -> in_proj_qkvz's weights packed for the front front: dict = msgspec.field(default_factory=dict) + gdn_trt_front: bool = False # TensorRT-LLM front built through four requests + # slot route_lin (numeric): layer id -> router-before-all-reduce packed weights + rlin: dict = msgspec.field(default_factory=dict) error: object = None # the front's sticky error word, int32, zero reason: str = "" # why a requested slot is not built @@ -103,6 +108,11 @@ def build(model, ctx) -> Bound: packed[gdn.attn.layer_id] = pack_out_proj(gdn.out_proj.weight.t()) except ValueError as e: return Bound(reason=f"layer {i}: {e}") + rlin = {} + if packed and "route_lin" in SLOTS: + rlin, reason = _route_lin(m, gdns) + if reason: + return Bound(reason=f"slot route_lin: {reason}") front_packed, error = {}, None if "gdn_front" in SLOTS: import torch @@ -138,6 +148,8 @@ def build(model, ctx) -> Bound: gdn_conv="gdn_conv" in SLOTS, packed=packed, front=front_packed, + gdn_trt_front="gdn_trt_front" in SLOTS, + rlin=rlin, error=error, ) @@ -202,3 +214,46 @@ def step(layer, x, st): core = core.reshape(*core.shape[:-2], core.shape[-2] * core.shape[-1]) out, _ = gdn.out_proj(core) return out + + +def _route_lin(m, gdns) -> tuple[dict, str]: + """Packed route-lin weights for every GDN layer, or why they cannot be built.""" + import torch + import torch.distributed as dist + from sglang.srt.layers.layer_boundary.fusions.norm import fused_norm_gamma + from sglang.srt.runtime_context import get_parallel + + from sgl_nvidia_mega.qwen3_5.sm103.hop import RL_WORLD, pack_router_lin + + group = get_parallel().tp_group + if group.world_size != RL_WORLD: + return {}, f"TP {group.world_size} (serves TP {RL_WORLD})" + rank = dist.get_rank(group.device_group) + out = {} + for i, gdn in gdns: + layer = m.layers[i] + mlp = getattr(layer, "mlp", None) + boundary = getattr(layer, "ffn_boundary", None) + norm = getattr(getattr(boundary, "plan", None), "norm", None) + gamma = fused_norm_gamma(norm) if norm is not None else None + if ( + mlp is None + or getattr(mlp, "gate", None) is None + or gamma is None + or layer.layer_id != gdn.attn.layer_id + or getattr(mlp, "layer_id", layer.layer_id) != layer.layer_id + ): + return {}, f"layer {i} has no router or fusible FFN-input norm" + try: + with torch.no_grad(): + out[gdn.attn.layer_id] = pack_router_lin( + mlp.gate.weight, + gamma, + gdn.out_proj.weight.t(), + gdn.out_proj.input_scale, + gdn.out_proj.weight_scale, + rank, + ) + except ValueError as error: + return {}, f"layer {i}: {error}" + return out, "" diff --git a/sgl_nvidia_mega/qwen3_5/gdn/conv.py b/sgl_nvidia_mega/qwen3_5/gdn/conv.py index 387fe5fc..21c0ae02 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/conv.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/conv.py @@ -227,6 +227,7 @@ def kernel( ): if USE_GDC: tl.extra.cuda.gdc_wait() + # Keep this release after the wait: ReplaySSM (no wait_for_gate_producer) reads a / b before its wait. tl.extra.cuda.gdc_launch_dependents() idx_seq = tl.program_id(0) idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) diff --git a/sgl_nvidia_mega/qwen3_5/gdn/front.py b/sgl_nvidia_mega/qwen3_5/gdn/front.py index 58c80184..2453d7b9 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/front.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/front.py @@ -31,6 +31,7 @@ # The SGLang names this module reads (test/unit/test_contract.py pins them) SGLANG_READS = { "srt/models/qwen3_5.py": ( + "Qwen3_5LinearDecoderLayer", "in_proj_qkvz", "in_proj_ba", "attn_tp_size", @@ -189,26 +190,100 @@ def build(model, gdns) -> dict: } +def prepare(lm, pools, trt: bool = False) -> None: + """Compile the requested TensorRT-LLM front on the first mamba-pool layouts, + before graph capture. Every GDN layer has the same layouts.""" + from sglang.srt.models.qwen3_5 import Qwen3_5LinearDecoderLayer + + from sgl_nvidia_mega.qwen3_5.sm103.gdn import trt_front + + if not trt: + return + m = lm.model + gdn = next( + m.layers[i].linear_attn + for i in range(m.start_layer, m.end_layer) + if isinstance(m.layers[i], Qwen3_5LinearDecoderLayer) + ) + li = pools.layer_ids.index(gdn.attn.layer_id) + trt_front.prepare( + gdn.in_proj_qkvz.weight.shape[1], + gdn.in_proj_ba.weight.shape[0], + gdn.attn.conv_weights, + pools.conv[0][li], + pools.intermediate_conv_window[0][li], + ) + + def step(gdn, x: torch.Tensor, st, opts) -> torch.Tensor: """``Qwen3_5GatedDeltaNet.forward`` of a target-verify step of 7 rows on the front: the TP-partial output [7, hidden] of the normed BF16 rows x; ``opts`` is slot gdn_verify's (the verify's tiles, and ``push`` for slot gdn_push).""" + from sgl_nvidia_mega.qwen3_5 import lowm from sgl_nvidia_mega.qwen3_5.gdn import replayssm from sgl_nvidia_mega.qwen3_5.gdn.front_norm import norm from sgl_nvidia_mega.qwen3_5.gdn.verify import _out_proj - from sgl_nvidia_mega.qwen3_5.sm103.gdn import front + from sgl_nvidia_mega.qwen3_5.sm103.gdn import core as core_kernel + from sgl_nvidia_mega.qwen3_5.sm103.gdn import front, trt_front - if x.shape[0] not in _ENGAGED: - _ENGAGED.add(x.shape[0]) + rows = x.shape[0] + requests = rows // ROWS + # TensorRT-LLM's front may read the hop's E4M3 copy of x (slot hop_fp8). + trt = st.slots.gdn_trt_front and opts.core + x8 = lowm.take_fp8(st, x) if trt and st.slots.hop_fp8 else None + # the L2 prefetch of the checkpoint / ring in front of slot gdn_core is left out with the early history (the core + # holds those loads ahead of its wait) and for the pinned one-request front (which completes later with it) + prefetch = not (opts.core and (core_kernel.EARLY or (front.PIN and requests == 1))) + if (rows, x8 is not None) not in _ENGAGED: + _ENGAGED.add((rows, x8 is not None)) logger.info( - "NVIDIA megakernel: Qwen3.5 slot gdn_front runs at %d rows", x.shape[0] + "NVIDIA megakernel: Qwen3.5 slot gdn_front runs at %d rows " + "(pinned=%d, core_early=%d, core_flat=%d, trt_front=%d, " + "fp8_x=%d, prefetch=%d)", + rows, + front.PIN, + core_kernel.EARLY, + core_kernel.FLAT, + trt, + x8 is not None, + int(prefetch and not trt), ) attn, bound = gdn.attn, st.bound["gdn"] meta, pools = st.view.gdn, st.view.mamba li = pools.layer_ids.index(attn.layer_id) - projected = x.new_empty((ROWS, 2560)) - conv = x.new_empty((ROWS, CONV_CHANNELS)) - ba = x.new_empty((ROWS, 16)) + if trt: + from sgl_nvidia_mega.qwen3_5.gdn.verify import _fused + + out = x.new_empty((rows, gdn.in_proj_qkvz.weight.shape[1])) + ba = x.new_empty((rows, gdn.in_proj_ba.weight.shape[0])) + trt_front.front( + x, + x8, + # the method keeps the weight as the [K, N] transpose of a row-major [N, K] + gdn.in_proj_qkvz.weight.t(), + gdn.in_proj_qkvz.input_scale, + gdn.in_proj_qkvz.weight_scale, + gdn.in_proj_ba.weight, + attn.conv_weights, + pools.conv[0][li], + pools.intermediate_conv_window[0][li], + meta.state_indices[:requests], + out, + ba, + ) + # convolved q | k | v, the raw z, and b / a in the front's b8 / a8 order + return _fused( + gdn, + out[:, :CONV_CHANNELS], + out[:, CONV_CHANNELS:].view(rows, attn.num_v_heads, attn.head_v_dim), + ba[:, attn.num_v_heads :], + ba[:, : attn.num_v_heads], + st, + opts, + ) + projected = x.new_empty((rows, 2560)) + conv = x.new_empty((rows, CONV_CHANNELS)) + ba = x.new_empty((rows, 16)) front.front( x, bound.front[attn.layer_id], @@ -230,6 +305,7 @@ def step(gdn, x: torch.Tensor, st, opts) -> torch.Tensor: history_g=pools.ring_g[li], ring_indices=st.fb.req_pool_indices[:1], pad_slot=meta.pad_slot_id, + prefetch=prefetch, ) if opts.core: from sgl_nvidia_mega.qwen3_5.gdn.verify import _fused diff --git a/sgl_nvidia_mega/qwen3_5/gdn/verify.py b/sgl_nvidia_mega/qwen3_5/gdn/verify.py index 178f258e..4a220508 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/verify.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/verify.py @@ -399,7 +399,11 @@ def _fused(gdn, qkv, z, a, b, st, opts) -> torch.Tensor: """Slot gdn_core: the ReplaySSM verify of the convolved rows ``qkv``, the gated norm and its static FP8 output in one launch (``sm103/gdn/core.py``), then ``out_proj`` (FlashInfer's BMM, or the hop's push GEMM with ``opts.push``).""" - from sgl_nvidia_mega.qwen3_5.sm103.gdn.core import core + from sgl_nvidia_mega.qwen3_5.sm103.gdn.core import ( + RELEASE_MAX_ROWS, + RELEASE_MIN_ROWS, + core, + ) attn, pools = gdn.attn, st.view.mamba li = pools.layer_ids.index(attn.layer_id) @@ -412,6 +416,19 @@ def _fused(gdn, qkv, z, a, b, st, opts) -> torch.Tensor: def ring(t): return None if t is None else t[li] + # cuBLASLt's programmatic-dependent out-projection may start at the core's entry + # over the validated row window. + release = ( + RELEASE_MIN_ROWS <= rows <= RELEASE_MAX_ROWS + and not opts.push + and not st.slots.lowm_outproj + ) + if release and ("release", rows) not in _ENGAGED: + _ENGAGED.add(("release", rows)) + logger.info( + "NVIDIA megakernel: Qwen3.5 slot gdn_core releases out_proj's BMM at its entry at %d rows", + rows, + ) core( qkv, z, @@ -436,10 +453,22 @@ def ring(t): x, rows_per_request=st.binding.tokens_per_request, num_k_heads=attn.num_k_heads, - early_dependents=opts.push, + early_dependents=opts.push or release, ) if opts.push: - return st.bound["boundary"].hop.out_proj_push( + hop = st.bound["boundary"].hop + h0 = _route_lin_h0(hop, attn.layer_id, rows, st) + if h0 is not None: + # Slot route_lin: the router partial rides the out_proj push. + return hop.out_proj_push_rl( + x_fp8=x, + packed_w=st.bound["gdn"].packed[attn.layer_id], + input_scale=gdn.out_proj.input_scale, + weight_scale=gdn.out_proj.weight_scale, + rl_w=st.bound["gdn"].rlin[attn.layer_id], + h0=h0, + ) + return hop.out_proj_push( x_fp8=x, packed_w=st.bound["gdn"].packed[attn.layer_id], input_scale=gdn.out_proj.input_scale, @@ -448,6 +477,16 @@ def ring(t): return _out_proj(gdn.out_proj, x, qkv.dtype, st, opts.lowm_push) +def _route_lin_h0(hop, layer_id: int, rows: int, st): + """The layer input when slot route_lin can ride the out_proj push; else None.""" + if not st.slots.route_lin or hop.rl is None or layer_id not in st.bound["gdn"].rlin: + return None + lid, h0 = st.scratch.get("rlin.h0", (None, None)) + if lid != layer_id or h0 is None or tuple(h0.shape) != (rows, 4096): + return None + return h0 + + def _core(attn, mixed_qkv, a, b, st, opts): """The backend's target-verify arm for the circular ReplaySSM ring.""" from sgl_nvidia_mega.qwen3_5.gdn import replayssm diff --git a/sgl_nvidia_mega/qwen3_5/lowm.py b/sgl_nvidia_mega/qwen3_5/lowm.py index 9151e006..2864aa30 100644 --- a/sgl_nvidia_mega/qwen3_5/lowm.py +++ b/sgl_nvidia_mega/qwen3_5/lowm.py @@ -413,7 +413,7 @@ def take_fp8(st, x: torch.Tensor): def layer_input_fp8_scale(layer, st, rows: int): """The static input scale of ``layer``'s mixer projection that reads the layer input in this step (slot hop_fp8), or None where none would read an E4M3 copy (gdn_front's and MKA's fused entries quantize - in their own launch).""" + in their own launch, except slot gdn_trt_front, which reads it).""" from sglang.srt.models.qwen3_5 import Qwen3_5LinearDecoderLayer from sgl_nvidia_mega.qwen3_5.attention import mka @@ -421,7 +421,9 @@ def layer_input_fp8_scale(layer, st, rows: int): if not st.slots.hop_fp8: return None if isinstance(layer, Qwen3_5LinearDecoderLayer): - if not st.slots.gdn_verify or st.slots.gdn_front: + if not st.slots.gdn_verify: + return None + if st.slots.gdn_front and not st.slots.gdn_trt_front: return None return layer.linear_attn.in_proj_qkvz.input_scale if st.slots.mka and rows <= mka.FUSED_MAX_ROWS: diff --git a/sgl_nvidia_mega/qwen3_5/moe/__init__.py b/sgl_nvidia_mega/qwen3_5/moe/__init__.py index b87ac113..2457e2b7 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/moe/__init__.py @@ -235,6 +235,10 @@ def _router_experts(mlp, x, defer: bool, st): from sglang.srt.layers.moe.topk import TopKOutputChecker if defer and st.slots.route: + rl = st.scratch.pop("rlin.moe", None) + if rl is not None and rl[0] == mlp.layer_id and st.slots.mkm: + # Slot route_lin produced the logits and NVFP4 input in the preceding hop. + return route.routed_lin(mlp, x, st.bound["moe"], *rl[1:]) return route.routed(mlp, x, st.bound["moe"] if st.slots.mkm else None) if defer and st.slots.mkm_m32: return route.routed_m32(mlp, x, st.bound["moe"], gemm=st.slots.route_gemm) diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index bfb73f64..b330960d 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -188,15 +188,20 @@ def build( for mlp in layers }, mkm_wide=wide, - identity=torch.arange( - mkm_kernel.M16_MAX_TOKENS * mkm_kernel.TOP_K, - dtype=torch.int32, - device="cuda", - ), + identity=identity_rows(mkm_kernel.M16_MAX_TOKENS * mkm_kernel.TOP_K), shared=shared, ) +def identity_rows(n: int): + """Return and register an int32 identity row map for the hop finalize.""" + from sgl_nvidia_mega.qwen3_5.sm103.hop import register_identity_rows + + rows = torch.arange(n, dtype=torch.int32, device="cuda") + register_identity_rows(rows) + return rows + + def build_m32( model, bound: Bound, @@ -245,11 +250,7 @@ def build_m32( or mkm_kernel.weight_maps(mlp.experts.w13_weight, mlp.experts.w2_weight) for mlp in layers }, - identity=torch.arange( - max(hi, mkm_kernel.M128_MAX_TOKENS) * mkm_kernel.TOP_K, - dtype=torch.int32, - device="cuda", - ), + identity=identity_rows(max(hi, mkm_kernel.M128_MAX_TOKENS) * mkm_kernel.TOP_K), ) @@ -373,6 +374,21 @@ def routed(mlp, x: torch.Tensor, fused: Bound | None = None, shared: bool = Fals return _routed_experts(mlp, x, fused, logits, x_fp4, x_sf) +def routed_lin(mlp, x: torch.Tensor, fused: Bound, logits, x_fp4, x_sf): + """``routed`` with the router front's outputs taken from the MoE input's hop. + + Slot ``route_lin`` is numeric for the logits from the mixer's router partials; its + NVFP4 input is bitwise equal to the router front's. Slot ``mkm`` runs on them. + """ + if ("route_lin", x.shape[0]) not in _ENGAGED: + _ENGAGED.add(("route_lin", x.shape[0])) + logger.info( + "NVIDIA megakernel: Qwen3.5 route-lin (router before the all-reduce) runs at %d rows", + x.shape[0], + ) + return _routed_experts(mlp, x, fused, logits, x_fp4, x_sf) + + def _routed_experts(mlp, x, fused, logits, x_fp4, x_sf): """The routed experts on the router front's outputs (``fused``: slot mkm's kernel), deferred.""" diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index 5c54a948..f812ba60 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -18,7 +18,9 @@ widths where FlashInfer's low-latency kernels, which it equals bit for bit, run. ``gdn_push`` (numeric: out_proj as the hop's push GEMM, on gdn_verify's FP8 norm) runs up to 8 rows. ``gdn_front`` (numeric: the GDN projections and conv in one launch, then -gdn_pdl's verify) covers one request of 7 rows, ``route_shared`` (numeric: the shared +gdn_pdl's verify) covers one request of 7 rows; ``gdn_trt_front`` replaces its projection +front through four requests. ``route_lin`` (numeric) starts the next router before the +current hidden all-reduce at exactly one request. ``route_shared`` (numeric: the shared expert's gate/up in route's launch) up to 8 rows. ``skv2`` runs the full attention's core of one request through the split-KV attention of SGLang branch nvidia/qwen35-agentx-gb300, ``mka`` the fused attention prologue and epilogue around it @@ -44,7 +46,9 @@ MULTI_REQUEST_TP = (8,) # the name that requests them in the slot list MULTI_REQUEST = "multi_request" -GDN_FRONT_ROWS = 7 # the rows slot gdn_front's front is built for +GDN_FRONT_ROWS = 7 # the one-request custom front's fixed width +GDN_TRT_FRONT_MAX_ROWS = 28 # TensorRT-LLM's front, through four requests + # Each slot's widest step: its launcher's bound (sm103/moe/router.py MAX_TOKENS, # sm103/moe/mkm.py MAX_TOKENS, sm103/hop OPROJ_MAX_M, sm103/attention/mka.py), and for # hop the rows up to which FlashInfer's MNNVL CuTe DSL fusion runs the one-shot @@ -107,11 +111,15 @@ class Slots(msgspec.Struct, frozen=True): gdn_push: bool = False # the projections and the conv in one launch (gdn/front.py; numeric) gdn_front: bool = False + # TensorRT-LLM's front replaces the custom front through four requests (numeric) + gdn_trt_front: bool = False # the shared expert's gate/up in the router launch (moe/route.py; numeric) route_shared: bool = False skv2: bool = False # the full attention's core through SKV2 (attention/split_kv.py) mka: bool = False # the fused attention prologue and epilogue (attention/mka.py) mka_push: bool = False # the epilogue pushes its rows into the hop + # the next router runs before the mixer's hidden all-reduce (numeric) + route_lin: bool = False # the input projections on the low-M launch (lowm.py; numeric) lowm_in: bool = False lowm_outproj: bool = False # GDN out_proj on the low-M launch (lowm.py; numeric) @@ -146,6 +154,7 @@ class Facts(msgspec.Struct, frozen=True): hop: bool = False # the decode hop, built for the widest step it covers gdn_push: bool = False # out_proj packed for the push GEMM on every GDN layer gdn_front: bool = False # in_proj_qkvz packed for the front on every GDN layer + gdn_trt_front: bool = False # TensorRT-LLM's front built for every GDN layer route_shared: bool = ( False # the router front with the shared stage, every MoE layer ) @@ -153,6 +162,7 @@ class Facts(msgspec.Struct, frozen=True): mka: bool = False # the fused prologue and epilogue for every full-attention layer mka_wide: bool = False # with its wide entries (four query heads a rank) mka_push: bool = False # with the epilogue's push into the hop + route_lin: bool = False # the router-before-all-reduce path built for every layer lowm_in: bool = False # the low-M input projections built for every layer lowm_out: bool = False # the low-M output projections built for every layer lowm_se: bool = False # the low-M shared expert built for every MoE layer @@ -168,13 +178,28 @@ def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> hop = facts.hop and rows <= HOP_MAX_ROWS mka = facts.mka and rows <= (MKA_WIDE_MAX_ROWS if facts.mka_wide else MKA_MAX_ROWS) lowm_se = facts.lowm_se and rows <= LOWM_SE_MAX_ROWS + mkm = ( + route + and facts.mkm + and rows <= (MKM_WIDE_MAX_ROWS if facts.mkm_wide else MKM_MAX_ROWS) + ) + gdn_front_base = facts.gdn_verify and facts.gdn_pdl and facts.gdn_front + gdn_trt_front = ( + gdn_front_base + and facts.gdn_trt_front + and rows % GDN_FRONT_ROWS == 0 + and rows <= GDN_TRT_FRONT_MAX_ROWS + ) + # At one request lowm_se replaces both route_shared's shared stage and + # SGLang's shared-expert tail. + route_shared = ( + route and facts.route_shared and rows <= ROUTE_SHARED_ROWS and not lowm_se + ) return Slots( route=route, route_topk=route_topk, route_gemm=route_topk and facts.route_gemm and rows <= ROUTE_GEMM_MAX_ROWS, - mkm=route - and facts.mkm - and rows <= (MKM_WIDE_MAX_ROWS if facts.mkm_wide else MKM_MAX_ROWS), + mkm=mkm, mkm_m32=route_topk and facts.mkm_m32 and MKM_WIDE_MAX_ROWS < rows <= MKM_M32_MAX_ROWS, @@ -199,20 +224,19 @@ def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> and hop and facts.gdn_push and rows <= GDN_PUSH_MAX_ROWS, - gdn_front=facts.gdn_verify - and facts.gdn_pdl - and facts.gdn_front - and rows == GDN_FRONT_ROWS, - # At one request lowm_se replaces both route_shared's shared stage and - # SGLang's shared-expert tail. - route_shared=route - and facts.route_shared - and rows <= ROUTE_SHARED_ROWS - and not lowm_se, + gdn_front=gdn_front_base and (rows == GDN_FRONT_ROWS or gdn_trt_front), + gdn_trt_front=gdn_trt_front, + route_shared=route_shared, # SKV2 attends one sequence skv2=facts.skv2 and rows <= tokens_per_request, mka=mka, mka_push=mka and facts.mka_push and hop, + route_lin=facts.route_lin + and rows == GDN_FRONT_ROWS + and route + and mkm + and hop + and not route_shared, lowm_in=facts.lowm_in and rows <= LOWM_IN_MAX_ROWS, lowm_outproj=facts.lowm_out and rows <= LOWM_OUTPROJ_MAX_ROWS, lowm_oproj=facts.lowm_out and rows <= LOWM_OPROJ_MAX_ROWS, diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 0471698e..d7ded9f3 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -36,7 +36,9 @@ "hop", "gdn_push", "gdn_front", + "gdn_trt_front", "route_shared", + "route_lin", "skv2", "mka", "mka_push", @@ -210,6 +212,11 @@ def slot_reasons() -> list[str]: reasons.append( f"{SLOTS_ENV}: slot 'gdn_front' runs slot 'gdn_pdl''s dependent verify" ) + if "gdn_trt_front" in SLOTS and not {"gdn_front", "gdn_core"} <= SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'gdn_trt_front' replaces slot 'gdn_front' in front " + "of slot 'gdn_core'" + ) if "gdn_core" in SLOTS and "gdn_verify" not in SLOTS: reasons.append( f"{SLOTS_ENV}: slot 'gdn_core' runs slot 'gdn_verify''s step in one launch" @@ -264,6 +271,27 @@ def slot_reasons() -> list[str]: f"{SLOTS_ENV}: slot 'mka_push' pushes slot 'mka''s epilogue into slot " "'hop''s mailbox" ) + if ( + "route_lin" in SLOTS + and not { + "route", + "mkm", + "hop", + "gdn_push", + "mka", + "mka_push", + } + <= SLOTS + ): + reasons.append( + f"{SLOTS_ENV}: slot 'route_lin' runs the router before the all-reduce " + "through slots 'route', 'mkm', 'hop', 'gdn_push', 'mka' and 'mka_push'" + ) + if "route_lin" in SLOTS and "route_shared" in SLOTS and "lowm_se" not in SLOTS: + reasons.append( + f"{SLOTS_ENV}: slot 'route_lin' needs 'lowm_se' to replace " + "'route_shared' at one request" + ) return reasons @@ -345,11 +373,16 @@ def build(model, ctx): hop=bound["boundary"].hop is not None, gdn_push=bool(bound["gdn"].packed), gdn_front=bool(bound["gdn"].front), + gdn_trt_front=bound["gdn"].gdn_trt_front, route_shared=bound["moe"].shared, skv2=bound["attention"].skv2, mka=bound["attention"].mka, mka_wide=bound["attention"].mka_wide, mka_push=bound["attention"].mka_push, + route_lin=bool(bound["gdn"].rlin) + and bound["attention"].route_lin + and bound["boundary"].hop is not None + and bound["boundary"].hop.rl is not None, lowm_in=bound["lowm"].lowm_in, lowm_out=bound["lowm"].lowm_out, lowm_se=bound["lowm"].se, diff --git a/sgl_nvidia_mega/qwen3_5/runtime/executor.py b/sgl_nvidia_mega/qwen3_5/runtime/executor.py index 24d7973c..e848f5ee 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/executor.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/executor.py @@ -208,6 +208,11 @@ def _check_pools(self, view): "NVIDIA megakernel: slot gdn_front cannot run: " + "; ".join(reasons) ) + front.prepare( + self.lm, + view.mamba, + trt=any(s.gdn_trt_front for s in self.slots.values()), + ) skv2 = any(s.skv2 for s in self.slots.values()) mka = any(s.mka for s in self.slots.values()) if skv2 or mka: diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/mka.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/mka.py index 81e514d9..a272b81a 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/attention/mka.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/mka.py @@ -39,6 +39,7 @@ "q35mka_prologue_late": 0, "q35mka_gate": 0, "q35mka_epilogue_q4": 0, + "q35mka_rl_router": 0, "q35mka_epilogue_q4_late": 0, "q35mka_epilogue_q8": 0, "q35mka_epilogue_q8_late": 0, @@ -283,3 +284,35 @@ def epilogue( "" if early_trigger else "_late" ) _launch(name, 16 * CLUSTER, p, attn.device) + + +def epilogue_rl(attn, gate, w_packed, sx, sw, out, site: Site, rl_args) -> None: + """``epilogue`` of a TP8 B1 step (4 q heads a rank, early trigger, pushing into the hop) plus the router partial + (slot ``route_lin``): 16 CTAs push this rank's router partial into the hop's router mailbox + (``rl_args``: sm103/hop RlArgs, from ``Hop.rl_push_args``).""" + m = int(attn.shape[0]) + if not (attn.is_contiguous() and gate.is_contiguous() and out.is_contiguous()): + raise ValueError("MKA epilogue: attn, gate and out must be contiguous") + if m > MAX_ROWS or attn.shape[1] != 4 * HEAD_DIM or m > site.m_cap: + raise ValueError(f"MKA router-lin epilogue: attn {tuple(attn.shape)}") + p = EpiParams( + attn=attn.data_ptr(), + gate=gate.data_ptr(), + w=w_packed.data_ptr(), + sx=sx.data_ptr(), + sw=sw.data_ptr(), + out=out.data_ptr(), + m=m, + push=1, + site=site, + ) + stream = torch.cuda.current_stream(attn.device).cuda_stream + kernels = _kernels(attn.device.index) + jit.launch( + kernels["q35mka_epilogue_q4"], + (16 * CLUSTER, 1, 1), + [p], + stream, + (CLUSTER, 1, 1), + ) + jit.launch(kernels["q35mka_rl_router"], (16, 1, 1), [p, rl_args], stream) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/skv2.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/skv2.py index 5e2c0542..100329cf 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/attention/skv2.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/skv2.py @@ -303,6 +303,7 @@ def kernel( copy_s = tcgen05.make_tmem_copy(copy_atom, s_plane).get_slice(local_tid) rlayout = copy_s.partition_D(s_plane).layout # Stock split rule: pages_per_split = cdiv(cdiv(length, 64), SPLITS); split covers whole pages. + # lengths / table are read before the dependency wait: they must be written before the graph launch. length = lengths[0] pps = cute.ceil_div(cute.ceil_div(length, 64), self.splits) begin_tok = split * pps * 64 diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu b/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu index 95eb922a..11530ff3 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu @@ -244,6 +244,7 @@ __device__ __forceinline__ void mka_prologue_kernel(const ProParams p) { if (!needs_norm && !is_gate) kvs = __ldg(p.v_scale); pdl_wait(); + // Keep this release after the wait: q35mka_gate reads x (the hop's output) before its own wait. if (kEarlyTrigger) pdl_trigger(); const float inv_sx = div_full(1.0f, sx); const float alpha = __fmul_rn(sx, sw); @@ -834,6 +835,168 @@ __device__ __forceinline__ void mka_epilogue_wide_kernel(const EpiParams p) { extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_prologue(const q35mka::ProParams p) { q35mka::mka_prologue_kernel(p); } +namespace q35mka { +// ---------------------------------------------------------------- slot route_lin: router before the all-reduce +// The attention layers' form of sm103/hop/src/rlin.cu (numeric): after the o_proj epilogue of a TP8 B1 +// step, 16 CTAs compute this rank's router partial +// (A W_o,r * sx * sw on the gated FP8 o_proj input, plus A[:, K_r] on the layer's residual input) and push it, fp32, +// into the hop's router mailbox; the MoE input's consume_rl sums the partials into the logits. The epilogue CTAs are +// unchanged. +struct RlArgs { // = q35rl::RlArgs (sm103/hop/src/rlin.cu) + void* rmc; + void* ruc; + const uint4* rw; + const uint16_t* h0; + int k0; + uint16_t* logits; + uint8_t* x_fp4; + uint8_t* x_sf; + const float* gs; + unsigned long long* trace; +}; +constexpr int kRlExperts = 512, kRlCoreK = 1024, kRlSliceK = 512, kRlSteps = 12, kRlCtas = 16; +constexpr uint32_t kRlSent = 0x80000000u; +struct RlSmem { + uint32_t xb[8][(kRlCoreK + kRlSliceK) / 2 + 4]; +}; +__device__ __forceinline__ void rl_mma_bf16(float (&d)[4], const uint4& a, uint32_t b0, uint32_t b1) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b0), "r"(b1)); +} +__device__ __forceinline__ uint32_t rl_e4m3x2_bf16x2(uint16_t v) { + float lo, hi; + asm("{\n\t.reg .b32 h2;\n\t.reg .b16 l, h;\n\t" + "cvt.rn.f16x2.e4m3x2 h2, %2;\n\t" + "mov.b32 {l, h}, h2;\n\t" + "cvt.f32.f16 %0, l;\n\t" + "cvt.f32.f16 %1, h;\n\t}" + : "=f"(lo), "=f"(hi) + : "h"(v)); + __nv_bfloat162 b = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&b); +} +__device__ __forceinline__ void rl_mc_st_v4f(float* mc, float a, float b, float c, float d) { + asm volatile("multimem.st.relaxed.sys.global.v4.f32 [%0], {%1,%2,%3,%4};" ::"l"(mc), "f"(a), "f"(b), "f"(c), "f"(d) + : "memory"); +} + +// Router partial rows [32c, 32c + 32); K is the gated FP8 o_proj input (1024), +// then h0[K_r] (512). +__device__ __forceinline__ uint2 rl_ld_cg2(const void* p) { + uint2 v; + asm volatile("ld.global.cg.v2.u32 {%0,%1}, [%2];" : "=r"(v.x), "=r"(v.y) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ uint4 rl_ld_cg4(const void* p) { + uint4 v; + asm volatile("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ void mka_epilogue_rl_router(const EpiParams& p, const RlArgs& r, float (&red)[8][32][8], + RlSmem& sm) { + const int c = blockIdx.x; + const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + const int g = lane >> 2, t = lane & 3; + uint4 a[kRlSteps][2]; + { + const uint4* wb = r.rw + ((size_t)c * 8 + warp) * (kRlSteps * 2 * 32) + lane; +#pragma unroll + for (int st = 0; st < kRlSteps; ++st) +#pragma unroll + for (int mt = 0; mt < 2; ++mt) a[st][mt] = ldg_nc_v4(wb + (st * 2 + mt) * 32); + } + const float sx = __ldg(p.sx); + pdl_trigger(); + const float inv_sx = div_full(1.0f, sx); + const int m = p.m; + // stage: 8 rows x 256 quads of the gated input (attn * sigmoid(gate) -> static e4m3, the epilogue's rounding), then + // 8 rows x 64 chunks of 8 bf16 of h0's slice; every load goes out before any is used + uint2 ao[8], ag[8]; + uint4 hv[2]; +#pragma unroll + for (int u = 0; u < 8; ++u) { + const int q = tid + u * 256, row = q >> 8, col = (q & 255) * 4; + ao[u] = make_uint2(0u, 0u); + ag[u] = make_uint2(0u, 0u); + if (row < m) { + ao[u] = rl_ld_cg2(p.attn + (size_t)row * kRlCoreK + col); + ag[u] = rl_ld_cg2(p.gate + (size_t)row * kRlCoreK + col); + } + } +#pragma unroll + for (int u = 0; u < 2; ++u) { + const int q = tid + u * 256, row = q >> 6, chk = q & 63; + hv[u] = row < m ? rl_ld_cg4(r.h0 + (size_t)row * 4096 + r.k0 + chk * 8) : make_uint4(0u, 0u, 0u, 0u); + } +#pragma unroll + for (int u = 0; u < 8; ++u) { + const int q = tid + u * 256, row = q >> 8, col = (q & 255) * 4; + const uint32_t f8 = row < m ? gate_quant4_r(ao[u], ag[u], inv_sx) : 0u; + sm.xb[row][col / 2] = rl_e4m3x2_bf16x2((uint16_t)(f8 & 0xFFFFu)); + sm.xb[row][col / 2 + 1] = rl_e4m3x2_bf16x2((uint16_t)(f8 >> 16)); + } +#pragma unroll + for (int u = 0; u < 2; ++u) { + const int q = tid + u * 256, row = q >> 6, chk = q & 63; + sm.xb[row][kRlCoreK / 2 + chk * 4 + 0] = hv[u].x; + sm.xb[row][kRlCoreK / 2 + chk * 4 + 1] = hv[u].y; + sm.xb[row][kRlCoreK / 2 + chk * 4 + 2] = hv[u].z; + sm.xb[row][kRlCoreK / 2 + chk * 4 + 3] = hv[u].w; + } + __syncthreads(); + float ch[2][4][4]; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int k = 0; k < 4; ++k) +#pragma unroll + for (int e = 0; e < 4; ++e) ch[mt][k][e] = 0.f; +#pragma unroll + for (int st = 0; st < kRlSteps; ++st) { + const int kk = (warp * kRlSteps + st) * 16; + const uint32_t b0 = sm.xb[g][kk / 2 + t]; + const uint32_t b1 = sm.xb[g][kk / 2 + 4 + t]; + rl_mma_bf16(ch[0][st & 3], a[st][0], b0, b1); + rl_mma_bf16(ch[1][st & 3], a[st][1], b0, b1); + } + { + float4* dst = reinterpret_cast(&red[warp][lane][0]); +#pragma unroll + for (int mt = 0; mt < 2; ++mt) { + float s4[4]; +#pragma unroll + for (int e = 0; e < 4; ++e) + s4[e] = __fadd_rn(__fadd_rn(ch[mt][0][e], ch[mt][1][e]), __fadd_rn(ch[mt][2][e], ch[mt][3][e])); + dst[mt] = make_float4(s4[0], s4[1], s4[2], s4[3]); + } + } + __syncthreads(); + const int tok = tid >> 2, j = tid & 3; + if (tok < m) { + float v[8]; +#pragma unroll + for (int q = 0; q < 8; ++q) { + const int rr = 8 * j + q; + const int l = (rr % 8) * 4 + tok / 2; + const int e = (rr / 16) * 4 + ((rr % 16) / 8) * 2 + (tok % 2); + float x = red[0][l][e]; +#pragma unroll + for (int w = 1; w < 8; ++w) x = __fadd_rn(x, red[w][l][e]); + v[q] = __float_as_uint(x) == kRlSent ? 0.f : x; + } + const q35hop::Site& s = p.site; + const uint32_t stage = q35hop::stage_peek(s); + float* dst = static_cast(r.rmc) + + (((size_t)stage * s.world + s.rank) * s.m_cap + tok) * kRlExperts + c * 32 + 8 * j; + rl_mc_st_v4f(dst, v[0], v[1], v[2], v[3]); + rl_mc_st_v4f(dst + 4, v[4], v[5], v[6], v[7]); + } +} +} // namespace q35mka + extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_prologue_late(const q35mka::ProParams p) { q35mka::mka_prologue_kernel(p); } @@ -843,6 +1006,13 @@ extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_gate(co extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_epilogue_q4(const q35mka::EpiParams p) { q35mka::mka_epilogue_kernel<4, true>(p); } +// The router partial runs immediately after q35mka_epilogue_q4 (16 CTAs, no cluster). +extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_rl_router(const q35mka::EpiParams p, + const q35mka::RlArgs r) { + __shared__ __align__(16) float red[8][32][8]; + __shared__ __align__(16) q35mka::RlSmem sm; + q35mka::mka_epilogue_rl_router(p, r, red, sm); +} extern "C" __global__ void __launch_bounds__(q35mka::kThreads, 1) q35mka_epilogue_q4_late(const q35mka::EpiParams p) { q35mka::mka_epilogue_kernel<4, false>(p); } diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/q1_kernel.py b/sgl_nvidia_mega/qwen3_5/sm103/draft/q1_kernel.py index 47de7417..1e11294c 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/q1_kernel.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/q1_kernel.py @@ -304,6 +304,7 @@ def kernel( copy_s = tcgen05.make_tmem_copy(copy_atom, s_plane).get_slice(local_tid) rlayout = copy_s.partition_D(s_plane).layout # trtllm-gen split: one CTA per SM, ceil(steps / CTAs) contiguous 256-token steps each. + # lengths / table are read before the dependency wait: they must be written before the graph launch. length = lengths[0] spc = cute.ceil_div(cute.ceil_div(length, 256), self.nsplit) begin_tok = split * spc * 256 diff --git a/sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py b/sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py index b4cbeccc..7b74ea03 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py @@ -9,6 +9,10 @@ It launches as a programmatic dependent: the checkpoint rows, ring cursors and static parameters load before ``griddepcontrol.wait``, every activation (q / k / v, z, a, b) after it. With ``early_dependents`` the next launch may start at once (it must wait before reading the output). +The testable ``early_hist`` variant can also hold committed history ahead of the wait; +production keeps ordinary history loads. Flat stores compute the ring entries and FP8 +outputs first and predicate only their stores, then store update vectors after the +output; the arithmetic is unchanged. """ from __future__ import annotations @@ -24,6 +28,11 @@ SOURCE = Path(__file__).parent / "src" / "core.cu" THREADS, K, V, RING, MAX_ROWS_PER_REQUEST = 128, 128, 128, 16, 8 +# Validated production policy: ordinary history loads, flat stores, and the early +# dependent release only over the measured 35-56-row window. +EARLY = False +RELEASE_MIN_ROWS, RELEASE_MAX_ROWS = 35, 56 +FLAT = True class Params(ctypes.Structure): @@ -81,6 +90,8 @@ def _kernel( prof: bool = False, warps: int = 4, min_blocks: int = 1, + early_hist: bool = False, + flat: bool = False, ) -> jit.Kernel: major, minor = torch.cuda.get_device_capability(device_index) defines = dict( @@ -94,6 +105,10 @@ def _kernel( NW=warps, MINB=min_blocks, ) + if early_hist: + defines["EARLY_HIST"] = 1 + if flat: + defines["FLAT_STORES"] = 1 source = "".join(f"#define {k} {v}\n" for k, v in defines.items()) cubin = jit.compile_source( source + SOURCE.read_text(), @@ -118,6 +133,19 @@ def build(device_index: int, rows: int, hk: int, hv: int, norm_f32: bool) -> Non trigger, False, min_blocks=min_blocks, + flat=FLAT, + ) + if EARLY: + _kernel( + device_index, + rows, + hk, + hv, + norm_f32, + trigger, + False, + early_hist=True, + flat=FLAT, ) @@ -166,6 +194,8 @@ def core( _profile: Optional[torch.Tensor] = None, warps: int = 4, min_blocks: Optional[int] = None, + early_hist: Optional[bool] = None, + flat: Optional[bool] = None, ) -> torch.Tensor: """Write the gated RMSNorm of the verify output as FP8 into ``out`` [rows, HV * V]. @@ -174,7 +204,9 @@ def core( row-strided views with one row stride; ``checkpoint`` [slots, HV, V, K] BF16; the ring [slots, HV | H, 16, V | K] (low parts optional) with cursors by ``replay_indices``; the first ``state_indices.numel()`` requests of ``rows_per_request`` rows each. ``o_out`` [rows, HV, V] - BF16 also receives the replay output (checks only).""" + BF16 also receives the replay output (checks only). ``early_hist`` (default ``EARLY``) loads the + committed history before the PDL wait when the grid runs one CTA an SM (same bits); ``flat`` + (default ``FLAT``) predicates the stores (same bits).""" rows, T = qkv.shape[0], rows_per_request n = rows // T HV = checkpoint.shape[1] @@ -262,6 +294,10 @@ def core( o_out=0 if o_out is None else o_out.data_ptr(), prof=0 if _profile is None else _profile.data_ptr(), ) + min_blocks = ( + _min_blocks(qkv.device.index, n * HV) if min_blocks is None else min_blocks + ) + early_hist = (EARLY if early_hist is None else early_hist) and min_blocks == 1 kernel = _kernel( qkv.device.index, T, @@ -272,7 +308,9 @@ def core( o_out is not None, _profile is not None, warps, - _min_blocks(qkv.device.index, n * HV) if min_blocks is None else min_blocks, + min_blocks, + early_hist, + FLAT if flat is None else flat, ) jit.launch( kernel, (n * HV, 1, 1), [p], torch.cuda.current_stream(qkv.device).cuda_stream diff --git a/sgl_nvidia_mega/qwen3_5/sm103/gdn/front.py b/sgl_nvidia_mega/qwen3_5/sm103/gdn/front.py index 5b04a108..d32face8 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/gdn/front.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/gdn/front.py @@ -10,9 +10,13 @@ It launches as a programmatic dependent and releases its own dependents at entry: the next launch must wait (``griddepcontrol.wait``) before it reads any output, b / a -included. By default (``tma``) each projection CTA's packed weights arrive by one bulk +included. The fallback (``tma``) has each projection CTA's packed weights arrive by one bulk copy into shared memory issued before that wait, instead of register loads; same MMAs, -same bits. +same bits. The validated one-request path (``pin``) takes them by +register loads held ahead of that wait (``q35_gdn_front_pin``; without the hold ptxas +sinks the loads below it), no shared-memory reload after it; same bits. ``prefetch`` +False leaves out the L2 prefetch of the verify's checkpoint and ring (when the core loads +them before its own wait, the prefetch only delays this launch's completion). """ from __future__ import annotations @@ -32,6 +36,8 @@ 8, 88, ) # 10 projection clusters + 1 side-projection cluster +# a request's launch with its register weight loads held ahead of the PDL wait (q35_gdn_front_pin) +PIN = True TIMEOUT = ( 500_000_000 # clock cycles a CTA waits for its cluster's partials before trapping ) @@ -113,7 +119,7 @@ def pack(weight_nk: torch.Tensor) -> torch.Tensor: @cache -def _kernel(device_index: int, tma: bool = True) -> jit.Kernel: +def _kernel(device_index: int, tma: bool = True, pin: bool = False) -> jit.Kernel: major, minor = torch.cuda.get_device_capability(device_index) cubin = jit.compile_source( SOURCE.read_text(), @@ -121,6 +127,8 @@ def _kernel(device_index: int, tma: bool = True) -> jit.Kernel: f"sm_{major}{minor}a", ("-std=c++17", "-default-device"), ) + if pin: + return jit.load(cubin, "q35_gdn_front_pin", (THREADS, 1, 1), 0) if tma: return jit.load(cubin, "q35_gdn_front_tma", (THREADS, 1, 1), WEIGHT_BYTES) return jit.load(cubin, "q35_gdn_front", (THREADS, 1, 1), 0) @@ -129,6 +137,8 @@ def _kernel(device_index: int, tma: bool = True) -> jit.Kernel: def build(device_index: int) -> None: """Compile the front (no compile inside a step).""" _kernel(device_index, True) + if PIN: + _kernel(device_index, True, True) def _prefetch(t: torch.Tensor, ring: int): @@ -169,6 +179,8 @@ def front( ring_indices: torch.Tensor, pad_slot: int = -1, tma: bool = True, + pin: bool | None = None, + prefetch: bool = True, ) -> None: """The front of the BF16 rows ``x`` [7, 4096] (contiguous) into the caller's buffers: ``projected`` [7, 2560] BF16 (q / k / v / z), ``conv_out`` [7, 1536] BF16, @@ -179,7 +191,9 @@ def front( capture, that a timed-out wait sets before trapping. The ReplaySSM verify's ``checkpoint`` (by state index) and ``history_*`` rings (by ``ring_indices[0]``) are prefetched into L2: a hint, nothing is written. ``tma`` False takes the weights by - register loads instead of the bulk copy (same values).""" + register loads instead of the bulk copy (same values); ``pin`` (default ``PIN``) + holds those register loads ahead of the PDL wait (same values); ``prefetch`` False + skips the L2 prefetch.""" if ( x.dtype != torch.bfloat16 or tuple(x.shape) != (ROWS, HIDDEN) @@ -198,7 +212,7 @@ def front( _prefetch(history_k, 1), _prefetch(history_g, 1), ) - if r is not None + if r is not None and prefetch ] ba = BaParams( input=x.data_ptr(), @@ -237,8 +251,9 @@ def front( ) p.state_stride[:] = list(conv_state.stride()) p.window_stride[:] = [windows.stride(0), 9] + pin = PIN if pin is None else pin jit.launch( - _kernel(x.device.index, tma), + _kernel(x.device.index, tma, pin), (CTAS, 1, 1), [p], torch.cuda.current_stream(x.device).cuda_stream, diff --git a/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/core.cu b/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/core.cu index c4d9f48f..f538b643 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/core.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/core.cu @@ -15,10 +15,14 @@ // // PDL: before griddepcontrol.wait the kernel reads only data of earlier steps or static data (indices, // ring cursors, the checkpoint rows, A_log / dt_bias / norm weight / FP8 scale); every global store -// and every read of the step's activations (q / k / v, z, a, b) follows the wait. +// and every read of the step's activations (q / k / v, z, a, b) follows the wait. The commit that writes the +// ring cursors and checkpoint rows must stay out of this launch's PDL window (it runs in the previous round). // // Defines: TT (rows a request), HK / HV (key / value heads), NORM_F32 (FP32 norm weight), TRIGGER (2: release dependents at entry, 0: at exit), -// STORE_O (also store the BF16 replay output, checks only). +// STORE_O (also store the BF16 replay output, checks only), EARLY_HIST (1: the committed history -- keys, log-decays, +// update vectors and their decays -- before the wait, held there; the same arithmetic, the same bits), FLAT_STORES +// (1: the ring entries and FP8 outputs computed first and only their stores predicated, one low-part test, the update +// vectors stored after the output; the same bits). typedef unsigned char u8; typedef unsigned short u16; @@ -30,6 +34,12 @@ typedef unsigned int u32; #ifndef MINB #define MINB 1 // resident CTAs an SM the registers are capped for #endif +#ifndef EARLY_HIST +#define EARLY_HIST 0 +#endif +#ifndef FLAT_STORES +#define FLAT_STORES 0 +#endif #ifndef NW #define NW 4 // warps; each owns VD / NW value columns as MT m16 tiles #endif @@ -130,6 +140,14 @@ __device__ __forceinline__ void kred(float (&acc)[4], const U4 (&A)[4], const U4 __device__ __forceinline__ void pdl_wait() { asm volatile("griddepcontrol.wait;" ::: "memory"); } __device__ __forceinline__ void pdl_launch() { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); } +// Holds the loads that `dep` reads (one word of each) ahead of griddepcontrol.wait, which ptxas may otherwise sink +// them below; `never` is a launch constant that is never 0. +__device__ __forceinline__ void pin_before_wait(u32 dep, unsigned long long never) { + asm volatile("{ .reg .pred a, b; setp.eq.u64 a, %1, 0; setp.eq.u32 b, %0, 1; and.pred a, a, b; @a trap; }" + :: "r"(dep), "l"(never) : "memory"); +} + + __device__ __forceinline__ long long load_index(const void* p, int i, bool i64) { return i64 ? reinterpret_cast(p)[i] : (long long)reinterpret_cast(p)[i]; } @@ -138,7 +156,17 @@ __device__ __forceinline__ long long load_index(const void* p, int i, bool i64) #define PROF 0 #endif #if PROF -#define STAMP(i, dep) do { if (lane == 0) { asm volatile("" ::"f"((float)(dep))); st_[i] = clock64(); } } while (0) +// PROF_GT (micro builds): %globaltimer (ns, comparable across SMs) instead of clock64 +__device__ __forceinline__ long long prof_time() { +#ifdef PROF_GT + unsigned long long t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return (long long)t; +#else + return clock64(); +#endif +} +#define STAMP(i, dep) do { if (lane == 0) { asm volatile("" ::"f"((float)(dep))); st_[i] = prof_time(); } } while (0) #else #define STAMP(i, dep) do { } while (0) #endif @@ -167,7 +195,7 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P #if PROF long long st_[8]; - st_[0] = clock64(); + st_[0] = prof_time(); #endif #if TRIGGER == 2 pdl_launch(); @@ -219,6 +247,77 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P #endif } const float rscale = rcp_rn(p.scale[0]); +#if EARLY_HIST + // the committed history (ring entries j < wp, written by earlier steps; this step writes j >= wp): the history keys + // of warps 2 / 3, and every lane's update-vector A fragments with their decays (the post-sync block's arithmetic) + U4 KHe[2][4]; + u32 hist_a[MT][4]; + float hist_td = 0.f; + { + u32 dep = 0; +#pragma unroll + for (int jt = 0; jt < 2; ++jt) +#pragma unroll + for (int c = 0; c < 4; ++c) KHe[jt][c] = zero4(); + if ((warp == 2 || warp == 3) && wp > 0) { +#pragma unroll + for (int jt = 0; jt < 2; ++jt) { + const int j = 8 * jt + r; + const u16* kh = p.ring_k + rid * p.k_slot + ((long long)hk * RING + ((base + j) & (RING - 1))) * KD; +#pragma unroll + for (int c = 0; c < 4; ++c) KHe[jt][c] = j < wp ? ld16(kh + 32 * c + 8 * q) : zero4(); + } + } +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < 4; ++i) hist_a[m][i] = 0u; + if (wp > 0) { + float gj = 0.f; + if (lane < RING && lane < wp) gj = p.ring_g[rid * p.g_slot + hv * RING + ((base + lane) & (RING - 1))]; +#pragma unroll + for (int d = 1; d < 32; d <<= 1) { + const float o = __shfl_up_sync(FULL, gj, d); + if (lane >= d) gj += o; + } + const float gtot = __shfl_sync(FULL, gj, 31); + const int jj[4] = {2 * q, 2 * q + 1, 2 * q + 8, 2 * q + 9}; + float dec[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float pj = __shfl_sync(FULL, gj, jj[i]); + dec[i] = jj[i] < wp ? __expf(gtot - pj) : 0.f; + } + hist_td = __expf(gtot); + const u16* dp = p.ring_d + rid * p.d_slot + (long long)hv * RING * VD; +#pragma unroll + for (int m = 0; m < MT; ++m) { + float e[2][4]; // [half][j index i] +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int i = 0; i < 4; ++i) + e[h][i] = jj[i] < wp + ? bf2f(dp[((base + jj[i]) & (RING - 1)) * VD + vb + 16 * m + 8 * h + r]) * dec[i] + : 0.f; + hist_a[m][0] = pack2(e[0][0], e[0][1]); + hist_a[m][1] = pack2(e[1][0], e[1][1]); + hist_a[m][2] = pack2(e[0][2], e[0][3]); + hist_a[m][3] = pack2(e[1][2], e[1][3]); + dep ^= hist_a[m][0] ^ hist_a[m][3]; + } + } +#pragma unroll + for (int jt = 0; jt < 2; ++jt) +#pragma unroll + for (int c = 0; c < 4; ++c) dep ^= KHe[jt][c].x; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int h = 0; h < 2; ++h) dep ^= H[m][h][0].x; + pin_before_wait(dep, (unsigned long long)p.write_pos); + } +#endif STAMP(1, bflo(H[1][1][3].w) + bflo(H[0][0][0].x) + rscale); // ---------------- the wait: the step's activations from here on ---------------- @@ -339,6 +438,28 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P const long long off = ((long long)hk * RING + ((base + wp + r) & (RING - 1))) * KD; u16* kp = p.ring_k + rid * p.k_slot + off; u16* klp = p.ring_klo ? p.ring_klo + rid * p.klo_slot + off : nullptr; +#if FLAT_STORES + U4 vh[4], vl[4]; +#pragma unroll + for (int c = 0; c < 4; ++c) { + u32 kh[4], kl[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const u32 wk = word(KR[c], i); + const float k0 = bflo(wk) * kmul, k1 = bfhi(wk) * kmul; + kh[i] = pack2(k0, k1); + kl[i] = pack2(k0 - bflo(kh[i]), k1 - bfhi(kh[i])); + } + vh[c].x = kh[0]; vh[c].y = kh[1]; vh[c].z = kh[2]; vh[c].w = kh[3]; + vl[c].x = kl[0]; vl[c].y = kl[1]; vl[c].z = kl[2]; vl[c].w = kl[3]; + } +#pragma unroll + for (int c = 0; c < 4; ++c) st16(kp + 32 * c + 8 * q, vh[c]); + if (klp) { +#pragma unroll + for (int c = 0; c < 4; ++c) st16(klp + 32 * c + 8 * q, vl[c]); + } +#else #pragma unroll for (int c = 0; c < 4; ++c) { u32 kh[4], kl[4]; @@ -355,6 +476,7 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P st16(kp + 32 * c + 8 * q, vh); if (klp) st16(klp + 32 * c + 8 * q, vl); } +#endif } // head-common products: warp 0 k k^T (A), warp 1 k q^T (F), warps 2 / 3 the history scores @@ -374,10 +496,15 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P U4 KHh[4]; #pragma unroll for (int jt = 0; jt < 2; ++jt) { +#if EARLY_HIST +#pragma unroll + for (int c = 0; c < 4; ++c) KHh[c] = KHe[jt][c]; +#else const int j = 8 * jt + r; const u16* kh = p.ring_k + rid * p.k_slot + ((long long)hk * RING + ((base + j) & (RING - 1))) * KD; #pragma unroll for (int c = 0; c < 4; ++c) KHh[c] = j < wp ? ld16(kh + 32 * c + 8 * q) : zero4(); +#endif float acc[4] = {0.f, 0.f, 0.f, 0.f}; if (warp == 2) { kred(acc, QN, KHh); @@ -400,6 +527,20 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P STAMP(5, 0.f); if (wp > 0) { +#if EARLY_HIST + const u32 bq0 = s_sc[0][0][lane], bq1 = s_sc[0][1][lane]; + const u32 bk0 = s_sc[1][0][lane], bk1 = s_sc[1][1][lane]; +#pragma unroll + for (int m = 0; m < MT; ++m) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + HQ[m][i] *= hist_td; + HKs[m][i] *= hist_td; + } + mma16816(HQ[m], hist_a[m][0], hist_a[m][1], hist_a[m][2], hist_a[m][3], bq0, bq1); + mma16816(HKs[m], hist_a[m][0], hist_a[m][1], hist_a[m][2], hist_a[m][3], bk0, bk1); + } +#else // committed history: hw = e^{g_total} hw + bf16(d_j e^{g_total - G_j}) . bf16(scores) float gj = 0.f; if (lane < RING && lane < wp) gj = p.ring_g[rid * p.g_slot + hv * RING + ((base + lane) & (RING - 1))]; @@ -440,6 +581,7 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P mma16816(HQ[m], a0, a1, a2, a3, bq0, bq1); mma16816(HKs[m], a0, a1, a2, a3, bk0, bk1); } +#endif } // rhs = beta (v - e^G hw_k); forward substitution (I + A) U = rhs per value column, every lane of a quad @@ -484,6 +626,7 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P } } STAMP(6, U[2 * MT - 1][TT - 1] + U[0][TT - 1]); +#if !FLAT_STORES // this step's update vectors (high, low parts) into the ring { u16* dp = p.ring_d + rid * p.d_slot + (long long)hv * RING * VD; @@ -507,6 +650,7 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P } } } +#endif // output O[v, t] = e^{G_t} hw_q + sum_{s <= t} U[v, s] F[s][t] float O[MT][4]; { @@ -530,6 +674,36 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P O[m][2 * h + 1] = eG1 * HQ[m][2 * h + 1] + a1; } } +#if FLAT_STORES + // this step's update vectors (high, low parts) into the ring, off the output's path + { + u16* dp = p.ring_d + rid * p.d_slot + (long long)hv * RING * VD; + const bool w0 = ok0 && wp + t0 < RING, w1 = ok1 && wp + t1 < RING; + const int ph0 = (base + wp + t0) & (RING - 1), ph1 = (base + wp + t1) & (RING - 1); + if (p.ring_dlo) { + u16* dlp = p.ring_dlo + rid * p.dlo_slot + (long long)hv * RING * VD; +#pragma unroll + for (int col = 0; col < 2 * MT; ++col) { + const int v = vb + 16 * (col >> 1) + 8 * (col & 1) + r; + const float d0 = pick8(U[col], t0), d1 = pick8(U[col], t1); + const u16 h0 = f2bf(d0), h1 = f2bf(d1); + const u16 l0 = f2bf(d0 - bf2f(h0)), l1 = f2bf(d1 - bf2f(h1)); + if (w0) dp[ph0 * VD + v] = h0; + if (w0) dlp[ph0 * VD + v] = l0; + if (w1) dp[ph1 * VD + v] = h1; + if (w1) dlp[ph1 * VD + v] = l1; + } + } else { +#pragma unroll + for (int col = 0; col < 2 * MT; ++col) { + const int v = vb + 16 * (col >> 1) + 8 * (col & 1) + r; + const u16 h0 = f2bf(pick8(U[col], t0)), h1 = f2bf(pick8(U[col], t1)); + if (w0) dp[ph0 * VD + v] = h0; + if (w1) dp[ph1 * VD + v] = h1; + } + } + } +#endif // gated RMSNorm over the head's 128 columns of the BF16-rounded replay output, then static FP8 float xb[MT][4]; @@ -572,6 +746,31 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P } const float rstd0 = rsqrtf(ss0 / (float)VD + p.eps); const float rstd1 = rsqrtf(ss1 / (float)VD + p.eps); +#if FLAT_STORES + { + u8* y0 = p.y + (row0 + t0) * p.y_stride + hv * VD + vb + r; + u8* y1 = p.y + (row0 + t1) * p.y_stride + hv * VD + vb + r; + u8 yq[MT][4]; +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < 4; ++i) { + float yv = xb[m][i] * ((i & 1) ? rstd1 : rstd0) * nw[m][i >> 1]; + yv *= zg[m][i]; + yq[m][i] = f2e4m3(bfround(yv) * rscale); + } +#pragma unroll + for (int m = 0; m < MT; ++m) +#pragma unroll + for (int i = 0; i < 4; ++i) { + if (i & 1) { + if (ok1) y1[16 * m + 8 * (i >> 1)] = yq[m][i]; + } else { + if (ok0) y0[16 * m + 8 * (i >> 1)] = yq[m][i]; + } + } + } +#else #pragma unroll for (int m = 0; m < MT; ++m) #pragma unroll @@ -586,9 +785,10 @@ extern "C" __global__ void __launch_bounds__(NW * 32, MINB) q35_gdn_core(const P p.y[(row0 + t) * p.y_stride + hv * VD + v] = f2e4m3(bfround(yv) * rscale); } } +#endif #if PROF if (lane == 0) { - const long long end = clock64(); + const long long end = prof_time(); long long* pw = p.prof + ((long long)cta * NW + warp) * 9; #pragma unroll for (int i = 0; i < 8; ++i) pw[i] = st_[i]; diff --git a/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/front.cu b/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/front.cu index fb7abde4..267401a9 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/front.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/front.cu @@ -10,7 +10,9 @@ // // Built with NVRTC: front_kernel is a device function, and the extern "C" entries at the end // instantiate its BF16-input form with the fused conv and the inline side projection, with the -// projection weights in registers (q35_gdn_front) or by bulk copy into shared memory (q35_gdn_front_tma). +// projection weights in registers (q35_gdn_front), by bulk copy into shared memory (q35_gdn_front_tma), +// or in registers held ahead of the PDL wait (q35_gdn_front_pin: ptxas otherwise sinks q35_gdn_front's +// non-coherent weight loads below griddepcontrol.wait). #include @@ -195,6 +197,21 @@ __device__ __forceinline__ uint4 ldg_nc_v4(const uint4* p) { return r; } +// Holds every weight load named in `dep` (one word of each) ahead of griddepcontrol.wait: ptxas otherwise sinks the +// loads below the wait (ACQBULK). The thread waits for its weights here; `never` is a launch constant that is never 0. +__device__ __forceinline__ void pin_before_wait(uint32_t dep, uint64_t never) { + asm volatile("{ .reg .pred a, b; setp.eq.u64 a, %1, 0; setp.eq.u32 b, %0, 1; and.pred a, a, b; @a trap; }" + :: "r"(dep), "l"(never) : "memory"); +} +__device__ __forceinline__ uint32_t weight_dep(const uint4 (&w)[16][2]) { + uint32_t dep = 0; +#pragma unroll + for (int step = 0; step < 16; ++step) +#pragma unroll + for (int block = 0; block < 2; ++block) dep ^= w[step][block].x; + return dep; +} + __device__ __forceinline__ uint32_t smem_u32(const void* ptr) { return static_cast(__cvta_generic_to_shared(ptr)); } @@ -245,6 +262,15 @@ __device__ __forceinline__ void mma_e4m3(float (&d)[4], const uint4& a, uint32_t } // stock _static_quant_fp8: y * div.full(1, s), clamp +-448, cvt.rn.satfinite.e4m3 +// quant4_static of four BF16 already loaded +__device__ __forceinline__ uint32_t quant4_raw(const uint2 raw, float inv) { + const __nv_bfloat162 v01 = *reinterpret_cast(&raw.x); + const __nv_bfloat162 v23 = *reinterpret_cast(&raw.y); + float f[4] = {__low2float(v01), __high2float(v01), __low2float(v23), __high2float(v23)}; +#pragma unroll + for (int i = 0; i < 4; ++i) f[i] = fminf(fmaxf(__fmul_rn(f[i], inv), -448.f), 448.f); + return e4m3x2(f[0], f[1]) | (e4m3x2(f[2], f[3]) << 16); +} __device__ __forceinline__ uint32_t quant4_static(const __nv_bfloat16* p, float inv) { const uint2 raw = *reinterpret_cast(p); const __nv_bfloat162 v01 = *reinterpret_cast(&raw.x); @@ -311,7 +337,8 @@ constexpr uint32_t kWeightBytes = 8 * 16 * 2 * 32 * 16; // Each cluster reduces its split-K partials in fixed rank order. // Tma: the CTA's weights arrive by one bulk copy into shared memory issued before the PDL wait (instead of 32 // register loads per thread), then each thread reads the same A fragments; same MMAs, same bits. -template +// Early (register weights): every weight load is held ahead of the PDL wait by pin_before_wait. +template __device__ __forceinline__ void front_kernel(const Params& p) { if (threadIdx.x == 0) pdl_trigger(); __shared__ __align__(16) float inbox[kCluster][256]; @@ -368,8 +395,28 @@ __device__ __forceinline__ void front_kernel(const Params& p) { for (int block = 0; block < 2; ++block) weights[step][block] = ldg_nc_v4(wp + (step * 2 + block) * 32); } + static_assert(!(Tma && Early), "Early holds register weight loads"); const float sx = __ldg(p.sx), alpha = __fmul_rn(sx, __ldg(p.sw)); + float inverse = 0.f; + if constexpr (Early) { + // the static scales and 1 / sx ahead of the wait too (ptxas sinks the scale loads below it otherwise) + inverse = div_full(1.0f, sx); + pin_before_wait(weight_dep(weights) ^ __float_as_uint(inverse) ^ __float_as_uint(alpha), p.timeout); + } pdl_wait(); + // Early: this thread's x loads go out first; warp 0 would otherwise wait out the sticky word and the conv slots (each + // consumed right after its load) before it issues them, and the CTA's x barrier with it. + uint2 raw[2][2] = {}; + if constexpr (Early) { +#pragma unroll + for (int i = 0; i < 2; ++i) + if (g < p.m) { + const auto* x = static_cast(p.x) + (req * p.m + g) * 4096 + rank * 512 + + (2 * warp + i) * 32 + 4 * t; + raw[i][0] = *reinterpret_cast(x); + raw[i][1] = *reinterpret_cast(x + 16); + } + } // Issued now, tested before the first push: no store precedes the check. const unsigned sticky = tid == 0 ? *reinterpret_cast(p.error) : 0u; // Both indices are read now (independently) and used only after the projection. @@ -380,7 +427,7 @@ __device__ __forceinline__ void front_kernel(const Params& p) { conv_replay = p.replay_index64 ? static_cast(p.replay_indices)[0] : static_cast(p.replay_indices)[0]; } - const float inverse = div_full(1.0f, sx); + if constexpr (!Early) inverse = div_full(1.0f, sx); #pragma unroll for (int i = 0; i < 2; ++i) { const int step = 2 * warp + i; @@ -391,6 +438,8 @@ __device__ __forceinline__ void front_kernel(const Params& p) { const uint8_t* x = static_cast(p.x) + offset; value = make_uint2(*reinterpret_cast(x), *reinterpret_cast(x + 16)); + } else if constexpr (Early) { + value = make_uint2(quant4_raw(raw[i][0], inverse), quant4_raw(raw[i][1], inverse)); } else { const auto* x = static_cast(p.x) + offset; value = make_uint2(quant4_static(x, inverse), quant4_static(x + 16, inverse)); @@ -539,3 +588,9 @@ extern "C" __global__ __launch_bounds__(sglang::gdn_pdl_front_ba_mkm_detail::kTh const __grid_constant__ sglang::gdn_pdl_front_ba_mkm_detail::Params p) { sglang::gdn_pdl_front_ba_mkm_detail::front_kernel(p); } + +// q35_gdn_front with its register weight loads held ahead of the PDL wait (same bits). +extern "C" __global__ __launch_bounds__(sglang::gdn_pdl_front_ba_mkm_detail::kThreads, 1) void q35_gdn_front_pin( + const __grid_constant__ sglang::gdn_pdl_front_ba_mkm_detail::Params p) { + sglang::gdn_pdl_front_ba_mkm_detail::front_kernel(p); +} diff --git a/sgl_nvidia_mega/qwen3_5/sm103/gdn/trt_front.py b/sgl_nvidia_mega/qwen3_5/sm103/gdn/trt_front.py new file mode 100644 index 00000000..c5fc2f5d --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/gdn/trt_front.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Adapted from TensorRT-LLM r3 (tekit 7df45abe5c), tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +# (Qwen3NextGatedDeltaNet._lowm_in_proj and its fused-front caller). +"""The GatedDeltaNet verify front of one to four requests on TensorRT-LLM r3's +``lowm_gemm`` (``sm103/trt_lowm``): in one launch the static-FP8 ``in_proj_qkvz`` on the hop's E4M3 rows +(``fp8``, TensorRT-LLM's mixed8 tactic) or on the BF16 rows quantized in registers (``quant``, its mixed tactic, +where no E4M3 copy exists), the BF16 ``in_proj_ba`` on the BF16 rows, and the causal conv update of the q / k / v +channels as the epilogue. This wrapper admits the validated TP8 shapes. + +Outputs: ``out`` [7, QKVZ] BF16 holds the convolved q | k | v then the raw z (the front's ``conv_out`` and +``projected[:, CONV:]`` in one buffer), ``ba`` [7, 2 HV] BF16 b then a; the conv state is updated at the request's +state index (none for a pad index, -1) and the windows after every token at intermediate row 0 (SGLang's +``verify_intermediate_state_indices`` gives batch row r row r). The sums run in TensorRT-LLM's order: numeric, not +SGLang's GEMMs bit for bit; the conv is the arithmetic of Triton's ``causal_conv1d_update``. + +The launch is a programmatic dependent that reads only weights, scales and conv weights before its +``griddepcontrol.wait`` (the rows, the state index and the conv state after it) and releases its dependents once +those weight loads are issued; the next launch waits for it before reading any output. ``prepare`` compiles both +activation forms where the pools first exist (before graph capture); ``front`` refuses an unprepared one. +""" + +from __future__ import annotations + +import torch + +MAX_REQUESTS = 4 +ROWS, HIDDEN = 7, 4096 +# TensorRT-LLM's _TACTICS mode of each activation form (lowm_linear.default_tactic) +MODES = {"fp8": "mixed8", "quant": "mixed"} + +_SPECS: dict = {} # (activation form, qkvz rows, requests) -> the prepared trt_lowm.Spec + + +def takes(requests: int) -> bool: + """Whether ``requests`` is inside the front's validated request range.""" + return 1 <= requests <= MAX_REQUESTS + + +def spec( + act: str, qkvz: int, ba: int, conv_weight, conv_state, windows, requests: int = 1 +): + """The ``trt_lowm`` launch of the front of ``requests`` requests of 7 rows with the ``act`` activations (``fp8`` / + ``quant``) over a [qkvz, 4096] FP8 and a [ba, 4096] BF16 weight, on these conv tensors' layouts: TensorRT-LLM's + tactic of the shape at M = 7 requests (BF16 segment: its k_warps and ksplit, one tile a CTA), as its GDN mixer + passes it; request r's windows at row r.""" + from sgl_nvidia_mega.qwen3_5.sm103 import trt_lowm + from sgl_nvidia_mega.qwen3_5.sm103.trt_lowm.lowm_linear import ( + default_tactic, + lowm_conv_supported, + ) + + m = ROWS * requests + tactic = default_tactic(m, qkvz, HIDDEN, mode=MODES[act]) + if not lowm_conv_supported(m, tactic, conv_weight.shape[0], ROWS): + raise ValueError( + f"TensorRT-LLM's {MODES[act]} tactic {tactic} has no conv epilogue for {m} x {qkvz}" + ) + return trt_lowm.spec( + m=m, + k=HIDDEN, + n_fp8=qkvz, + n_bf16=ba, + fp8_act=act, + tactic=tactic, + tactic_bf16=(tactic[0], 1, tactic[2] if len(tactic) > 2 else 1), + split_out=qkvz, + out_dtype=torch.bfloat16, + conv=trt_lowm.conv_spec( + conv_weight, conv_state, windows, requests=requests, tokens=ROWS + ), + ) + + +def prepare(qkvz: int, ba: int, conv_weight, conv_state, windows) -> None: + """Compile both activation forms for 1..MAX_REQUESTS requests on these layouts (not inside a graph capture).""" + from sgl_nvidia_mega.qwen3_5.sm103 import trt_lowm + + for requests in range(1, MAX_REQUESTS + 1): + for act in MODES: + s = spec(act, qkvz, ba, conv_weight, conv_state, windows, requests) + trt_lowm.prepare(s) + _SPECS[act, qkvz, requests] = s + + +def front( + x: torch.Tensor, + x_fp8, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + ba_weight: torch.Tensor, + conv_weight: torch.Tensor, + conv_state: torch.Tensor, + windows: torch.Tensor, + state_indices: torch.Tensor, + out: torch.Tensor, + ba: torch.Tensor, +) -> None: + """The front of R requests' BF16 rows ``x`` [7 R, 4096] (and their E4M3 copy ``x_fp8``, quantized with + ``input_scale``, or None) over ``weight`` [QKVZ, 4096] E4M3 (row-major [N, K]) and ``ba_weight`` [2 HV, 4096] + BF16 into ``out`` [7 R, QKVZ] and ``ba`` [7 R, 2 HV]; ``conv_state`` [slots, C, 3] / ``windows`` [rows, 7, C, 3] + updated at ``state_indices`` [R] int32 and intermediate rows 0..R-1.""" + from sgl_nvidia_mega.qwen3_5.sm103 import trt_lowm + + act = "fp8" if x_fp8 is not None else "quant" + requests = x.shape[0] // ROWS + s = _SPECS.get((act, weight.shape[0], requests)) + if s is None: + raise RuntimeError( + f"the TensorRT-LLM GDN front ({act}, {weight.shape[0]} rows, {requests} requests) was not prepared" + ) + trt_lowm.run( + s, + out=out, + out2=ba, + x16=x, + x8=x_fp8, + w_fp8=weight, + weight_scale=weight_scale, + input_scale=input_scale, + w_bf16=ba_weight, + conv=(conv_weight, conv_state, state_indices, windows), + ) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/hop/__init__.py b/sgl_nvidia_mega/qwen3_5/sm103/hop/__init__.py index aa7a6033..9d2e143a 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/hop/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/hop/__init__.py @@ -29,10 +29,27 @@ SOURCE = Path(__file__).parent / "src" / "hop.cu" HEADER = Path(__file__).parent / "src" / "q35_hop.cuh" +# Slot route_lin (numeric): the router logits of a TP8 B1 step out of the mixer's +# all-reduce (src/rlin.cu), through a router mailbox next to the hidden one. +RL_SOURCE = Path(__file__).parent / "src" / "rlin.cu" +RL_EXPERTS, RL_WORLD, RL_CTAS, RL_STEPS = 512, 8, 16, 12 +RL_SENTINEL = -(2**31) # f32 -0.0 as int32: an empty router slot HIDDEN, STAGES, ROW_CTAS, ROW_THREADS, QUAD_CTAS, MAX_TOP_K = 4096, 3, 4, 128, 8, 16 OPROJ_THREADS, OPROJ_MAX_M, OPROJ_STEPS = 256, 8, (4, 8) # K = 256 * steps SENTINEL = -32768 # bf16 -0.0 as int16: an empty mailbox slot WORLDS = (2, 4, 8) +# int32 0, 1, 2, ... buffers (data_ptr -> length) whose prefixes are identity row maps of a finalize +_IDENTITY_ROWS: dict = {} +# the registered buffers, held for the process: a freed one (route.build_m32 replaces route.build's) would give its +# address to another tensor, whose rows a finalize then reads as the identity map +_IDENTITY_KEEP: list = [] + + +def register_identity_rows(rows: torch.Tensor) -> None: + """Record ``rows`` (int32 0, 1, 2, ...): a finalize whose ``permuted_indices`` is a contiguous prefix of it + passes no row map, so its routed rows need not wait for the map's load. The registry holds ``rows``.""" + _IDENTITY_KEEP.append(rows) + _IDENTITY_ROWS[rows.data_ptr()] = rows.numel() class Site(ctypes.Structure): @@ -77,6 +94,49 @@ class OprojArgs(ctypes.Structure): ) +class RlArgs(ctypes.Structure): + _fields_ = [ + ("rmc", ctypes.c_void_p), + ("ruc", ctypes.c_void_p), + ("rw", ctypes.c_void_p), + ("h0", ctypes.c_void_p), + ("k0", ctypes.c_int), + ("logits", ctypes.c_void_p), + ("x_fp4", ctypes.c_void_p), + ("x_sf", ctypes.c_void_p), + ("gs", ctypes.c_void_p), + ] + + +def pack_router_lin( + router_w: torch.Tensor, + gamma: torch.Tensor, + out_w_nk: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + rank: int, +) -> torch.Tensor: + """The router-before-all-reduce weights of one layer on this rank, packed as the + rl push's mma.m16n8k16 bf16 A fragments: [A W_out,r * sx * sw | A[:, K_r]] (512 x + (1024 + 512)), A = router_w [512, 4096] scaled by the post-attention norm's applied + gamma per column, W_out,r [4096, 1024] this rank's e4m3 out_proj (row-major), K_r = + columns [512 rank, 512 rank + 512); packed [16 cta][8 warp][12 step][2 mtile][32 lane] + x 8 bf16 (a0..a3 of lane (g, t): rows g / g + 8, columns 2t / 2t + 8 of the 16 x 16 + tile).""" + a = router_w.float() * gamma.float()[None, :] + n, k = out_w_nk.shape + if tuple(a.shape) != (RL_EXPERTS, HIDDEN) or (n, k) != (HIDDEN, 1024): + raise ValueError(f"router {tuple(a.shape)}, out_proj {(n, k)}") + core = (a @ out_w_nk.float()) * ( + input_scale.float() * weight_scale.float() + ).reshape(1, 1) + sl = HIDDEN // RL_WORLD + w = torch.cat([core, a[:, sl * rank : sl * (rank + 1)]], dim=1).to(torch.bfloat16) + # rows = cta * 32 + mtile * 16 + rh * 8 + g; cols = (warp * 12 + step) * 16 + kh * 8 + t * 2 + p + v = w.view(RL_CTAS, 2, 2, 8, 8, RL_STEPS, 2, 4, 2) + return v.permute(0, 4, 5, 1, 3, 7, 6, 2, 8).contiguous().view(-1) + + def pack_out_proj(weight_nk: torch.Tensor) -> torch.Tensor: """[4096, K] e4m3 row-major -> the push kernel's mma.m16n8k32 A fragments: [cta 128][warp 8][k-step S][block 2][lane 32][16 B], K = 8 * S * 32.""" @@ -116,11 +176,40 @@ def _kernels(device_index: int, deadline_ns: int) -> dict[str, jit.Kernel]: return kernels +@cache +def _rl_kernels(device_index: int, deadline_ns: int) -> dict[str, jit.Kernel]: + major, minor = torch.cuda.get_device_capability(device_index) + cubin = jit.compile_source( + RL_SOURCE.read_text(), + "q35_hop_rl", + f"sm_{major}{minor}a", + ("-std=c++17", "-default-device", f"-DQ35_HOP_DEADLINE_NS={deadline_ns}ull"), + headers(), + ) + return { + "q35hop_oproj_push_rl_s4": jit.load( + cubin, "q35hop_oproj_push_rl_s4", (OPROJ_THREADS, 1, 1), 0 + ), + "q35hop_consume_rl_w8": jit.load( + cubin, "q35hop_consume_rl_w8", (ROW_THREADS, 1, 1), 0 + ), + } + + class Hop: """One process's mailbox and control block, shared by every hop of every step; construct on every rank of ``group`` (a collective) before any capture.""" - def __init__(self, *, max_m: int, eps: float, group, device, deadline_s: float): + def __init__( + self, + *, + max_m: int, + eps: float, + group, + device, + deadline_s: float, + route_lin: bool = False, + ): import torch.distributed as dist import torch.distributed._symmetric_memory as symm_mem @@ -156,6 +245,25 @@ def __init__(self, *, max_m: int, eps: float, group, device, deadline_s: float): self.world, self.max_m, ) + # Slot route_lin's router mailbox [stages, world, rows, 512] f32. + self.rl = None + self._pushed_rl = False + if route_lin and self.world == RL_WORLD: + self._rl_box = symm_mem.empty( + (STAGES, self.world, self.max_m, RL_EXPERTS), + dtype=torch.int32, + device=self.device, + ) + rl_handle = symm_mem.rendezvous(self._rl_box, group) + rl_mc = int(rl_handle.multicast_ptr or 0) + if not rl_mc or rl_mc % 16: + raise RuntimeError( + "the router mailbox needs an NVLink multicast mapping" + ) + self._rl_box.fill_(RL_SENTINEL) + self._rl_handle = rl_handle + self.rl = (rl_mc, self._rl_box.data_ptr()) + self.rl_kernels = _rl_kernels(self.device.index, int(deadline_s * 1e9)) torch.cuda.synchronize(self.device) dist.barrier(group=group) @@ -169,6 +277,10 @@ def all_reduce_residual_rms_norm( partial out_proj_push returned is already in the mailbox.""" m = int(local_contribution.shape[0]) args, outs = self._args(m, residual, gamma, fp8_scale) + if self._pushed_rl: + raise RuntimeError( + "a router-lin push must be consumed by all_reduce_residual_rms_norm_rl" + ) if self.has_pushed(local_contribution): self._pushed = None else: @@ -205,6 +317,101 @@ def out_proj_push(self, *, x_fp8, packed_w, input_scale, weight_scale): self.mark_pushed(out) return out + def out_proj_push_rl(self, *, x_fp8, packed_w, input_scale, weight_scale, rl_w, h0): + """``out_proj_push`` of a TP8 B1 step plus this rank's router partial (16 more CTAs: + ``rl_w`` from ``pack_router_lin``, ``h0`` [M, 4096] the layer's residual input), + pushed into the router mailbox; ``all_reduce_residual_rms_norm_rl`` must come next.""" + m = int(x_fp8.shape[0]) + steps = packed_w.numel() // (HIDDEN * 256) + if self.rl is None or not 1 <= m <= min(OPROJ_MAX_M, self.max_m) or steps != 4: + raise ValueError( + f"the router-lin push serves TP8 out_proj s4 of 1-{OPROJ_MAX_M} rows, not {m}" + ) + if ( + h0.dtype != torch.bfloat16 + or tuple(h0.shape) != (m, HIDDEN) + or not h0.is_contiguous() + ): + raise ValueError(f"router-lin h0 {h0.dtype} {tuple(h0.shape)}") + self._check_consumed() + out = torch.empty((m, HIDDEN), dtype=torch.bfloat16, device=self.device) + args = OprojArgs( + site=self.site, + w=packed_w.data_ptr(), + x=x_fp8.data_ptr(), + sx=input_scale.data_ptr(), + sw=weight_scale.data_ptr(), + out=out.data_ptr(), + m=m, + ) + rargs = RlArgs( + rmc=self.rl[0], + ruc=self.rl[1], + rw=rl_w.data_ptr(), + h0=h0.data_ptr(), + k0=(HIDDEN // RL_WORLD) * self.rank, + ) + jit.launch( + self.rl_kernels["q35hop_oproj_push_rl_s4"], + (HIDDEN // 32 + RL_CTAS, 1, 1), + [args, rargs], + torch.cuda.current_stream(self.device).cuda_stream, + ) + self.mark_pushed(out) + self._pushed_rl = True + return out + + def rl_push_args(self, rl_w, h0) -> RlArgs: + """The router-partial launch arguments of another producer (the attention epilogue) that pushes this rank's + router partial; it then calls ``mark_pushed_rl`` on its hidden partial.""" + if self.rl is None or self.world != RL_WORLD: + raise ValueError("no router mailbox") + return RlArgs( + rmc=self.rl[0], + ruc=self.rl[1], + rw=rl_w.data_ptr(), + h0=h0.data_ptr(), + k0=(HIDDEN // RL_WORLD) * self.rank, + ) + + def mark_pushed_rl(self, partial: torch.Tensor) -> None: + """``mark_pushed`` of a producer that also pushed its router partial.""" + self.mark_pushed(partial) + self._pushed_rl = True + + def all_reduce_residual_rms_norm_rl( + self, *, local_contribution, residual, gamma, global_scale, fp8_scale=None + ): + """``all_reduce_residual_rms_norm`` of a router-lin push, plus the router's bf16 + logits [M, 512] and the routed experts' NVFP4 input (x_fp4 [M, 2048], x_sf + [M, 256], ``global_scale`` [1]): ((norm, residual_out[, fp8]), (logits, x_fp4, + x_sf)).""" + m = int(local_contribution.shape[0]) + if not (self._pushed_rl and self.has_pushed(local_contribution)): + raise RuntimeError("no router-lin push of this partial to consume") + args, outs = self._args(m, residual, gamma, fp8_scale) + logits = torch.empty((m, RL_EXPERTS), dtype=torch.bfloat16, device=self.device) + x_fp4 = torch.empty((m, HIDDEN // 2), dtype=torch.uint8, device=self.device) + x_sf = torch.empty((m, HIDDEN // 16), dtype=torch.uint8, device=self.device) + rargs = RlArgs( + rmc=self.rl[0], + ruc=self.rl[1], + logits=logits.data_ptr(), + x_fp4=x_fp4.data_ptr(), + x_sf=x_sf.data_ptr(), + gs=global_scale.data_ptr(), + ) + self._pushed = None + self._pushed_rl = False + jit.launch( + self.rl_kernels["q35hop_consume_rl_w8"], + (m, ROW_CTAS, 1), + [args, rargs], + torch.cuda.current_stream(self.device).cuda_stream, + (1, ROW_CTAS, 1), + ) + return outs, (logits, x_fp4, x_sf) + def mark_pushed(self, partial: torch.Tensor) -> None: """Record that a producer kernel, launched on this stream after the last all-reduce, pushed the bf16 ``partial`` [M, 4096] into the mailbox (through @@ -239,7 +446,13 @@ def moe_finalize_all_reduce_rms_norm( self._check_consumed() args, outs = self._args(m, residual, gamma, fp8_scale) args.routed, args.expert_w = routed_output.data_ptr(), expert_weights.data_ptr() - args.permuted, args.top_k = permuted_indices.data_ptr(), top_k + identity = ( + permuted_indices.is_contiguous() + and _IDENTITY_ROWS.get(permuted_indices.data_ptr(), 0) + >= permuted_indices.numel() + ) + args.permuted = None if identity else permuted_indices.data_ptr() + args.top_k = top_k args.shared = jit.null_or_pointer(gated_shared_output).value name = "q35hop_publish_fin_top10" if top_k == 10 else "q35hop_publish_fin" self._launch(name, args, QUAD_CTAS) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/hop/src/hop.cu b/sgl_nvidia_mega/qwen3_5/sm103/hop/src/hop.cu index 331575a0..d638e76a 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/hop/src/hop.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/hop/src/hop.cu @@ -55,21 +55,16 @@ __device__ __forceinline__ size_t row_elem(int token, int frag) { // sum_k w_k * routed[perm_k] in k order with separate multiply and add, then the shared output, // as the reference finalize does. TOPK > 0 fixes top-k at compile time, so every routed row is // loaded into registers before any accumulation (TOPK = 0 reads top_k at run time). + +// The routing slots are read after griddepcontrol.wait: the MoE that runs right before (mkm's expert weights, the +// router's tables) writes them, and PDL makes a predecessor's writes visible only after the wait. A null `permuted` +// is the identity (row token * top_k + k): then the routed rows are requested at once, with the weights. template __device__ __forceinline__ void publish_fin_quad_kernel(const Args& a) { constexpr int kSlots = TOPK > 0 ? TOPK : kMaxTopK; - __shared__ int32_t rows_s[kSlots]; - __shared__ float wts_s[kSlots]; const int top_k = TOPK > 0 ? TOPK : a.top_k; const int token = blockIdx.x; const int quad = blockIdx.y * kRowThreads + threadIdx.x; - if (threadIdx.x < top_k) { - const int i = token * top_k + threadIdx.x; - const int r = a.permuted[i]; - rows_s[threadIdx.x] = r; - wts_s[threadIdx.x] = r < 0 ? 0.f : __uint_as_float(static_cast(a.expert_w[i]) << 16); - } - __syncthreads(); const size_t col = static_cast(quad) * 4; pdl_wait(); const uint32_t stage = stage_peek(a.site); @@ -79,17 +74,24 @@ __device__ __forceinline__ void publish_fin_quad_kernel(const Args& a) { if (a.shared != nullptr) sh = *reinterpret_cast(a.shared + static_cast(token) * kHidden + col); int rows[kSlots]; #pragma unroll - for (int k = 0; k < kSlots; ++k) rows[k] = k < top_k ? rows_s[k] : -1; + for (int k = 0; k < kSlots; ++k) { + const int i = token * top_k + k; + rows[k] = k < top_k ? (a.permuted != nullptr ? a.permuted[i] : i) : -1; + } uint2 src[kSlots]; #pragma unroll for (int k = 0; k < kSlots; ++k) src[k] = rows[k] >= 0 ? *reinterpret_cast(a.routed + static_cast(rows[k]) * kHidden + col) : make_uint2(0, 0); + float wts[kSlots]; +#pragma unroll + for (int k = 0; k < kSlots; ++k) + wts[k] = rows[k] >= 0 ? __uint_as_float(static_cast(a.expert_w[token * top_k + k]) << 16) : 0.f; float acc[4] = {0.f, 0.f, 0.f, 0.f}; #pragma unroll for (int k = 0; k < kSlots; ++k) { if (rows[k] >= 0) { - const float w = wts_s[k]; + const float w = wts[k]; const float f[4] = {bf_lo(src[k].x), bf_hi(src[k].x), bf_lo(src[k].y), bf_hi(src[k].y)}; #pragma unroll for (int e = 0; e < 4; ++e) acc[e] = __fadd_rn(acc[e], __fmul_rn(f[e], w)); @@ -246,6 +248,7 @@ extern "C" __global__ void __launch_bounds__(q35hop::kRowThreads) q35hop_publish q35hop_kernels::publish_fin_quad_kernel<0>(a); } + extern "C" __global__ void __launch_bounds__(q35hop::kRowThreads) q35hop_consume_w2(Args a) { q35hop_kernels::consume_kernel<2>(a); } diff --git a/sgl_nvidia_mega/qwen3_5/sm103/hop/src/rlin.cu b/sgl_nvidia_mega/qwen3_5/sm103/hop/src/rlin.cu new file mode 100644 index 00000000..c3116600 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/hop/src/rlin.cu @@ -0,0 +1,447 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Slot route_lin (numeric): the MoE router's logits of a TP8 B1 verify step come +// out of the mixer's all-reduce instead of a router launch after it. RMSNorm scales its input by one number per row, +// so the router GEMV of the post-attention norm is +// logits = W_router . (h1 * inv_rms(h1) * gamma) = inv_rms(h1) * A . (h0 + sum_r attn_r), A = W_router diag(gamma). +// Each rank's out_proj launch adds 16 CTAs that compute its router partial (A W_out,r) . core_r + A[:, K_r] . h0[K_r] +// (the out_proj's static FP8 scales folded into the packed weights, fp32 accumulation) and push it, fp32, into a second +// mailbox with the hidden partial's stage; the consume sums the partials in rank order, scales them by the inv_rms it +// computes for the norm and writes the bf16 logits, and also writes the routed experts' NVFP4 input (the router front's +// quantizer, bit for bit). The hidden path (residual, norm, FP8 output) is the plugin's oproj_push_s4 + consume_w8, +// unchanged. +// +// q35hop_oproj_push_rl_s4 grid 128 + 16: CTAs 0..127 are oproj_push_s4's, CTAs 128..143 the router partials +// q35hop_consume_rl_w8 grid (M, 4), cluster (1, 4, 1): consume_w8 + the logits + the NVFP4 input +// +// Built with NVRTC on the hop's header (q35_hop.cuh). + +#include "q35_hop.cuh" + +namespace q35rl { + +using namespace q35hop; + +// = hop.cu's Args / OprojArgs (the ctypes structures of sm103/hop/__init__.py) +struct Args { + Site site; + const uint16_t* partial; + const uint16_t* routed; + const uint16_t* expert_w; + const int32_t* permuted; + const uint16_t* shared; + const uint16_t* residual_in; + const uint16_t* gamma; + const float* fp8_scale; + uint16_t* residual_out; + uint16_t* norm_out; + uint8_t* fp8_out; + int m, top_k; + float eps; +}; + +struct OprojArgs { + Site site; + const uint4* w; + const uint8_t* x; + const float* sx; + const float* sw; + uint16_t* out; + int m; +}; + +struct RlArgs { + void* rmc; // multicast base of the router mailbox [kStages][world][m_cap][kExperts] f32 + void* ruc; // this rank's view of it + const uint4* rw; // packed [A W_out,r * sx * sw | A[:, K_r]]: [16 cta][8 warp][12 step][2 mtile][32 lane] x 16 B + const uint16_t* h0; // [M, 4096] bf16: the layer's residual input + int k0; // this rank's h0 slice [k0, k0 + 512) + uint16_t* logits; // [M, 512] bf16 (consume) + uint8_t* x_fp4; // [M, 2048] (consume) + uint8_t* x_sf; // [M, 256] (consume) + const float* gs; // [1] the routed experts' NVFP4 global scale (consume) +}; + +constexpr int kExperts = 512; +constexpr int kHidCtas = kHidden / 32; // 128 +constexpr int kRlCtas = kExperts / 32; // 16 +static_assert(kRlCtas == 16, "the launcher's grid is 128 + 16 CTAs"); +constexpr int kCoreK = 1024; // core_r columns at TP8 +constexpr int kSliceK = kHidden / 8; // h0 columns a rank: 512 +constexpr int kRlSteps = (kCoreK + kSliceK) / 16 / 8; // k16 steps a warp: 12 +__device__ __forceinline__ bool is_router_cta() { return blockIdx.x >= kHidCtas; } +__device__ __forceinline__ int hid_cta() { return static_cast(blockIdx.x); } +__device__ __forceinline__ int rl_cta() { return static_cast(blockIdx.x) - kHidCtas; } +constexpr uint32_t kSentF32 = 0x80000000u; // f32 -0.0: an empty router slot + +__device__ __forceinline__ size_t rslot(const Site& s, uint32_t stage, int src, int token, int e) { + return ((static_cast(stage) * s.world + src) * s.m_cap + token) * kExperts + e; +} + +__device__ __forceinline__ size_t row_elem(int token, int frag) { + return static_cast(token) * kHidden + static_cast(frag) * kVec; +} +__device__ __forceinline__ uint4 ld_v4(const void* p) { return *static_cast(p); } +__device__ __forceinline__ uint4 ldg_nc_v4(const uint4* q) { + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(q)); + return r; +} + +__device__ __forceinline__ void mma_e4m3(float (&d)[4], const uint4& a, uint32_t b0, uint32_t b1) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b0), "r"(b1)); +} +__device__ __forceinline__ void mma_bf16(float (&d)[4], const uint4& a, uint32_t b0, uint32_t b1) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b0), "r"(b1)); +} +// two e4m3 (the low byte first) as a bf16 pair, exactly +__device__ __forceinline__ uint32_t e4m3x2_bf16x2(uint16_t v) { + float lo, hi; + asm("{\n\t.reg .b32 h2;\n\t.reg .b16 l, h;\n\t" + "cvt.rn.f16x2.e4m3x2 h2, %2;\n\t" + "mov.b32 {l, h}, h2;\n\t" + "cvt.f32.f16 %0, l;\n\t" + "cvt.f32.f16 %1, h;\n\t}" + : "=f"(lo), "=f"(hi) + : "h"(v)); + return pack_bf16x2(lo, hi); +} +__device__ __forceinline__ void mc_st_v4f(float* mc, float a, float b, float c, float d) { + asm volatile("multimem.st.relaxed.sys.global.v4.f32 [%0], {%1,%2,%3,%4};" ::"l"(mc), "f"(a), "f"(b), "f"(c), "f"(d) + : "memory"); +} +__device__ __forceinline__ float sanitize_f32(float x) { return __float_as_uint(x) == kSentF32 ? 0.f : x; } + +// ---------------------------------------------------------------- out_proj + router partials + +// = hop.cu's oproj_push_kernel<4>: CTA c owns output rows [32c, 32c + 32) over the full K +__device__ __forceinline__ void oproj_rows(const OprojArgs& o, const RlArgs& r, float (&red)[8][32][8]) { + constexpr int S = 4; + constexpr int kK = 8 * S * 32; + const int cta = hid_cta(); + const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + const int g = lane >> 2, t = lane & 3; + uint4 a[S][2]; + { + const uint4* wb = o.w + (static_cast(cta) * 8 + warp) * (S * 2 * 32) + lane; +#pragma unroll + for (int st = 0; st < S; ++st) +#pragma unroll + for (int b = 0; b < 2; ++b) a[st][b] = ldg_nc_v4(wb + (st * 2 + b) * 32); + } + const float alpha = __fmul_rn(*o.sx, *o.sw); + pdl_wait(); + const int m = o.m; + const bool tok_ok = g < m; + const uint8_t* xr = o.x + static_cast(g) * kK + warp * S * 32 + 4 * t; + float ch[2][2][4] = {}; +#pragma unroll + for (int st = 0; st < S; ++st) { + const uint32_t b0 = tok_ok ? *reinterpret_cast(xr + st * 32) : 0u; + const uint32_t b1 = tok_ok ? *reinterpret_cast(xr + st * 32 + 16) : 0u; + mma_e4m3(ch[0][st & 1], a[st][0], b0, b1); + mma_e4m3(ch[1][st & 1], a[st][1], b0, b1); + } + { + float4* dst = reinterpret_cast(&red[warp][lane][0]); + dst[0] = make_float4(__fadd_rn(ch[0][0][0], ch[0][1][0]), __fadd_rn(ch[0][0][1], ch[0][1][1]), + __fadd_rn(ch[0][0][2], ch[0][1][2]), __fadd_rn(ch[0][0][3], ch[0][1][3])); + dst[1] = make_float4(__fadd_rn(ch[1][0][0], ch[1][1][0]), __fadd_rn(ch[1][0][1], ch[1][1][1]), + __fadd_rn(ch[1][0][2], ch[1][1][2]), __fadd_rn(ch[1][0][3], ch[1][1][3])); + } + __syncthreads(); + const int tok = tid >> 2, j = tid & 3; + if (tok < m) { + float v[8]; +#pragma unroll + for (int q = 0; q < 8; ++q) { + const int r = 8 * j + q; + const int l = (r % 8) * 4 + tok / 2; + const int e = (r / 16) * 4 + ((r % 16) / 8) * 2 + (tok % 2); + float x = red[0][l][e]; +#pragma unroll + for (int w = 1; w < 8; ++w) x = __fadd_rn(x, red[w][l][e]); + v[q] = __fmul_rn(x, alpha); + } + const uint4 frag = pack8(v); + const int fr = cta * 4 + j; + push_frag(o.site, stage_peek(o.site), tok, fr, frag); + *reinterpret_cast(o.out + row_elem(tok, fr)) = frag; + } + pdl_trigger(); +} + +// CTA 128 + c: router partial rows [32c, 32c + 32) of this rank, K = core_r (1024, e4m3 as bf16) then h0[K_r] (512). +// The B operands (x [M, 1024] e4m3 and h0's slice [M, 512] bf16) are staged in shared memory as bf16 first (two +// 16 B loads a thread), then each warp runs its 12 k-steps in 4 accumulator chains a row tile. +struct RlSmem { + uint32_t xb[8][(kCoreK + kSliceK) / 2 + 4]; // [token][k pair] bf16x2, row padded against bank conflicts +}; +__device__ __forceinline__ void router_rows(const OprojArgs& o, const RlArgs& r, float (&red)[8][32][8], RlSmem& sm) { + const int c = rl_cta(); + const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + const int g = lane >> 2, t = lane & 3; + uint4 a[kRlSteps][2]; + { + const uint4* wb = r.rw + (static_cast(c) * 8 + warp) * (kRlSteps * 2 * 32) + lane; +#pragma unroll + for (int st = 0; st < kRlSteps; ++st) +#pragma unroll + for (int mt = 0; mt < 2; ++mt) a[st][mt] = ldg_nc_v4(wb + (st * 2 + mt) * 32); + } + pdl_wait(); + const int m = o.m; + // stage: token rows < m; x: 7 x 1024 B = 64 chunks of 16 B a row; h0 slice: 7 x 1024 B = 64 chunks a row. All four + // loads of a thread go out before any is used (one L2 round trip, not four). + uint4 ld[4]; +#pragma unroll + for (int u = 0; u < 4; ++u) { + const int i = tid + u * 256, row = i >> 7, chk = i & 127; + ld[u] = make_uint4(0u, 0u, 0u, 0u); + if (row < m) { + ld[u] = chk < 64 ? *reinterpret_cast(o.x + static_cast(row) * kCoreK + chk * 16) + : *reinterpret_cast(r.h0 + static_cast(row) * kHidden + r.k0 + (chk - 64) * 8); + } + } +#pragma unroll + for (int u = 0; u < 4; ++u) { + const int i = tid + u * 256, row = i >> 7, chk = i & 127; + const uint32_t vw[4] = {ld[u].x, ld[u].y, ld[u].z, ld[u].w}; + if (chk < 64) { +#pragma unroll + for (int q = 0; q < 4; ++q) { + sm.xb[row][chk * 8 + 2 * q] = e4m3x2_bf16x2(static_cast(vw[q] & 0xFFFFu)); + sm.xb[row][chk * 8 + 2 * q + 1] = e4m3x2_bf16x2(static_cast(vw[q] >> 16)); + } + } else { +#pragma unroll + for (int q = 0; q < 4; ++q) sm.xb[row][kCoreK / 2 + (chk - 64) * 4 + q] = vw[q]; + } + } + __syncthreads(); + float ch[2][4][4] = {}; +#pragma unroll + for (int st = 0; st < kRlSteps; ++st) { + const int kk = (warp * kRlSteps + st) * 16; + const uint32_t b0 = sm.xb[g][kk / 2 + t]; + const uint32_t b1 = sm.xb[g][kk / 2 + 4 + t]; + mma_bf16(ch[0][st & 3], a[st][0], b0, b1); + mma_bf16(ch[1][st & 3], a[st][1], b0, b1); + } + { + float4* dst = reinterpret_cast(&red[warp][lane][0]); +#pragma unroll + for (int mt = 0; mt < 2; ++mt) { + float s4[4]; +#pragma unroll + for (int e = 0; e < 4; ++e) + s4[e] = __fadd_rn(__fadd_rn(ch[mt][0][e], ch[mt][1][e]), __fadd_rn(ch[mt][2][e], ch[mt][3][e])); + dst[mt] = make_float4(s4[0], s4[1], s4[2], s4[3]); + } + } + __syncthreads(); + const int tok = tid >> 2, j = tid & 3; + if (tok < m) { + float v[8]; +#pragma unroll + for (int q = 0; q < 8; ++q) { + const int rr = 8 * j + q; + const int l = (rr % 8) * 4 + tok / 2; + const int e = (rr / 16) * 4 + ((rr % 16) / 8) * 2 + (tok % 2); + float x = red[0][l][e]; +#pragma unroll + for (int w = 1; w < 8; ++w) x = __fadd_rn(x, red[w][l][e]); + v[q] = sanitize_f32(x); + } + float* dst = static_cast(r.rmc) + rslot(o.site, stage_peek(o.site), o.site.rank, tok, c * 32 + 8 * j); + mc_st_v4f(dst, v[0], v[1], v[2], v[3]); + mc_st_v4f(dst + 4, v[4], v[5], v[6], v[7]); + } + pdl_trigger(); +} + +// ---------------------------------------------------------------- consume + logits + NVFP4 input + +__device__ __forceinline__ uint32_t bf16x2_abs_max(uint32_t a, uint32_t b) { + uint32_t out; + asm("{\n.reg .b32 ta, tb;\nand.b32 ta, %1, 0x7FFF7FFF;\nand.b32 tb, %2, 0x7FFF7FFF;\nmax.bf16x2 %0, ta, tb;\n}" + : "=r"(out) + : "r"(a), "r"(b)); + return out; +} +__device__ __forceinline__ uint32_t bf16x2_max(uint32_t a, uint32_t b) { + uint32_t out; + asm("max.bf16x2 %0, %1, %2;" : "=r"(out) : "r"(a), "r"(b)); + return out; +} +__device__ __forceinline__ float rcp_approx_ftz(float a) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(a)); + return y; +} +__device__ __forceinline__ uint32_t cvt_e4m3(float a) { + uint32_t out; + asm("{\n.reg .b16 pair;\n.reg .f32 zero;\nmov.f32 zero, 0f00000000;\n" + "cvt.rn.satfinite.e4m3x2.f32 pair, zero, %1;\ncvt.u32.u16 %0, pair;\n}" + : "=r"(out) + : "f"(a)); + return out & 0xFFu; +} +__device__ __forceinline__ float e4m3_to_f32(uint32_t v) { + float out; + asm("{\n.reg .b16 pair;\n.reg .b32 h2;\n.reg .b16 lo, hi;\ncvt.u16.u32 pair, %1;\n" + "cvt.rn.f16x2.e4m3x2 h2, pair;\nmov.b32 {lo, hi}, h2;\ncvt.f32.f16 %0, lo;\n}" + : "=f"(out) + : "r"(v)); + return out; +} +__device__ __forceinline__ uint32_t cvt_e2m1x8(const float (&v)[8]) { + uint32_t out; + asm("{\n.reg .b8 b0, b1, b2, b3;\n" + "cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\ncvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\ncvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" + "mov.b32 %0, {b0, b1, b2, b3};\n}" + : "=r"(out) + : "f"(v[0]), "f"(v[1]), "f"(v[2]), "f"(v[3]), "f"(v[4]), "f"(v[5]), "f"(v[6]), "f"(v[7])); + return out; +} + +// pull_frag (the hidden fragment's rank-order sum, bit for bit) and logit e's router partials polled together: both +// rounds of loads go out at once, then only the pending ones again; deadline as pull_frag. +template +__device__ __forceinline__ void pull_both(const RlArgs& r, const Site& s, uint32_t stage, int token, int frag, + float acc[8], float& lsum) { + const uint16_t* uc = static_cast(s.uc); + const uint32_t* ru = static_cast(r.ruc); + uint4 v[W]; + uint32_t u[W]; +#pragma unroll + for (int k = 0; k < W; ++k) v[k] = ld_volatile_v4(uc + slot(s, stage, k, token, frag)); +#pragma unroll + for (int k = 0; k < W; ++k) u[k] = ld_volatile_u32(ru + rslot(s, stage, k, token, frag)); + uint32_t pending = 0, rpend = 0; +#pragma unroll + for (int k = 0; k < W; ++k) pending |= frag_has_negzero(v[k]) ? (1u << k) : 0u; +#pragma unroll + for (int k = 0; k < W; ++k) rpend |= u[k] == kSentF32 ? (1u << k) : 0u; + unsigned long long t0 = 0; + uint32_t spins = 0; + while (pending | rpend) { +#pragma unroll + for (int k = 0; k < W; ++k) + if (pending & (1u << k)) v[k] = ld_volatile_v4(uc + slot(s, stage, k, token, frag)); +#pragma unroll + for (int k = 0; k < W; ++k) + if (rpend & (1u << k)) u[k] = ld_volatile_u32(ru + rslot(s, stage, k, token, frag)); + uint32_t still = 0, rstill = 0; +#pragma unroll + for (int k = 0; k < W; ++k) + if ((pending & (1u << k)) && frag_has_negzero(v[k])) still |= 1u << k; +#pragma unroll + for (int k = 0; k < W; ++k) + if ((rpend & (1u << k)) && u[k] == kSentF32) rstill |= 1u << k; + pending = still; + rpend = rstill; + if (Q35_HOP_DEADLINE_NS != 0 && (pending | rpend) && ((++spins & 255u) == 0)) { + const unsigned long long now = gtimer(); + if (t0 == 0) { + t0 = now; + } else if (now - t0 > Q35_HOP_DEADLINE_NS) { + fault_at_null(); + } + } + } +#pragma unroll + for (int e = 0; e < 8; ++e) acc[e] = 0.f; +#pragma unroll + for (int k = 0; k < W; ++k) { + float f[8]; + unpack8(v[k], f); +#pragma unroll + for (int e = 0; e < 8; ++e) acc[e] = __fadd_rn(acc[e], f[e]); + } + lsum = 0.f; +#pragma unroll + for (int k = 0; k < W; ++k) lsum = __fadd_rn(lsum, __uint_as_float(u[k])); +} + +// = hop.cu's consume_kernel, plus: logit `frag` of the token from the router partials, and the NVFP4 input of the +// norm row (block frag / 2 is this thread's 8 values and its lane neighbour's), as the router front quantizes it +template +__device__ __forceinline__ void consume_rl_kernel(const Args& a, const RlArgs& r) { + __shared__ RowSmem rs; + const int token = blockIdx.x; + const int frag = blockIdx.y * kRowThreads + threadIdx.x; + pdl_trigger(); + float gam[8]; + unpack8(ld_v4(a.gamma + static_cast(frag) * kVec), gam); + const float inv_s = a.fp8_out != nullptr ? inv_scale_of(*a.fp8_scale) : 1.f; + const float gs = *r.gs; + row_init(rs); + pdl_wait(); + const Hop h = hop_enter(a.site); + const uint4 res = ld_v4(a.residual_in + row_elem(token, frag)); + float acc[8], lsum = 0.f; + pull_both(r, a.site, h.stage, token, frag, acc, lsum); + float pre[8]; + const uint4 res_out = residual_add(acc, res, pre); + const float total = row_sumsq_fold(sumsq8(pre), rs); + const float inv = inv_rms(total, a.eps); + const uint4 norm = norm_apply(pre, inv, gam); + *reinterpret_cast(a.residual_out + row_elem(token, frag)) = res_out; + if (a.norm_out != nullptr) *reinterpret_cast(a.norm_out + row_elem(token, frag)) = norm; + if (a.fp8_out != nullptr) *reinterpret_cast(a.fp8_out + row_elem(token, frag)) = quant_fp8(norm, inv_s); + r.logits[static_cast(token) * kExperts + frag] = static_cast(pack_bf16x2(__fmul_rn(lsum, inv), 0.f)); + { + const uint32_t w4[4] = {norm.x, norm.y, norm.z, norm.w}; + const uint32_t mine = bf16x2_max(bf16x2_abs_max(w4[0], w4[1]), bf16x2_abs_max(w4[2], w4[3])); + const uint32_t peer = __shfl_xor_sync(0xffffffffu, mine, 1); + const uint32_t m = (frag & 1) ? bf16x2_max(peer, mine) : bf16x2_max(mine, peer); + const float amax = fmaxf(bf_lo(m), bf_hi(m)); + const float scale = __fmul_rn(gs, __fmul_rn(amax, rcp_approx_ftz(6.0f))); + const uint32_t sf = cvt_e4m3(scale); + const float sf_f = e4m3_to_f32(sf); + const float out_scale = (sf_f == 0.0f) ? 0.0f : rcp_approx_ftz(__fmul_rn(sf_f, rcp_approx_ftz(gs))); + float q[8]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + q[2 * i] = __fmul_rn(bf_lo(w4[i]), out_scale); + q[2 * i + 1] = __fmul_rn(bf_hi(w4[i]), out_scale); + } + const int blk = frag >> 1; + *reinterpret_cast(r.x_fp4 + static_cast(token) * (kHidden / 2) + blk * 8 + (frag & 1) * 4) = + cvt_e2m1x8(q); + if (!(frag & 1)) r.x_sf[static_cast(token) * (kHidden / 16) + blk] = static_cast(sf); + } + rearm_frag(a.site, h.stage, token, frag); + uint32_t* uc = static_cast(r.ruc); +#pragma unroll + for (int k = 0; k < W; ++k) uc[rslot(a.site, h.stage, k, token, frag)] = kSentF32; + hop_leave(a.site, h, gridDim.x * gridDim.y); +} + +} // namespace q35rl + +extern "C" __global__ void __launch_bounds__(256, 1) q35hop_oproj_push_rl_s4(q35rl::OprojArgs o, q35rl::RlArgs r) { + __shared__ __align__(16) float red[8][32][8]; + __shared__ __align__(16) q35rl::RlSmem sm; + if (!q35rl::is_router_cta()) { + q35rl::oproj_rows(o, r, red); + } else { + q35rl::router_rows(o, r, red, sm); + } +} + +extern "C" __global__ void __launch_bounds__(q35hop::kRowThreads) q35hop_consume_rl_w8(q35rl::Args a, q35rl::RlArgs r) { + q35rl::consume_rl_kernel<8>(a, r); +} diff --git a/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/__init__.py b/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/__init__.py new file mode 100644 index 00000000..53b39994 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/__init__.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""TensorRT-LLM r3's low-M GEMM (``LowMMmaGemmKernel``, CuTe DSL: ``lowm_mma_gemm.py`` with ``utils.py`` and +``lowm_linear.py``, vendored from tekit 7df45abe5c) for callers that give their own shapes and tactics. + +``spec(...)`` names one launch by TensorRT-LLM's compile key (``lowm_linear._compiled``): rows, K, the FP8 and BF16 +weight segments with their tactics, the activation operand, the output split, the activations' row strides and the +GDN conv epilogue's tensor layouts (``conv_spec``); the kernel is specialized on all of them. ``prepare(spec)`` +compiles it and refuses inside a graph capture: prepare every spec a step can take before the first capture (in a +package's build, or where the pools first exist for a spec that names them). ``run(spec, ...)`` launches a prepared +spec into the caller's buffers and refuses a spec that was not prepared, so a step never compiles. + +Numerics are TensorRT-LLM's (``lowm_mma_gemm.py`` module doc): FP32 sums in its order, not SGLang's GEMMs bit for +bit. Every launch is a programmatic dependent that reads only weights, scales and conv weights ahead of its +``griddepcontrol.wait`` and, with ``early_trigger``, releases its dependents once those loads are issued (they wait +for its completion before reading its outputs). +""" + +from __future__ import annotations + +import logging +import time +from typing import NamedTuple, Optional + +import torch + +logger = logging.getLogger(__name__) + + +class Spec(NamedTuple): + """One compiled launch: the arguments of ``lowm_linear._compiled``, in its order.""" + + m: int + k: int + k_warps: int + n_fp8: int + n_tiles: int + fp8_act: str + swiglu_fp8: bool + n_bf16: int + n_tiles_bf16: int + swiglu_bf16: bool + split_out: int + n_valid: int + quant_out: bool + n_out: int + n_out2: int + out_dtype: torch.dtype + a16_stride: int + a8_stride: int + use_pdl: bool + early_trigger: bool + release_after_wait: bool + unsafe_skip_wait: bool = False + ksplit: int = 1 + swap_ab: bool = False + min_blocks: int = 1 + conv_channels: int = 0 + conv_tokens: int = 1 + conv_write_state: bool = True + conv_shapes: tuple = () + + +def conv_spec( + weight, state, windows, requests: int, tokens: int, write_state: bool = True +) -> tuple: + """The GDN conv epilogue of a spec, from the tensors ``run`` will get (the kernel is specialized on their shapes + and strides): ``weight`` [C, 4] BF16, ``state`` [slots, C, 3] BF16, ``windows`` [>= requests, tokens, C, 3] BF16; + ``requests`` of ``tokens`` rows (``lowm_linear.lowm_gemm``'s ``conv``).""" + if weight.dim() != 2 or weight.shape[1] != 4 or weight.dtype != torch.bfloat16: + raise ValueError(f"conv weight {weight.dtype} {tuple(weight.shape)}") + if ( + state.dim() != 3 + or tuple(state.shape[1:]) != (weight.shape[0], 3) + or state.dtype != torch.bfloat16 + ): + raise ValueError(f"conv state {state.dtype} {tuple(state.shape)}") + if ( + windows.dim() != 4 + or tuple(windows.shape[1:]) != (tokens, weight.shape[0], 3) + or windows.dtype != torch.bfloat16 + or windows.shape[0] < requests + ): + raise ValueError(f"conv windows {windows.dtype} {tuple(windows.shape)}") + return ( + weight.shape[0], + tokens, + write_state, + ( + (tuple(weight.shape), tuple(weight.stride())), + (tuple(state.shape), tuple(state.stride())), + requests, + (tuple(windows.shape), tuple(windows.stride())), + ), + ) + + +def spec( + *, + m: int, + k: int, + n_fp8: int = 0, + n_bf16: int = 0, + fp8_act: str = "fp8", + tactic: Optional[tuple] = None, + tactic_bf16: Optional[tuple] = None, + swiglu_fp8: bool = False, + swiglu_bf16: bool = False, + split_out: int = 0, + n_valid: int = 0, + quant_out: bool = False, + out_dtype: torch.dtype = torch.bfloat16, + a16_stride: Optional[int] = None, + a8_stride: Optional[int] = None, + pdl: bool = True, + early_trigger: bool = True, + release_after_wait: bool = False, + conv: Optional[tuple] = None, +) -> Spec: + """A launch of ``m`` rows of K ``k`` over ``n_fp8`` FP8 and ``n_bf16`` BF16 weight rows, resolved as + ``lowm_linear.lowm_gemm`` does. ``tactic`` / ``tactic_bf16``: (k_warps, n_tiles[, ksplit[, swap_ab[, + min_blocks]]]) of the FP8 / BF16 segment, required for a non-empty segment (TensorRT-LLM's choice per shape: + ``lowm_linear.default_tactic``); one k_warps and ksplit for both (the BF16 segment takes one tile a CTA + otherwise). ``fp8_act``: ``fp8`` (E4M3 rows ``x8``), ``quant`` (BF16 rows quantized in registers) or ``bf16``. + ``a16_stride`` / ``a8_stride``: the row strides (elements) of the activations ``run`` will get (default ``k``). + ``pdl`` / ``early_trigger`` / ``release_after_wait`` select the dependent-launch + behavior explicitly. ``conv``: ``conv_spec(...)``.""" + from sgl_nvidia_mega.qwen3_5.sm103.trt_lowm.lowm_linear import _tactic3 + + if (n_fp8 and tactic is None) or (n_bf16 and tactic_bf16 is None): + raise ValueError("every non-empty segment needs its tactic") + k_warps, n_tiles, ksplit, swap, min_blocks = ( + _tactic3(tactic) if n_fp8 else (0, 1, 1, 0, 1) + ) + kw16, nt16, ks16, swap16, mb16 = ( + _tactic3(tactic_bf16) if n_bf16 else (0, 1, 1, 0, 1) + ) + if not n_fp8: + k_warps, ksplit, swap, min_blocks = kw16, ks16, swap16, mb16 + elif n_bf16 and (kw16, ks16) != (k_warps, ksplit): + nt16 = 1 + n_logical = (n_fp8 // 2 if swiglu_fp8 else n_fp8) + ( + n_bf16 // 2 if swiglu_bf16 else n_bf16 + ) + if n_valid: + if not split_out < n_valid <= n_logical: + raise ValueError( + f"n_valid {n_valid} with split_out {split_out} of {n_logical} columns" + ) + n_logical = n_valid + uses_a16 = n_bf16 > 0 or (n_fp8 > 0 and fp8_act != "fp8") + uses_a8 = n_fp8 > 0 and fp8_act == "fp8" + conv_args = conv if conv is not None else (0, 1, True, ()) + return Spec( + m, + k, + k_warps, + n_fp8, + n_tiles, + fp8_act, + swiglu_fp8, + n_bf16, + nt16, + swiglu_bf16, + split_out, + n_valid, + quant_out, + split_out if split_out else n_logical, + n_logical - split_out if split_out else 0, + out_dtype, + (a16_stride or k) if uses_a16 else 0, + (a8_stride or k) if uses_a8 else 0, + pdl, + early_trigger and not release_after_wait, + release_after_wait, + False, + ksplit, + bool(swap) and m <= 8, + min_blocks, + *conv_args, + ) + + +_PREPARED: dict = {} + + +def prepare(s: Spec) -> None: + """Compile ``s`` once (CuTe DSL), outside graph capture.""" + if s in _PREPARED: + return + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"TensorRT-LLM low-M GEMM {s}: prepare it before graph capture" + ) + from sgl_nvidia_mega.qwen3_5.sm103.trt_lowm.lowm_linear import _compiled, _dummy + + dev = torch.device("cuda", torch.cuda.current_device()) + # the placeholder operands of unused segments, made here and not inside a capture + for kind in ("a16", "a8", "b8", "b16", "s", "c", "cw", "cs", "slot", "win"): + _dummy(kind, dev) + start = time.perf_counter() + _PREPARED[s] = _compiled(*s) + logger.info( + "NVIDIA megakernel: TensorRT-LLM low-M GEMM compiled in %.1f s (m=%d k=%d fp8=%d/%s bf16=%d, " + "k_warps=%d n_tiles=%d/%d ksplit=%d swap_ab=%d min_blocks=%d conv=%d)", + time.perf_counter() - start, + s.m, + s.k, + s.n_fp8, + s.fp8_act, + s.n_bf16, + s.k_warps, + s.n_tiles, + s.n_tiles_bf16, + s.ksplit, + s.swap_ab, + s.min_blocks, + s.conv_channels, + ) + + +def prepared(s: Spec) -> bool: + return s in _PREPARED + + +def _check( + t: torch.Tensor, name: str, shape: tuple, stride: tuple, dtypes: tuple +) -> None: + if t is None: + raise ValueError(f"{name} is required") + if ( + tuple(t.shape) != shape + or tuple(t.stride()) != stride + or t.dtype not in dtypes + or t.data_ptr() % 16 + ): + raise ValueError( + f"{name} {t.dtype} {tuple(t.shape)} {tuple(t.stride())} is not the prepared " + f"{dtypes[0]} {shape} {stride} (16-byte aligned)" + ) + + +def run( + s: Spec, + *, + out: torch.Tensor, + out2: Optional[torch.Tensor] = None, + x16: Optional[torch.Tensor] = None, + x8: Optional[torch.Tensor] = None, + w_fp8: Optional[torch.Tensor] = None, + weight_scale: Optional[torch.Tensor] = None, + input_scale: Optional[torch.Tensor] = None, + w_bf16: Optional[torch.Tensor] = None, + out_scale: Optional[torch.Tensor] = None, + conv: Optional[tuple] = None, +) -> None: + """Launch the prepared ``s`` on the current stream: ``out`` [m, n_out] (``out_dtype``, E4M3 with ``quant_out``) + and ``out2`` [m, n_out2] contiguous; ``x16`` [m, k] BF16 / ``x8`` [m, k] E4M3 rows at the spec's row strides; + ``w_fp8`` [n_fp8, k] E4M3 with ``weight_scale`` and ``input_scale`` (FP32 [1]); ``w_bf16`` [n_bf16, k]; + ``out_scale`` with ``quant_out``; ``conv`` = (weight, state, slots [requests] int32, windows) at the layouts of + the spec's ``conv_spec`` (TensorRT-LLM ``lowm_gemm``'s operands).""" + import cuda.bindings.driver as cuda + + from sgl_nvidia_mega.qwen3_5.sm103.trt_lowm.lowm_linear import _cute, _dummy + + fn = _PREPARED.get(s) + if fn is None: + raise RuntimeError( + f"TensorRT-LLM low-M GEMM {s} was not prepared (no compile inside a step)" + ) + dev = out.device + e4m3 = (torch.float8_e4m3fn, torch.uint8) + _check( + out, + "out", + (s.m, s.n_out), + (s.n_out, 1), + e4m3 if s.quant_out else (s.out_dtype,), + ) + if s.n_out2: + _check(out2, "out2", (s.m, s.n_out2), (s.n_out2, 1), (s.out_dtype,)) + if s.a16_stride: + _check(x16, "x16", (s.m, s.k), (s.a16_stride, 1), (torch.bfloat16,)) + if s.a8_stride: + _check(x8, "x8", (s.m, s.k), (s.a8_stride, 1), e4m3) + if s.n_fp8: + _check(w_fp8, "w_fp8", (s.n_fp8, s.k), (s.k, 1), e4m3) + if s.n_bf16: + _check(w_bf16, "w_bf16", (s.n_bf16, s.k), (s.k, 1), (torch.bfloat16,)) + if s.conv_channels: + (cw, cws), (cs, css), n_req, (win, wins) = s.conv_shapes + weight, state, slots, windows = conv + _check(weight, "conv weight", cw, cws, (torch.bfloat16,)) + _check(state, "conv state", cs, css, (torch.bfloat16,)) + _check(slots, "conv slots", (n_req,), (1,), (torch.int32,)) + _check(windows, "conv windows", win, wins, (torch.bfloat16,)) + conv_tensors = (weight, state, slots, windows) + else: + conv_tensors = ( + _dummy("cw", dev), + _dummy("cs", dev), + _dummy("slot", dev), + _dummy("win", dev), + ) + ones = _dummy("s", dev) + fn( + _cute(x16 if s.a16_stride else _dummy("a16", dev)), + _cute(x8.view(torch.uint8) if s.a8_stride else _dummy("a8", dev)), + _cute(w_fp8.view(torch.uint8) if s.n_fp8 else _dummy("b8", dev)), + _cute(weight_scale.reshape(1) if s.n_fp8 else ones), + _cute(input_scale.reshape(1) if (s.n_fp8 and s.fp8_act != "bf16") else ones), + _cute(out_scale.reshape(1) if s.quant_out else ones), + _cute(w_bf16 if s.n_bf16 else _dummy("b16", dev)), + _cute(out.view(torch.uint8) if s.quant_out else out), + _cute(out2 if s.n_out2 else _dummy("c", dev)), + *(_cute(t) for t in conv_tensors), + cuda.CUstream(torch.cuda.current_stream(dev).cuda_stream), + ) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/lowm_linear.py b/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/lowm_linear.py new file mode 100644 index 00000000..707cc780 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/lowm_linear.py @@ -0,0 +1,291 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# Adapted from TensorRT-LLM r3 (tekit 7df45abe5c), +# tensorrt_llm/_torch/modules/lowm_linear.py. +"""The TensorRT-LLM low-M compiler helpers used by the Qwen3.5 GDN front. + +Only the validated TP8 GDN projection shapes and TensorRT-LLM's graph-safe CuTe +compiler adapter are kept here. Slot selection, row bounds and dependent-launch +policy belong to the plugin plan and its caller, rather than environment switches. +""" + +from __future__ import annotations + +import functools +from typing import Optional, Tuple + +import torch + +_NUM_SMS_DEFAULT = 148 + + +@functools.lru_cache(maxsize=None) +def _num_sms() -> int: + if not torch.cuda.is_available(): + return _NUM_SMS_DEFAULT + return torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + + +Tactic = Tuple[int, int, int, int, int] + + +def _tactic3(t) -> Tactic: + """Normalize a two- to five-entry tactic to five entries.""" + return ( + t[0], + t[1], + t[2] if len(t) > 2 else 1, + t[3] if len(t) > 3 else 0, + t[4] if len(t) > 4 else 1, + ) + + +# TensorRT-LLM r3 tactics used by the TP8 GDN front. The final entry is also +# the fixed fallback beyond its measured crossover, through four requests. +_TACTICS = { + (2560, 4096, "mixed"): ((8, (16, 4, 1, 1)), (16, (16, 4))), + (2560, 4096, "mixed8"): ((8, (16, 2, 1, 1, 2)), (16, (8, 4))), +} + + +def lowm_conv_supported(m: int, tactic, conv_channels: int, tokens: int) -> bool: + """Whether the tactic supports the GDN causal-conv epilogue.""" + k_warps, n_tiles = tactic[0], tactic[1] + ksplit = tactic[2] if len(tactic) > 2 else 1 + cols = 8 * n_tiles + return ( + ksplit == 1 + and tokens > 0 + and m % tokens == 0 + and conv_channels % cols == 0 + and (m // tokens) * cols <= 32 * k_warps + ) + + +def default_tactic( + m: int, + n: int, + k: int, + fp8: bool = True, + even_tiles: bool = False, + mode: Optional[str] = None, +) -> Tactic: + """Return TensorRT-LLM's fixed tactic for an ``m x n x k`` GEMM.""" + entries = _TACTICS.get((n, k, mode or ("quant" if fp8 else "bf16"))) + if entries: + for max_rows, tactic in entries: + if m <= max_rows or max_rows == entries[-1][0]: + normalized = _tactic3(tactic) + k_warps, n_tiles, ksplit, swap, min_blocks = normalized + if even_tiles and n_tiles % 2: + n_tiles *= 2 + if n % (8 * n_tiles) == 0 and (k // 64) % (k_warps * ksplit) == 0: + return k_warps, n_tiles, ksplit, swap, min_blocks + break + + chunks = k // 64 + per_chunk = 16 if fp8 else 32 + sms = _num_sms() + best = None + for n_tiles in (1, 2, 4): + if n % (8 * n_tiles) or (even_tiles and n_tiles % 2): + continue + ctas = n // (8 * n_tiles) + for k_warps in (1, 2, 4, 8, 16): + if chunks % k_warps: + continue + per_warp = chunks // k_warps + if per_warp * n_tiles * per_chunk > 256: + continue + waves = -(-ctas // (2 * sms)) + key = (waves, per_warp > 8, -min(ctas, 2 * sms), per_warp, -k_warps) + if best is None or key < best[0]: + best = (key, (k_warps, n_tiles, 1, 0, 1)) + if best is None: + raise ValueError(f"no low-M tactic for N={n}, K={k}") + return best[1] + + +class _GraphSafeDlpack: + """DLPack export that does not synchronize streams (CUDA-graph capture safe).""" + + def __init__(self, tensor: torch.Tensor): + self._tensor = tensor + + def __dlpack__(self, stream=None): + return self._tensor.__dlpack__(stream=-1) + + def __dlpack_device__(self): + return self._tensor.__dlpack_device__() + + +def _cute(t: torch.Tensor): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(_GraphSafeDlpack(t), assumed_align=16) + + +_DUMMIES: dict = {} + + +def _dummy(kind: str, device) -> torch.Tensor: + """Placeholder operands with fixed layouts (unused by the compiled variant).""" + key = (kind, device) + if key not in _DUMMIES: + if kind == "a16": + t = torch.zeros((1, 64), dtype=torch.bfloat16, device=device) + elif kind == "a8": + t = torch.zeros((1, 64), dtype=torch.uint8, device=device) + elif kind == "b8": + t = torch.zeros((8, 64), dtype=torch.uint8, device=device) + elif kind == "b16": + t = torch.zeros((8, 64), dtype=torch.bfloat16, device=device) + elif kind == "s": + t = torch.ones((1,), dtype=torch.float32, device=device) + elif kind == "c": + t = torch.zeros((1, 8), dtype=torch.bfloat16, device=device) + elif kind == "cw": + t = torch.zeros((8, 4), dtype=torch.bfloat16, device=device) + elif kind == "cs": + t = torch.zeros((1, 8, 3), dtype=torch.bfloat16, device=device) + elif kind == "slot": + t = torch.zeros((1,), dtype=torch.int32, device=device) + elif kind == "win": + t = torch.zeros((1, 1, 8, 3), dtype=torch.bfloat16, device=device) + else: + raise KeyError(kind) + _DUMMIES[key] = t + return _DUMMIES[key] + + +def _row_view(rows: int, cols: int, stride: int, dtype, device) -> torch.Tensor: + return torch.empty((rows, stride), dtype=dtype, device=device)[:, :cols] + + +@functools.lru_cache(maxsize=None) +def _compiled( + m: int, + k: int, + k_warps: int, + n_fp8: int, + n_tiles: int, + fp8_act: str, + swiglu_fp8: bool, + n_bf16: int, + n_tiles_bf16: int, + swiglu_bf16: bool, + split_out: int, + n_valid: int, + quant_out: bool, + n_out: int, + n_out2: int, + out_dtype: torch.dtype, + a16_stride: int, + a8_stride: int, + use_pdl: bool, + early_trigger: bool, + release_after_wait: bool, + unsafe_skip_wait: bool = False, + ksplit: int = 1, + swap_ab: bool = False, + min_blocks: int = 1, + conv_channels: int = 0, + conv_tokens: int = 1, + conv_write_state: bool = True, + conv_shapes: tuple = (), +): + import cuda.bindings.driver as cuda + import cutlass.cute as cute + + from .lowm_mma_gemm import LowMMmaGemmKernel + + kernel = LowMMmaGemmKernel( + num_rows=m, + k_extent=k, + k_warps=k_warps, + use_pdl=use_pdl, + n_fp8=n_fp8, + n_tiles=n_tiles, + fp8_act=fp8_act, + swiglu_fp8=swiglu_fp8, + n_bf16=n_bf16, + n_tiles_bf16=n_tiles_bf16, + swiglu_bf16=swiglu_bf16, + split_out=split_out, + n_valid=n_valid, + quant_out=quant_out, + early_trigger=early_trigger, + release_after_wait=release_after_wait, + unsafe_skip_wait=unsafe_skip_wait, + ksplit=ksplit, + swap_ab=swap_ab, + min_blocks=min_blocks, + conv_channels=conv_channels, + conv_tokens=conv_tokens, + conv_write_state=conv_write_state, + ) + dev = torch.device("cuda") + if conv_channels: + (cw_shape, cw_stride), (cs_shape, cs_stride), n_req, (win_shape, win_stride) = ( + conv_shapes + ) + cw = torch.empty_strided(cw_shape, cw_stride, dtype=torch.bfloat16, device=dev) + cs = torch.empty_strided(cs_shape, cs_stride, dtype=torch.bfloat16, device=dev) + slots = torch.zeros((n_req,), dtype=torch.int32, device=dev) + win = torch.empty_strided( + win_shape, win_stride, dtype=torch.bfloat16, device=dev + ) + else: + cw, cs, slots, win = ( + _dummy("cw", dev), + _dummy("cs", dev), + _dummy("slot", dev), + _dummy("win", dev), + ) + a16 = ( + _row_view(m, k, a16_stride, torch.bfloat16, dev) + if a16_stride + else _dummy("a16", dev) + ) + a8 = ( + _row_view(m, k, a8_stride, torch.uint8, dev) if a8_stride else _dummy("a8", dev) + ) + b = ( + torch.empty((n_fp8, k), dtype=torch.uint8, device=dev) + if n_fp8 + else _dummy("b8", dev) + ) + b2 = ( + torch.empty((n_bf16, k), dtype=torch.bfloat16, device=dev) + if n_bf16 + else _dummy("b16", dev) + ) + s = _dummy("s", dev) + c = torch.empty( + (m, n_out), dtype=torch.uint8 if quant_out else out_dtype, device=dev + ) + c2 = ( + torch.empty((m, n_out2), dtype=out_dtype, device=dev) + if n_out2 + else _dummy("c", dev) + ) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile( + kernel, + _cute(a16), + _cute(a8), + _cute(b), + _cute(s), + _cute(s), + _cute(s), + _cute(b2), + _cute(c), + _cute(c2), + _cute(cw), + _cute(cs), + _cute(slots), + _cute(win), + stream, + ) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/lowm_mma_gemm.py b/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/lowm_mma_gemm.py new file mode 100644 index 00000000..636bbd45 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/lowm_mma_gemm.py @@ -0,0 +1,1518 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# Vendored from TensorRT-LLM r3 (tekit 7df45abe5c), tensorrt_llm/_torch/cute_dsl_kernels/blackwell/lowm_mma_gemm.py +# (git archive of that commit; md5 9dfebcf3dd06b7c91b33d9f5a5ed4b61). One change: the GDN conv epilogue reads +# its conv-state slot (``gSlot``) after griddepcontrol.wait instead of before it, so only weights, scales and +# conv weights are read ahead of the wait. +"""Warp-MMA GEMM for decode-sized M over per-tensor FP8 (W8A8) and BF16 weights. + +Low-M GEMM with register-prefetched weights for Qwen3.5's per-tensor FP8 (W8A8) projections. +``C[M, :] = A[M, K] @ W.T`` where the weight rows come in two segments, each a whole number of CTAs: + +* FP8 segment ``B[n_fp8, K]`` (E4M3, FP32 per-tensor weight scale ``Sw[0]``). ``fp8_act`` selects + the activation operand: + + - ``"fp8"``: ``A8[M, K]`` E4M3 activations quantized by the producer with the static input scale + ``Sin[0]``; ``mma.sync.m16n8k32`` E4M3 x E4M3 -> FP32, output ``acc * Sin * Sw``. Same numerics + as the activation-quantize kernel + cuBLAS FP8 GEMM up to the FP32 summation order. + - ``"quant"``: ``A16[M, K]`` BF16 activations quantized in registers exactly like + ``scaleMatrixPerTensorVec`` (``bf16 -> fp32``, ``* (1 / Sin)``, ``cvt.rn.satfinite`` to E4M3), + then as ``"fp8"``: no separate activation-quantize kernel, same quantized values. + - ``"bf16"``: ``A16`` BF16 activations times the E4M3 weights converted to BF16 in registers + (exact) on ``mma.sync.m16n8k16`` BF16, output ``acc * Sw`` (no activation quantization). + +* BF16 segment ``B2[n_bf16, K]`` with ``A16`` (for example GDN ``in_proj_ba`` next to the FP8 + ``in_proj_qkvz``, the shared-expert gate next to its FP8 gate_up, or a whole BF16 projection). + +Logical output columns are the FP8 segment's outputs followed by the BF16 segment's; columns below +``split_out`` go to ``C``, the rest to ``C2`` (column ``- split_out``). A segment can end in a SwiGLU: +its weight rows are then stacked as (8 gate rows, the 8 matching up rows) pairs and each pair gives 8 +outputs ``silu(gate) * up`` (FP32). + +Layout of the work: + +* A CTA owns ``n_tiles`` groups of 8 weight rows (one ``m16n8`` tile column each) and splits K over + its ``k_warps`` warps; the warps reduce through shared memory at the end. +* K is processed in 64-element chunks. Lane ``(r, q) = (lane // 4, lane % 4)`` holds the weights + ``W[n0 + r, k + 16q : k + 16q + 16]`` of a chunk (16 bytes of FP8, 32 of BF16) and the activations + of the same 16 columns for rows ``16t + r`` and ``16t + r + 8``. Those registers *are* the MMA + fragments once the reduction order inside a chunk is permuted the same way for A and W: + FP8 (two ``m16n8k32`` steps ``j``): logical k ``4q + i`` / ``4q + 16 + i`` -> physical + ``16q + 8j + i`` / ``16q + 8j + 4 + i``; BF16 (four ``m16n8k16`` steps ``j``): logical k + ``2q + i`` / ``2q + 8 + i`` -> physical ``16q + 4j + i`` / ``16q + 4j + 2 + i``. +* ``swap_ab`` (M <= 8, an even number of tiles per CTA): the weights are the MMA's 16-row A operand (two + 8-row tiles) and the activations its 8-column B operand, so a tile pair takes one MMA per K step + instead of two (half the MMAs and half the accumulator dependency chain). +* ``ksplit`` > 1 (cluster split-K): a cluster of ``ksplit`` CTAs shares a column group, each CTA + takes a ``1/ksplit`` K slice (so it reads and quantizes only that slice of A), and the partial sums + are reduced through distributed shared memory in a fixed rank order. +* All weights and scales of a warp are loaded before ``griddepcontrol.wait`` (programmatic dependent + launch); activations only after it. With ``early_trigger`` the grid releases its dependent right + after issuing the weight loads (the dependent still waits for this grid's completion before it + reads the output). +* ``conv_channels`` > 0 (Qwen3.5 GDN in_proj): the first ``conv_channels`` output columns of + the FP8 segment (the q|k|v channels of ``in_proj_qkvz``) leave through the GDN causal conv1d update of + the MTP target-verify step instead of raw: the rows are ``num_rows / conv_tokens`` requests of + ``conv_tokens`` tokens; per request and channel the window is the conv state (3 columns, slot from + ``state_indices``) followed by the BF16-rounded GEMM outputs, and the kernel writes + ``silu(sum_j w_j * window_{t+j})`` (BF16 products, FP32 sum, ``ex2.approx`` / ``div.full``: the + arithmetic Triton generates for ``causal_conv1d_update``), the intermediate windows after every token + and the conv state after the last one. The conv weights (static) are loaded before the dependency wait; + the slots (``state_indices``) and the conv states (written by earlier kernels: the previous step's + update, the state promotion) only after it, so their freshness does not depend on how the preceding + grid releases its dependents. Their latency hides behind the main loop. +""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float32, Int32, Uint8, Uint32, const_expr +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +from .utils import griddepcontrol_launch_dependents, griddepcontrol_wait + +try: + from cutlass.memory import SmemAllocator as _SmemAllocator +except ImportError: # older CuTe DSL + from cutlass.utils import SmemAllocator as _SmemAllocator + +CHUNK = 64 # K elements per lane quad (16 per lane) +N_TILE = 8 +FP8_ACTS = ("fp8", "quant", "bf16") + + +@dsl_user_op +def _e4m3x4_to_bf16x4(x: Uint32, *, loc=None, ip=None): + """Four packed E4M3 -> four BF16 (exact) as two ``Uint32`` bf16x2 (element 0 in the low half).""" + out = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 2), + [Uint32(x).ir_value(loc=loc, ip=ip)], + "{\n\t" + ".reg .b16 lo, hi, h0, h1, h2, h3;\n\t" + ".reg .b32 p01, p23;\n\t" + ".reg .f32 f0, f1, f2, f3;\n\t" + "mov.b32 {lo, hi}, $2;\n\t" + "cvt.rn.f16x2.e4m3x2 p01, lo;\n\t" + "cvt.rn.f16x2.e4m3x2 p23, hi;\n\t" + "mov.b32 {h0, h1}, p01;\n\t" + "mov.b32 {h2, h3}, p23;\n\t" + "cvt.f32.f16 f0, h0;\n\t" + "cvt.f32.f16 f1, h1;\n\t" + "cvt.f32.f16 f2, h2;\n\t" + "cvt.f32.f16 f3, h3;\n\t" + "cvt.rn.bf16x2.f32 $0, f1, f0;\n\t" + "cvt.rn.bf16x2.f32 $1, f3, f2;\n\t" + "}\n", + "=r,=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Uint32(llvm.extractvalue(T.i32(), out, [0], loc=loc, ip=ip)), + Uint32(llvm.extractvalue(T.i32(), out, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _bf16x4_to_e4m3x4( + x01: Uint32, x23: Uint32, factor: Float32, *, loc=None, ip=None +) -> Uint32: + """Four BF16 (two bf16x2 words, element 0 in the low half) times ``factor`` -> four E4M3 packed with + element 0 in the low byte: FP32 multiply then ``cvt.rn.satfinite``, the arithmetic of the static + per-tensor quantize kernel (``scaleMatrixPerTensorVec``).""" + out = llvm.inline_asm( + T.i32(), + [ + Uint32(x01).ir_value(loc=loc, ip=ip), + Uint32(x23).ir_value(loc=loc, ip=ip), + Float32(factor).ir_value(loc=loc, ip=ip), + ], + "{\n\t" + ".reg .b32 u0, u1, u2, u3;\n\t" + ".reg .f32 f0, f1, f2, f3;\n\t" + ".reg .b16 l, h;\n\t" + "shl.b32 u0, $1, 16;\n\t" + "and.b32 u1, $1, 0xFFFF0000;\n\t" + "shl.b32 u2, $2, 16;\n\t" + "and.b32 u3, $2, 0xFFFF0000;\n\t" + "mov.b32 f0, u0;\n\t" + "mov.b32 f1, u1;\n\t" + "mov.b32 f2, u2;\n\t" + "mov.b32 f3, u3;\n\t" + "mul.rn.f32 f0, f0, $3;\n\t" + "mul.rn.f32 f1, f1, $3;\n\t" + "mul.rn.f32 f2, f2, $3;\n\t" + "mul.rn.f32 f3, f3, $3;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 l, f1, f0;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 h, f3, f2;\n\t" + "mov.b32 $0, {l, h};\n\t" + "}\n", + "=r,r,r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Uint32(out) + + +@dsl_user_op +def _f32_to_e4m3(x: Float32, *, loc=None, ip=None) -> Uint32: + """One FP32 -> E4M3 (``cvt.rn.satfinite``) in the low byte of a ``Uint32``.""" + out = llvm.inline_asm( + T.i32(), + [Float32(x).ir_value(loc=loc, ip=ip)], + "{\n\t" + ".reg .b16 h;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 h, 0f00000000, $1;\n\t" + "cvt.u32.u16 $0, h;\n\t" + "}\n", + "=r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Uint32(out) + + +@dsl_user_op +def _rcp_rn(x: Float32, *, loc=None, ip=None) -> Float32: + """IEEE ``1 / x`` (``div.rn.f32``), as ``1.0f / s`` in the static quantize kernel.""" + out = llvm.inline_asm( + T.f32(), + [Float32(x).ir_value(loc=loc, ip=ip)], + "div.rn.f32 $0, 0f3F800000, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Float32(out) + + +@dsl_user_op +def _gdn_conv_silu(e0, e1, e2, e3, w0, w1, w2, w3, *, loc=None, ip=None) -> Float32: + """One GDN causal-conv1d output: ``silu(((0 + w0 e0) + w1 e1) + w2 e2) + w3 e3)`` rounded to BF16 (returned + as the exact FP32 value). Inputs are FP32 holding BF16 values. Matches the SASS of Triton's + ``_causal_conv1d_update_kernel``: BF16 products (round to nearest), FP32 sum from 0, + ``exp(-acc)`` = ``ex2.approx.f32(-acc * log2(e))``, ``acc / (1 + e)`` = ``div.full.f32``, BF16 round.""" + out = llvm.inline_asm( + T.f32(), + [Float32(v).ir_value(loc=loc, ip=ip) for v in (e0, e1, e2, e3, w0, w1, w2, w3)], + "{\n\t" + ".reg .f32 p, acc, t, e, d, y;\n\t" + ".reg .b16 b, z;\n\t" + "mov.b16 z, 0;\n\t" + "mul.rn.f32 p, $1, $5;\n\t" + "cvt.rn.bf16.f32 b, p;\n\t" + "mov.b32 p, {z, b};\n\t" + "add.rn.f32 acc, 0f00000000, p;\n\t" + "mul.rn.f32 p, $2, $6;\n\t" + "cvt.rn.bf16.f32 b, p;\n\t" + "mov.b32 p, {z, b};\n\t" + "add.rn.f32 acc, acc, p;\n\t" + "mul.rn.f32 p, $3, $7;\n\t" + "cvt.rn.bf16.f32 b, p;\n\t" + "mov.b32 p, {z, b};\n\t" + "add.rn.f32 acc, acc, p;\n\t" + "mul.rn.f32 p, $4, $8;\n\t" + "cvt.rn.bf16.f32 b, p;\n\t" + "mov.b32 p, {z, b};\n\t" + "add.rn.f32 acc, acc, p;\n\t" + "sub.rn.f32 t, 0f00000000, acc;\n\t" + "mul.rn.f32 t, t, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e, t;\n\t" + "add.rn.f32 d, e, 0f3F800000;\n\t" + "div.full.f32 y, acc, d;\n\t" + "cvt.rn.bf16.f32 b, y;\n\t" + "mov.b32 $0, {z, b};\n\t" + "}\n", + "=f,f,f,f,f,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Float32(out) + + +@dsl_user_op +def _mapa_shared_cluster(smem_ptr, peer_rank, *, loc=None, ip=None) -> Int32: + """Shared-memory address of ``smem_ptr`` in cluster CTA ``peer_rank`` (``mapa.shared::cluster``).""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + Int32(peer_rank).ir_value(loc=loc, ip=ip), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_shared_cluster_f32(addr, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Int32(addr).ir_value(loc=loc, ip=ip)], + "ld.shared::cluster.f32 $0, [$1];", + "=f,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +def _mma_asm(shape_types: str): + @dsl_user_op + def _mma(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + out = llvm.inline_asm( + llvm.StructType.get_literal([T.f32()] * 4), + [Uint32(v).ir_value(loc=loc, ip=ip) for v in (a0, a1, a2, a3, b0, b1)] + + [Float32(v).ir_value(loc=loc, ip=ip) for v in (c0, c1, c2, c3)], + f"mma.sync.aligned.{shape_types} " + "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return tuple( + Float32(llvm.extractvalue(T.f32(), out, [i], loc=loc, ip=ip)) + for i in range(4) + ) + + return _mma + + +_mma_bf16_m16n8k16 = _mma_asm("m16n8k16.row.col.f32.bf16.bf16.f32") +_mma_e4m3_m16n8k32 = _mma_asm("m16n8k32.row.col.f32.e4m3.e4m3.f32") + + +def _g2r_atom(dtype, streaming: bool): + """128-bit global -> register copy atom (streaming cache hint for weights when this DSL has it).""" + modes = getattr(cute.nvgpu, "LoadCacheMode", None) + mode = ( + getattr(modes, "STREAMING" if streaming else "ALWAYS", None) if modes else None + ) + if mode is not None: + try: + return cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + dtype, + num_bits_per_copy=128, + load_cache_mode=mode, + ) + except TypeError: + pass + return cute.make_copy_atom(cute.nvgpu.CopyG2ROp(), dtype, num_bits_per_copy=128) + + +class LowMMmaGemmKernel: + """``num_rows`` (<= 64) x K GEMM over an FP8 and a BF16 weight segment (see the module doc).""" + + def __init__( + self, + *, + num_rows: int, + k_extent: int, + k_warps: int, + use_pdl: bool, + n_fp8: int = 0, + n_tiles: int = 1, + fp8_act: str = "fp8", + swiglu_fp8: bool = False, + n_bf16: int = 0, + n_tiles_bf16: int = 1, + swiglu_bf16: bool = False, + split_out: int = 0, + n_valid: int = 0, + quant_out: bool = False, + early_trigger: bool = False, + release_after_wait: bool = False, + unsafe_skip_wait: bool = False, + ksplit: int = 1, + swap_ab: bool = False, + min_blocks: int = 1, + conv_channels: int = 0, + conv_tokens: int = 1, + conv_write_state: bool = True, + ) -> None: + """``n_fp8`` / ``n_bf16``: weight rows of the FP8 / BF16 segment (multiples of ``8 * n_tiles`` + / ``8 * n_tiles_bf16``; either may be 0). ``fp8_act``: activation operand of the FP8 segment + (module doc). ``swiglu_*``: the segment's rows are (gate, up) 8-row pairs and it outputs + ``silu(gate) * up`` (half as many columns; ``n_tiles`` even). ``split_out``: logical output + columns ``>= split_out`` go to ``C2`` (0: everything to ``C``). ``n_valid``: logical columns + ``>= n_valid`` (zero-padded weight rows) are not stored (0: all are). ``quant_out``: the + columns stored to ``C`` are quantized to E4M3 with the static scale ``So[0]`` (``x * (1/So)``, + ``cvt.rn.satfinite``; ``C`` is a uint8 tensor). ``early_trigger``: release the + dependent grid once the weight loads are issued (before this grid's own dependency wait) + instead of after the stores. ``release_after_wait``: release it right after the dependency + wait (overrides ``early_trigger``). ``unsafe_skip_wait``: tests only (the positive control of + the delayed-producer test): no dependency wait, so the activations may be read too early. + ``ksplit``: CTAs per cluster splitting K (1, 2, 4 or 8; module doc). ``swap_ab``: swapped MMA + operands for segments with an even tile count when ``num_rows`` <= 8 (module doc). + ``min_blocks``: launch-bounds occupancy hint (CTAs per SM the register allocation must allow). + ``conv_channels`` / ``conv_tokens`` / ``conv_write_state``: GDN conv1d epilogue on the leading FP8 + columns (module doc).""" + if ksplit not in (1, 2, 4, 8): + raise ValueError(f"ksplit={ksplit} must be 1, 2, 4 or 8") + if k_extent % (CHUNK * k_warps * ksplit): + raise ValueError( + f"K={k_extent} must be a multiple of {CHUNK * k_warps * ksplit}" + ) + if not 1 <= num_rows <= 64: + raise ValueError(f"num_rows={num_rows} must be in [1, 64]") + for tiles in (n_tiles, n_tiles_bf16): + if tiles not in (1, 2, 4): + raise ValueError(f"n_tiles={tiles} must be 1, 2 or 4") + if n_fp8 % (N_TILE * n_tiles) or n_bf16 % (N_TILE * n_tiles_bf16): + raise ValueError( + f"segment rows {n_fp8}/{n_bf16} must be multiples of the CTA tiles " + f"{N_TILE * n_tiles}/{N_TILE * n_tiles_bf16}" + ) + if n_fp8 + n_bf16 == 0: + raise ValueError("empty GEMM") + if fp8_act not in FP8_ACTS: + raise ValueError(f"fp8_act={fp8_act!r} must be one of {FP8_ACTS}") + if (swiglu_fp8 and n_tiles % 2) or (swiglu_bf16 and n_tiles_bf16 % 2): + raise ValueError("a SwiGLU segment needs an even number of tiles per CTA") + self.num_rows = num_rows + self.row_tiles = (num_rows + 15) // 16 + self.k_extent = k_extent + self.k_warps = k_warps + self.ksplit = ksplit + self.swap_ab = swap_ab and num_rows <= 8 + self.min_blocks = min_blocks + self.chunks_per_warp = k_extent // (CHUNK * k_warps * ksplit) + self.n_fp8 = n_fp8 + self.n_bf16 = n_bf16 + self.n_tiles = n_tiles + self.n_tiles_bf16 = n_tiles_bf16 + self.fp8_act = fp8_act + self.swiglu_fp8 = swiglu_fp8 + self.swiglu_bf16 = swiglu_bf16 + self.fp8_ctas = n_fp8 // (N_TILE * n_tiles) + self.bf16_ctas = n_bf16 // (N_TILE * n_tiles_bf16) + self.fp8_out = n_fp8 // 2 if swiglu_fp8 else n_fp8 + self.bf16_out = n_bf16 // 2 if swiglu_bf16 else n_bf16 + self.split_out = split_out + self.n_valid = n_valid if n_valid else self.fp8_out + self.bf16_out + self.quant_out = quant_out + self.use_pdl = use_pdl + self.early_trigger = early_trigger and not release_after_wait + self.release_after_wait = release_after_wait + self.unsafe_skip_wait = unsafe_skip_wait + cols = 0 + if n_fp8: + cols = N_TILE * n_tiles + if n_bf16: + cols = max(cols, N_TILE * n_tiles_bf16) + self.max_cols = cols + self.conv_channels = conv_channels + self.conv_tokens = conv_tokens + self.conv_write_state = conv_write_state + if conv_channels: + fp8_cols = N_TILE * n_tiles + if ( + ksplit != 1 + or swiglu_fp8 + or quant_out + or conv_channels % fp8_cols + or conv_channels > self.fp8_out + or (split_out and conv_channels > split_out) + or num_rows % conv_tokens + or (num_rows // conv_tokens) * fp8_cols > 32 * k_warps + ): + raise ValueError( + f"conv epilogue: channels={conv_channels} tokens={conv_tokens} rows={num_rows} " + f"cols/CTA={fp8_cols} warps={k_warps} ksplit={ksplit} needs whole FP8 CTAs, ksplit 1, no " + "SwiGLU / quantized output, and one (request, channel) pair per thread" + ) + + @cute.jit + def __call__( + self, + gA16: cute.Tensor, + gA8: cute.Tensor, + gB: cute.Tensor, + gSw: cute.Tensor, + gSin: cute.Tensor, + gSo: cute.Tensor, + gB2: cute.Tensor, + gC: cute.Tensor, + gC2: cute.Tensor, + gCw: cute.Tensor, + gCs: cute.Tensor, + gSlot: cute.Tensor, + gWin: cute.Tensor, + stream, + ) -> None: + copy_a16 = _g2r_atom(BFloat16, streaming=False) + copy_a8 = _g2r_atom(Uint8, streaming=False) + copy_b = _g2r_atom(Uint8, streaming=True) + copy_b2 = _g2r_atom(BFloat16, streaming=True) + smem = ( + ( + self.k_warps + + (1 if self.ksplit > 1 else 0) + + (1 if self.conv_channels else 0) + ) + * self.row_tiles + * 16 + * self.max_cols + * 4 + ) + grid = [(self.fp8_ctas + self.bf16_ctas) * self.ksplit, 1, 1] + kernel = self.kernel( + gA16, + gA8, + gB, + gSw, + gSin, + gSo, + gB2, + gC, + gC2, + gCw, + gCs, + gSlot, + gWin, + copy_a16, + copy_a8, + copy_b, + copy_b2, + ) + if const_expr(self.ksplit > 1): + kernel.launch( + grid=grid, + block=[32 * self.k_warps, 1, 1], + cluster=[self.ksplit, 1, 1], + smem=smem, + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=self.min_blocks, + ) + else: + kernel.launch( + grid=grid, + block=[32 * self.k_warps, 1, 1], + smem=smem, + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=self.min_blocks, + ) + + @cute.kernel + def kernel( + self, + gA16: cute.Tensor, + gA8: cute.Tensor, + gB: cute.Tensor, + gSw: cute.Tensor, + gSin: cute.Tensor, + gSo: cute.Tensor, + gB2: cute.Tensor, + gC: cute.Tensor, + gC2: cute.Tensor, + gCw: cute.Tensor, + gCs: cute.Tensor, + gSlot: cute.Tensor, + gWin: cute.Tensor, + copy_a16: cute.CopyAtom, + copy_a8: cute.CopyAtom, + copy_b: cute.CopyAtom, + copy_b2: cute.CopyAtom, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + smem = _SmemAllocator() + red = smem.allocate_tensor( + Float32, + cute.make_layout( + (self.k_warps, self.row_tiles * 16, self.max_cols), + stride=(self.row_tiles * 16 * self.max_cols, self.max_cols, 1), + ), + byte_alignment=16, + ) + part = red + xs = red + if const_expr(self.conv_channels > 0): + # BF16-rounded GEMM outputs of the conv channels (the conv window), FP32 storage. + xs = smem.allocate_tensor( + Float32, + cute.make_layout( + (self.row_tiles * 16, self.max_cols), stride=(self.max_cols, 1) + ), + byte_alignment=16, + ) + group = bidx + kslice = 0 + if const_expr(self.ksplit > 1): + part = smem.allocate_tensor( + Float32, + cute.make_layout( + (self.row_tiles * 16, self.max_cols), stride=(self.max_cols, 1) + ), + byte_alignment=16, + ) + # Clusters are tiles of consecutive block indices: rank = bidx % ksplit. + group = bidx // self.ksplit + kslice = bidx % self.ksplit + if const_expr(self.fp8_ctas > 0 and self.bf16_ctas > 0): + if group < self.fp8_ctas: + self._cta( + gA16, + gA8, + gB, + gSw, + gSin, + gSo, + gB2, + gC, + gC2, + gCw, + gCs, + gSlot, + gWin, + red, + part, + xs, + copy_a16, + copy_a8, + copy_b, + copy_b2, + tidx, + group, + kslice, + True, + ) + else: + self._cta( + gA16, + gA8, + gB, + gSw, + gSin, + gSo, + gB2, + gC, + gC2, + gCw, + gCs, + gSlot, + gWin, + red, + part, + xs, + copy_a16, + copy_a8, + copy_b, + copy_b2, + tidx, + group - self.fp8_ctas, + kslice, + False, + ) + elif const_expr(self.fp8_ctas > 0): + self._cta( + gA16, + gA8, + gB, + gSw, + gSin, + gSo, + gB2, + gC, + gC2, + gCw, + gCs, + gSlot, + gWin, + red, + part, + xs, + copy_a16, + copy_a8, + copy_b, + copy_b2, + tidx, + group, + kslice, + True, + ) + else: + self._cta( + gA16, + gA8, + gB, + gSw, + gSin, + gSo, + gB2, + gC, + gC2, + gCw, + gCs, + gSlot, + gWin, + red, + part, + xs, + copy_a16, + copy_a8, + copy_b, + copy_b2, + tidx, + group, + kslice, + False, + ) + + @cute.jit + def _cta( + self, + gA16: cute.Tensor, + gA8: cute.Tensor, + gB: cute.Tensor, + gSw: cute.Tensor, + gSin: cute.Tensor, + gSo: cute.Tensor, + gB2: cute.Tensor, + gC: cute.Tensor, + gC2: cute.Tensor, + gCw: cute.Tensor, + gCs: cute.Tensor, + gSlot: cute.Tensor, + gWin: cute.Tensor, + red: cute.Tensor, + part: cute.Tensor, + xs: cute.Tensor, + copy_a16: cute.CopyAtom, + copy_a8: cute.CopyAtom, + copy_b: cute.CopyAtom, + copy_b2: cute.CopyAtom, + tidx, + cta, + kslice, + fp8: cutlass.Constexpr, + ) -> None: + warp = tidx // 32 + lane = tidx % 32 + r = lane // 4 + q = lane % 4 + + num_rows: cutlass.Constexpr = self.num_rows + row_tiles: cutlass.Constexpr = self.row_tiles + k_warps: cutlass.Constexpr = self.k_warps + chunks: cutlass.Constexpr = self.chunks_per_warp + n_tiles: cutlass.Constexpr = self.n_tiles if fp8 else self.n_tiles_bf16 + swiglu: cutlass.Constexpr = self.swiglu_fp8 if fp8 else self.swiglu_bf16 + # FP8 weights on the FP8 MMA (activations FP8 from the producer or quantized here). + fp8_mma: cutlass.Constexpr = fp8 and self.fp8_act != "bf16" + cols: cutlass.Constexpr = N_TILE * n_tiles + out_cols: cutlass.Constexpr = cols // 2 if swiglu else cols + out_col0: cutlass.Constexpr = 0 if fp8 else self.fp8_out + + n0 = cta * cols + # This CTA's K slice (cluster split-K), then this warp's part of it. + k_warp0 = (kslice * k_warps + warp) * (chunks * CHUNK) + + # ---- weights and scales of every tile: independent of the producer kernel ---- + w8 = cute.make_rmem_tensor( + cute.make_layout((n_tiles * chunks, 16), stride=(16, 1)), Uint8 + ) + w16 = cute.make_rmem_tensor( + cute.make_layout((2 * n_tiles * chunks, 8), stride=(8, 1)), BFloat16 + ) + alpha = Float32(1.0) + factor = Float32(1.0) + if const_expr(fp8): + gB_vec = cute.logical_divide(gB, (None, 16)) + for nt in cutlass.range_constexpr(n_tiles): + b_row = n0 + N_TILE * nt + r + for c in cutlass.range_constexpr(chunks): + k0 = k_warp0 + c * CHUNK + 16 * q + cute.copy( + copy_b, + gB_vec[b_row, (None, k0 // 16)], + w8[nt * chunks + c, None], + ) + alpha = gSw[0] + if const_expr(fp8_mma): + s_in = gSin[0] + alpha = s_in * alpha + if const_expr(self.fp8_act == "quant"): + factor = _rcp_rn(s_in) + else: + gB2_vec = cute.logical_divide(gB2, (None, 8)) + for nt in cutlass.range_constexpr(n_tiles): + b_row = n0 + N_TILE * nt + r + for c in cutlass.range_constexpr(chunks): + k0 = k_warp0 + c * CHUNK + 16 * q + slot = nt * chunks + c + cute.copy( + copy_b2, gB2_vec[b_row, (None, k0 // 8)], w16[2 * slot, None] + ) + cute.copy( + copy_b2, + gB2_vec[b_row, (None, k0 // 8 + 1)], + w16[2 * slot + 1, None], + ) + + out_factor = Float32(1.0) + if const_expr(self.quant_out): + out_factor = _rcp_rn(gSo[0]) + + # ---- GDN conv epilogue inputs: one (request, channel) pair per thread, indices clamped so every thread + # loads in bounds (CTAs past the conv channels ignore them). Weights (static) before the dependency + # wait; the slots and the conv state after it (module doc) ---- + conv: cutlass.Constexpr = fp8 and self.conv_channels > 0 + conv_t: cutlass.Constexpr = self.conv_tokens + conv_pairs: cutlass.Constexpr = (num_rows // conv_t) * out_cols if conv else 1 + cv_req = Int32(0) + cv_col = Int32(0) + cv_ch = Int32(0) + cv_slot = Int32(0) + cv_s0 = Float32(0.0) + cv_s1 = Float32(0.0) + cv_s2 = Float32(0.0) + cv_w0 = Float32(0.0) + cv_w1 = Float32(0.0) + cv_w2 = Float32(0.0) + cv_w3 = Float32(0.0) + if const_expr(conv): + pair = cutlass.min(tidx, Int32(conv_pairs - 1)) + cv_req = pair // out_cols + cv_col = pair % out_cols + cv_ch = cutlass.min(cta * out_cols + cv_col, Int32(self.conv_channels - 1)) + cv_w0 = gCw[cv_ch, 0].to(Float32) + cv_w1 = gCw[cv_ch, 1].to(Float32) + cv_w2 = gCw[cv_ch, 2].to(Float32) + cv_w3 = gCw[cv_ch, 3].to(Float32) + + if const_expr(self.use_pdl and self.early_trigger): + griddepcontrol_launch_dependents() + if const_expr(self.use_pdl and not self.unsafe_skip_wait): + griddepcontrol_wait() + if const_expr(self.use_pdl and self.release_after_wait): + # Every launch before this one has completed here. Do not move this above the wait. + griddepcontrol_launch_dependents() + if const_expr(conv): + # Slot and conv state after the wait (written by earlier kernels); these loads are in flight during + # the main loop. + cv_slot = gSlot[cv_req] + slot_c = cv_slot - cutlass.min(cv_slot, Int32(0)) # max(slot, 0) + cv_s0 = gCs[slot_c, cv_ch, 0].to(Float32) + cv_s1 = gCs[slot_c, cv_ch, 1].to(Float32) + cv_s2 = gCs[slot_c, cv_ch, 2].to(Float32) + + acc = cute.make_rmem_tensor( + cute.make_layout((n_tiles * row_tiles, 4), stride=(4, 1)), Float32 + ) + acc.fill(0.0) + swap: cutlass.Constexpr = self.swap_ab and n_tiles % 2 == 0 + if const_expr(fp8_mma): + if const_expr(swap): + self._mainloop_fp8_swap( + acc, + gA16, + gA8, + w8, + factor, + copy_a16, + copy_a8, + k_warp0, + r, + q, + n_tiles, + ) + else: + self._mainloop_fp8( + acc, + gA16, + gA8, + w8, + factor, + copy_a16, + copy_a8, + k_warp0, + r, + q, + n_tiles, + ) + else: + if const_expr(swap): + self._mainloop_bf16_swap( + acc, gA16, w8, w16, copy_a16, k_warp0, r, q, n_tiles, fp8 + ) + else: + self._mainloop_bf16( + acc, gA16, w8, w16, copy_a16, k_warp0, r, q, n_tiles, fp8 + ) + + # ---- cross-warp (split-K) reduction through shared memory ---- + if const_expr(swap): + # acc[p] of tile pair p: (act row 2q, weight row 16p + r), (2q + 1, 16p + r), then + 8 rows. + for pair in cutlass.range_constexpr(n_tiles // 2): + col = 2 * N_TILE * pair + r + red[warp, 2 * q, col] = acc[pair, 0] + red[warp, 2 * q + 1, col] = acc[pair, 1] + red[warp, 2 * q, col + N_TILE] = acc[pair, 2] + red[warp, 2 * q + 1, col + N_TILE] = acc[pair, 3] + else: + for nt in cutlass.range_constexpr(n_tiles): + for t in cutlass.range_constexpr(row_tiles): + col = N_TILE * nt + 2 * q + red[warp, 16 * t + r, col] = acc[nt * row_tiles + t, 0] + red[warp, 16 * t + r, col + 1] = acc[nt * row_tiles + t, 1] + red[warp, 16 * t + r + 8, col] = acc[nt * row_tiles + t, 2] + red[warp, 16 * t + r + 8, col + 1] = acc[nt * row_tiles + t, 3] + cute.arch.sync_threads() + if const_expr(conv): + if cta * out_cols < self.conv_channels: + self._conv_epilogue( + gC, + gCs, + gWin, + red, + xs, + tidx, + cta, + alpha, + out_cols, + conv_pairs, + cv_req, + cv_col, + cv_ch, + cv_slot, + cv_s0, + cv_s1, + cv_s2, + cv_w0, + cv_w1, + cv_w2, + cv_w3, + ) + else: + self._epilogue_k1( + gC, + gC2, + red, + tidx, + cta, + alpha, + out_factor, + out_cols, + out_col0, + swiglu, + ) + elif const_expr(self.ksplit == 1): + for idx in cutlass.range(tidx, num_rows * out_cols, 32 * k_warps): + row = idx // out_cols + col = idx % out_cols + value = Float32(0.0) + if const_expr(swiglu): + # Output column 8p + j of the CTA: gate in tile 2p, up in tile 2p + 1. + gcol = (col // N_TILE) * (2 * N_TILE) + col % N_TILE + gate = Float32(0.0) + up = Float32(0.0) + for w in cutlass.range_constexpr(k_warps): + gate = gate + red[w, row, gcol] + for w in cutlass.range_constexpr(k_warps): + up = up + red[w, row, gcol + N_TILE] + value = self._swiglu(gate, up, alpha) + else: + for w in cutlass.range_constexpr(k_warps): + value = value + red[w, row, col] + value = value * alpha + self._finish( + gC, gC2, row, out_col0 + cta * out_cols + col, value, out_factor + ) + else: + # This CTA's partial sums (its warps, in order) ... + for idx in cutlass.range(tidx, num_rows * cols, 32 * k_warps): + row = idx // cols + col = idx % cols + total = Float32(0.0) + for w in cutlass.range_constexpr(k_warps): + total = total + red[w, row, col] + part[row, col] = total + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + # ... summed over the cluster in rank order through distributed shared memory; rank r + # finishes the outputs idx = r (mod ksplit). + for idx in cutlass.range( + kslice + tidx * self.ksplit, + num_rows * out_cols, + 32 * k_warps * self.ksplit, + ): + row = idx // out_cols + col = idx % out_cols + value = Float32(0.0) + if const_expr(swiglu): + gcol = (col // N_TILE) * (2 * N_TILE) + col % N_TILE + gate = self._cluster_sum(part, row, gcol) + up = self._cluster_sum(part, row, gcol + N_TILE) + value = self._swiglu(gate, up, alpha) + else: + value = self._cluster_sum(part, row, col) * alpha + self._finish( + gC, gC2, row, out_col0 + cta * out_cols + col, value, out_factor + ) + # Keep this CTA's shared memory alive until every peer has read it. + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + if const_expr( + self.use_pdl and not self.early_trigger and not self.release_after_wait + ): + griddepcontrol_launch_dependents() + + @cute.jit + def _epilogue_k1( + self, + gC, + gC2, + red, + tidx, + cta, + alpha, + out_factor, + out_cols: cutlass.Constexpr, + out_col0: cutlass.Constexpr, + swiglu: cutlass.Constexpr, + ) -> None: + """The ksplit == 1 epilogue (warp sums in order, alpha, store), for CTAs of a conv-epilogue launch + past the conv channels.""" + num_rows: cutlass.Constexpr = self.num_rows + k_warps: cutlass.Constexpr = self.k_warps + for idx in cutlass.range(tidx, num_rows * out_cols, 32 * k_warps): + row = idx // out_cols + col = idx % out_cols + value = Float32(0.0) + if const_expr(swiglu): + gcol = (col // N_TILE) * (2 * N_TILE) + col % N_TILE + gate = Float32(0.0) + up = Float32(0.0) + for w in cutlass.range_constexpr(k_warps): + gate = gate + red[w, row, gcol] + for w in cutlass.range_constexpr(k_warps): + up = up + red[w, row, gcol + N_TILE] + value = self._swiglu(gate, up, alpha) + else: + for w in cutlass.range_constexpr(k_warps): + value = value + red[w, row, col] + value = value * alpha + self._finish( + gC, gC2, row, out_col0 + cta * out_cols + col, value, out_factor + ) + + @cute.jit + def _conv_epilogue( + self, + gC, + gCs, + gWin, + red, + xs, + tidx, + cta, + alpha, + out_cols: cutlass.Constexpr, + conv_pairs: cutlass.Constexpr, + cv_req, + cv_col, + cv_ch, + cv_slot, + cv_s0, + cv_s1, + cv_s2, + cv_w0, + cv_w1, + cv_w2, + cv_w3, + ) -> None: + """GDN causal conv1d update of this CTA's channels (module doc): BF16 GEMM outputs into shared memory, + then one thread per (request, channel) writes the conv outputs of all tokens, the windows after every + token and the conv state after the last one.""" + num_rows: cutlass.Constexpr = self.num_rows + k_warps: cutlass.Constexpr = self.k_warps + conv_t: cutlass.Constexpr = self.conv_tokens + for idx in cutlass.range(tidx, num_rows * out_cols, 32 * k_warps): + row = idx // out_cols + col = idx % out_cols + value = Float32(0.0) + for w in cutlass.range_constexpr(k_warps): + value = value + red[w, row, col] + value = value * alpha + xs[row, col] = value.to(BFloat16).to(Float32) + cute.arch.sync_threads() + if tidx < conv_pairs: + row0 = cv_req * conv_t + ext = [cv_s0, cv_s1, cv_s2] + for t in cutlass.range_constexpr(conv_t): + ext.append(xs[row0 + t, cv_col]) + for t in cutlass.range_constexpr(conv_t): + y = _gdn_conv_silu( + ext[t], + ext[t + 1], + ext[t + 2], + ext[t + 3], + cv_w0, + cv_w1, + cv_w2, + cv_w3, + ) + gC[row0 + t, cv_ch] = y.to(gC.element_type) + if cv_slot >= 0: + for t in cutlass.range_constexpr(conv_t): + gWin[cv_req, t, cv_ch, 0] = ext[t + 1].to(gWin.element_type) + gWin[cv_req, t, cv_ch, 1] = ext[t + 2].to(gWin.element_type) + gWin[cv_req, t, cv_ch, 2] = ext[t + 3].to(gWin.element_type) + if const_expr(self.conv_write_state): + gCs[cv_slot, cv_ch, 0] = ext[conv_t].to(gCs.element_type) + gCs[cv_slot, cv_ch, 1] = ext[conv_t + 1].to(gCs.element_type) + gCs[cv_slot, cv_ch, 2] = ext[conv_t + 2].to(gCs.element_type) + + @cute.jit + def _mainloop_fp8( + self, + acc: cute.Tensor, + gA16: cute.Tensor, + gA8: cute.Tensor, + w8: cute.Tensor, + factor: Float32, + copy_a16: cute.CopyAtom, + copy_a8: cute.CopyAtom, + k_warp0, + r, + q, + n_tiles: cutlass.Constexpr, + ) -> None: + """E4M3 x E4M3 ``m16n8k32`` main loop (activations FP8, or BF16 quantized here).""" + num_rows: cutlass.Constexpr = self.num_rows + row_tiles: cutlass.Constexpr = self.row_tiles + chunks: cutlass.Constexpr = self.chunks_per_warp + quant: cutlass.Constexpr = self.fp8_act == "quant" + + # FP8 activations of rows 16t + r (lo) and 16t + r + 8 (hi), columns 16q .. 16q + 15. + a_lo = cute.make_rmem_tensor(cute.make_layout((16,)), Uint8) + a_hi = cute.make_rmem_tensor(cute.make_layout((16,)), Uint8) + a_lo_u32 = cute.recast_tensor(a_lo, Uint32) + a_hi_u32 = cute.recast_tensor(a_hi, Uint32) + # BF16 activations (``quant``): columns 16q + 0..7 (0) and 16q + 8..15 (1). + x_lo0 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + x_lo1 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + x_hi0 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + x_hi1 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + x_lo0_u32 = cute.recast_tensor(x_lo0, Uint32) + x_lo1_u32 = cute.recast_tensor(x_lo1, Uint32) + x_hi0_u32 = cute.recast_tensor(x_hi0, Uint32) + x_hi1_u32 = cute.recast_tensor(x_hi1, Uint32) + gA16_vec = cute.logical_divide(gA16, (None, 8)) + gA8_vec = cute.logical_divide(gA8, (None, 16)) + + for c in cutlass.range_constexpr(chunks): + k0 = k_warp0 + c * CHUNK + 16 * q + for t in cutlass.range_constexpr(row_tiles): + has_hi: cutlass.Constexpr = 16 * t + 8 < num_rows + # Rows past num_rows load the last row instead (no divergent branch around the loads and + # the quantization); their outputs are computed but never stored. + row_lo = self._clamp_row(16 * t + r, 16 * t + 8) + row_hi = self._clamp_row(16 * t + r + 8, 16 * t + 16) + if const_expr(not has_hi): + a_hi_u32.fill(0) + if const_expr(quant): + cute.copy(copy_a16, gA16_vec[row_lo, (None, k0 // 8)], x_lo0) + cute.copy(copy_a16, gA16_vec[row_lo, (None, k0 // 8 + 1)], x_lo1) + a_lo_u32[0] = _bf16x4_to_e4m3x4(x_lo0_u32[0], x_lo0_u32[1], factor) + a_lo_u32[1] = _bf16x4_to_e4m3x4(x_lo0_u32[2], x_lo0_u32[3], factor) + a_lo_u32[2] = _bf16x4_to_e4m3x4(x_lo1_u32[0], x_lo1_u32[1], factor) + a_lo_u32[3] = _bf16x4_to_e4m3x4(x_lo1_u32[2], x_lo1_u32[3], factor) + if const_expr(has_hi): + cute.copy(copy_a16, gA16_vec[row_hi, (None, k0 // 8)], x_hi0) + cute.copy( + copy_a16, gA16_vec[row_hi, (None, k0 // 8 + 1)], x_hi1 + ) + a_hi_u32[0] = _bf16x4_to_e4m3x4( + x_hi0_u32[0], x_hi0_u32[1], factor + ) + a_hi_u32[1] = _bf16x4_to_e4m3x4( + x_hi0_u32[2], x_hi0_u32[3], factor + ) + a_hi_u32[2] = _bf16x4_to_e4m3x4( + x_hi1_u32[0], x_hi1_u32[1], factor + ) + a_hi_u32[3] = _bf16x4_to_e4m3x4( + x_hi1_u32[2], x_hi1_u32[3], factor + ) + else: + cute.copy(copy_a8, gA8_vec[row_lo, (None, k0 // 16)], a_lo) + if const_expr(has_hi): + cute.copy(copy_a8, gA8_vec[row_hi, (None, k0 // 16)], a_hi) + for nt in cutlass.range_constexpr(n_tiles): + w = cute.recast_tensor(w8[nt * chunks + c, None], Uint32) + c0 = acc[nt * row_tiles + t, 0] + c1 = acc[nt * row_tiles + t, 1] + c2 = acc[nt * row_tiles + t, 2] + c3 = acc[nt * row_tiles + t, 3] + for j in cutlass.range_constexpr(2): + c0, c1, c2, c3 = _mma_e4m3_m16n8k32( + a_lo_u32[2 * j], + a_hi_u32[2 * j], + a_lo_u32[2 * j + 1], + a_hi_u32[2 * j + 1], + w[2 * j], + w[2 * j + 1], + c0, + c1, + c2, + c3, + ) + acc[nt * row_tiles + t, 0] = c0 + acc[nt * row_tiles + t, 1] = c1 + acc[nt * row_tiles + t, 2] = c2 + acc[nt * row_tiles + t, 3] = c3 + + @cute.jit + def _mainloop_fp8_swap( + self, + acc: cute.Tensor, + gA16: cute.Tensor, + gA8: cute.Tensor, + w8: cute.Tensor, + factor: Float32, + copy_a16: cute.CopyAtom, + copy_a8: cute.CopyAtom, + k_warp0, + r, + q, + n_tiles: cutlass.Constexpr, + ) -> None: + """Swapped E4M3 ``m16n8k32`` main loop (num_rows <= 8): A = the weights of tiles 2p (rows 0..7) and + 2p + 1 (rows 8..15), B = activation rows 0..7 (lane group r holds row r).""" + chunks: cutlass.Constexpr = self.chunks_per_warp + quant: cutlass.Constexpr = self.fp8_act == "quant" + a_row = cute.make_rmem_tensor(cute.make_layout((16,)), Uint8) + a_row_u32 = cute.recast_tensor(a_row, Uint32) + x0 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + x1 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + x0_u32 = cute.recast_tensor(x0, Uint32) + x1_u32 = cute.recast_tensor(x1, Uint32) + gA16_vec = cute.logical_divide(gA16, (None, 8)) + gA8_vec = cute.logical_divide(gA8, (None, 16)) + row = self._clamp_row(r, 8) + for c in cutlass.range_constexpr(chunks): + k0 = k_warp0 + c * CHUNK + 16 * q + if const_expr(quant): + cute.copy(copy_a16, gA16_vec[row, (None, k0 // 8)], x0) + cute.copy(copy_a16, gA16_vec[row, (None, k0 // 8 + 1)], x1) + a_row_u32[0] = _bf16x4_to_e4m3x4(x0_u32[0], x0_u32[1], factor) + a_row_u32[1] = _bf16x4_to_e4m3x4(x0_u32[2], x0_u32[3], factor) + a_row_u32[2] = _bf16x4_to_e4m3x4(x1_u32[0], x1_u32[1], factor) + a_row_u32[3] = _bf16x4_to_e4m3x4(x1_u32[2], x1_u32[3], factor) + else: + cute.copy(copy_a8, gA8_vec[row, (None, k0 // 16)], a_row) + for pair in cutlass.range_constexpr(n_tiles // 2): + w_lo = cute.recast_tensor(w8[(2 * pair) * chunks + c, None], Uint32) + w_hi = cute.recast_tensor(w8[(2 * pair + 1) * chunks + c, None], Uint32) + c0 = acc[pair, 0] + c1 = acc[pair, 1] + c2 = acc[pair, 2] + c3 = acc[pair, 3] + for j in cutlass.range_constexpr(2): + c0, c1, c2, c3 = _mma_e4m3_m16n8k32( + w_lo[2 * j], + w_hi[2 * j], + w_lo[2 * j + 1], + w_hi[2 * j + 1], + a_row_u32[2 * j], + a_row_u32[2 * j + 1], + c0, + c1, + c2, + c3, + ) + acc[pair, 0] = c0 + acc[pair, 1] = c1 + acc[pair, 2] = c2 + acc[pair, 3] = c3 + + @cute.jit + def _mainloop_bf16_swap( + self, + acc: cute.Tensor, + gA16: cute.Tensor, + w8: cute.Tensor, + w16: cute.Tensor, + copy_a16: cute.CopyAtom, + k_warp0, + r, + q, + n_tiles: cutlass.Constexpr, + fp8_weights: cutlass.Constexpr, + ) -> None: + """Swapped BF16 ``m16n8k16`` main loop (num_rows <= 8; see ``_mainloop_fp8_swap``).""" + chunks: cutlass.Constexpr = self.chunks_per_warp + a0 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + a1 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + a0_u32 = cute.recast_tensor(a0, Uint32) + a1_u32 = cute.recast_tensor(a1, Uint32) + w_lo = cute.make_rmem_tensor(cute.make_layout((8,)), Uint32) + w_hi = cute.make_rmem_tensor(cute.make_layout((8,)), Uint32) + gA16_vec = cute.logical_divide(gA16, (None, 8)) + row = self._clamp_row(r, 8) + for c in cutlass.range_constexpr(chunks): + k0 = k_warp0 + c * CHUNK + 16 * q + cute.copy(copy_a16, gA16_vec[row, (None, k0 // 8)], a0) + cute.copy(copy_a16, gA16_vec[row, (None, k0 // 8 + 1)], a1) + for pair in cutlass.range_constexpr(n_tiles // 2): + self._weights_bf16(w_lo, w8, w16, (2 * pair) * chunks + c, fp8_weights) + self._weights_bf16( + w_hi, w8, w16, (2 * pair + 1) * chunks + c, fp8_weights + ) + c0 = acc[pair, 0] + c1 = acc[pair, 1] + c2 = acc[pair, 2] + c3 = acc[pair, 3] + for j in cutlass.range_constexpr(4): + if const_expr(j < 2): + c0, c1, c2, c3 = _mma_bf16_m16n8k16( + w_lo[2 * j], + w_hi[2 * j], + w_lo[2 * j + 1], + w_hi[2 * j + 1], + a0_u32[2 * j], + a0_u32[2 * j + 1], + c0, + c1, + c2, + c3, + ) + else: + c0, c1, c2, c3 = _mma_bf16_m16n8k16( + w_lo[2 * j], + w_hi[2 * j], + w_lo[2 * j + 1], + w_hi[2 * j + 1], + a1_u32[2 * j - 4], + a1_u32[2 * j - 3], + c0, + c1, + c2, + c3, + ) + acc[pair, 0] = c0 + acc[pair, 1] = c1 + acc[pair, 2] = c2 + acc[pair, 3] = c3 + + @cute.jit + def _mainloop_bf16( + self, + acc: cute.Tensor, + gA16: cute.Tensor, + w8: cute.Tensor, + w16: cute.Tensor, + copy_a16: cute.CopyAtom, + k_warp0, + r, + q, + n_tiles: cutlass.Constexpr, + fp8_weights: cutlass.Constexpr, + ) -> None: + """BF16 ``m16n8k16`` main loop: BF16 activations times BF16 weights (or E4M3 weights converted + to BF16 exactly).""" + num_rows: cutlass.Constexpr = self.num_rows + row_tiles: cutlass.Constexpr = self.row_tiles + chunks: cutlass.Constexpr = self.chunks_per_warp + + a_lo0 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + a_lo1 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + a_hi0 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + a_hi1 = cute.make_rmem_tensor(cute.make_layout((8,)), BFloat16) + a_lo0_u32 = cute.recast_tensor(a_lo0, Uint32) + a_lo1_u32 = cute.recast_tensor(a_lo1, Uint32) + a_hi0_u32 = cute.recast_tensor(a_hi0, Uint32) + a_hi1_u32 = cute.recast_tensor(a_hi1, Uint32) + w_bf16 = cute.make_rmem_tensor(cute.make_layout((8,)), Uint32) + gA16_vec = cute.logical_divide(gA16, (None, 8)) + + for c in cutlass.range_constexpr(chunks): + k0 = k_warp0 + c * CHUNK + 16 * q + if const_expr(n_tiles == 1): + self._weights_bf16(w_bf16, w8, w16, c, fp8_weights) + for t in cutlass.range_constexpr(row_tiles): + has_hi: cutlass.Constexpr = 16 * t + 8 < num_rows + # Rows past num_rows load the last row instead (no divergent branch); never stored. + row_lo = self._clamp_row(16 * t + r, 16 * t + 8) + row_hi = self._clamp_row(16 * t + r + 8, 16 * t + 16) + cute.copy(copy_a16, gA16_vec[row_lo, (None, k0 // 8)], a_lo0) + cute.copy(copy_a16, gA16_vec[row_lo, (None, k0 // 8 + 1)], a_lo1) + if const_expr(has_hi): + cute.copy(copy_a16, gA16_vec[row_hi, (None, k0 // 8)], a_hi0) + cute.copy(copy_a16, gA16_vec[row_hi, (None, k0 // 8 + 1)], a_hi1) + else: + a_hi0_u32.fill(0) + a_hi1_u32.fill(0) + for nt in cutlass.range_constexpr(n_tiles): + if const_expr(n_tiles > 1): + # Per row tile so that one tile's BF16 weights are live at a time. + self._weights_bf16( + w_bf16, w8, w16, nt * chunks + c, fp8_weights + ) + c0 = acc[nt * row_tiles + t, 0] + c1 = acc[nt * row_tiles + t, 1] + c2 = acc[nt * row_tiles + t, 2] + c3 = acc[nt * row_tiles + t, 3] + for j in cutlass.range_constexpr(4): + # A: a0 = (row_lo, 4j..4j+1), a1 = (row_hi, same), a2/a3 = 4j+2..4j+3. + if const_expr(j < 2): + c0, c1, c2, c3 = _mma_bf16_m16n8k16( + a_lo0_u32[2 * j], + a_hi0_u32[2 * j], + a_lo0_u32[2 * j + 1], + a_hi0_u32[2 * j + 1], + w_bf16[2 * j], + w_bf16[2 * j + 1], + c0, + c1, + c2, + c3, + ) + else: + c0, c1, c2, c3 = _mma_bf16_m16n8k16( + a_lo1_u32[2 * j - 4], + a_hi1_u32[2 * j - 4], + a_lo1_u32[2 * j - 3], + a_hi1_u32[2 * j - 3], + w_bf16[2 * j], + w_bf16[2 * j + 1], + c0, + c1, + c2, + c3, + ) + acc[nt * row_tiles + t, 0] = c0 + acc[nt * row_tiles + t, 1] = c1 + acc[nt * row_tiles + t, 2] = c2 + acc[nt * row_tiles + t, 3] = c3 + + @cute.jit + def _clamp_row(self, row, rows_needed: cutlass.Constexpr): + """``row`` when every row the lane quad group touches (< ``rows_needed``) exists, else + ``min(row, num_rows - 1)``.""" + result = row + if const_expr(rows_needed > self.num_rows): + result = cutlass.min(row, Int32(self.num_rows - 1)) + return result + + @cute.jit + def _weights_bf16( + self, + w_bf16: cute.Tensor, + w8: cute.Tensor, + w16: cute.Tensor, + slot: cutlass.Constexpr, + fp8_weights: cutlass.Constexpr, + ) -> None: + """The 16 weights of register slot ``slot`` as 8 bf16x2 words (E4M3 converted exactly).""" + if const_expr(fp8_weights): + raw = cute.recast_tensor(w8[slot, None], Uint32) + for i in cutlass.range_constexpr(4): + lo, hi = _e4m3x4_to_bf16x4(raw[i]) + w_bf16[2 * i] = lo + w_bf16[2 * i + 1] = hi + else: + first = cute.recast_tensor(w16[2 * slot, None], Uint32) + second = cute.recast_tensor(w16[2 * slot + 1, None], Uint32) + for i in cutlass.range_constexpr(4): + w_bf16[i] = first[i] + w_bf16[4 + i] = second[i] + + @cute.jit + def _swiglu(self, gate: Float32, up: Float32, alpha: Float32) -> Float32: + """``silu(gate) * up`` on the scaled sums, each rounded to BF16 first like the unfused path's + GEMM output before its SwiGLU kernel.""" + gate = (gate * alpha).to(BFloat16).to(Float32) + up = (up * alpha).to(BFloat16).to(Float32) + return gate / (Float32(1.0) + cute.math.exp(-gate, fastmath=False)) * up + + @cute.jit + def _cluster_sum(self, part: cute.Tensor, row, col) -> Float32: + """``sum_p part_p[row, col]`` over the cluster's CTAs p = 0 .. ksplit - 1, in that order.""" + offset = cute.crd2idx((row, col), part.layout) * 4 + total = Float32(0.0) + for p in cutlass.range_constexpr(self.ksplit): + total = total + _ld_shared_cluster_f32( + _mapa_shared_cluster(part.iterator, p) + offset + ) + return total + + @cute.jit + def _finish( + self, + gC: cute.Tensor, + gC2: cute.Tensor, + row, + out_col, + value: Float32, + out_factor: Float32, + ) -> None: + if const_expr(self.n_valid < self.fp8_out + self.bf16_out): + if out_col < self.n_valid: + self._store(gC, gC2, row, out_col, value, out_factor) + else: + self._store(gC, gC2, row, out_col, value, out_factor) + + @cute.jit + def _store( + self, + gC: cute.Tensor, + gC2: cute.Tensor, + row, + n, + value: Float32, + out_factor: Float32, + ) -> None: + if const_expr(self.split_out == 0): + self._put_c(gC, row, n, value, out_factor) + else: + if n < self.split_out: + self._put_c(gC, row, n, value, out_factor) + else: + gC2[row, n - self.split_out] = value.to(gC2.element_type) + + @cute.jit + def _put_c( + self, gC: cute.Tensor, row, n, value: Float32, out_factor: Float32 + ) -> None: + if const_expr(self.quant_out): + gC[row, n] = _f32_to_e4m3(value * out_factor).to(Uint8) + else: + gC[row, n] = value.to(gC.element_type) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/utils.py b/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/utils.py new file mode 100644 index 00000000..0dc2b6f2 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/trt_lowm/utils.py @@ -0,0 +1,483 @@ +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# Vendored from TensorRT-LLM r3 (tekit 7df45abe5c), tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py +# (git archive of that commit; md5 0525f46fdf54cb190250815a1a5260e7). +# The unused process-global PDL environment switch is omitted; callers pass policy. +# +# 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. + +# Copyright (c) 2025 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. + +# This file is copied and modified from cutlass https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/core.py + +import ctypes +from typing import Union + +import cutlass +import cutlass._mlir.dialects.cute as _cute_ir +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, nvvm +from cutlass.cute.typing import AddressSpace, Numeric, Pointer, Type +from cutlass.cutlass_dsl import T, dsl_user_op + + +# WAR for CuTeDSL make_ptr implementation +class _Pointer(Pointer): + """Represents a runtime pointer that can interoperate with various data structures, + including numpy arrays and device memory. + + Args: + pointer (int or pointer-like object): The pointer to the data. + dtype (Type): Data type of the elements pointed to. + mem_space (_cute_ir.AddressSpace, optional): Memory space where the pointer resides. Defaults to generic. + assumed_align (int, optional): Alignment of the input pointer in bytes. Defaults to None. + + Attributes: + _pointer: The underlying pointer. + _dtype: Data type of the elements. + _addr_space: Memory space of the pointer. + _assumed_align: Alignment of the pointer in bytes. + _desc: C-type descriptor for the pointer. + _c_pointer: C-compatible pointer representation. + """ + + def __init__( + self, + pointer, + dtype, + mem_space: _cute_ir.AddressSpace = _cute_ir.AddressSpace.generic, + assumed_align=None, + ): + self._pointer = pointer + self._dtype = dtype + self._addr_space = mem_space + + if assumed_align is None: + self._assumed_align = dtype.width // 8 + else: + self._assumed_align = assumed_align + + self._desc = None + self._c_pointer = None + assert int(self._pointer) % self._assumed_align == 0, ( + f"pointer must be {self._assumed_align} bytes aligned" + ) + + def size_in_bytes(self) -> int: + return ctypes.sizeof(ctypes.c_void_p(int(self._pointer))) + + def __get_mlir_types__(self): + return [self.mlir_type] + + def __c_pointers__(self): + if self._c_pointer is None: + self._desc = ctypes.c_void_p(int(self._pointer)) + self._c_pointer = ctypes.addressof(self._desc) + return [self._c_pointer] + + def __new_from_mlir_values__(self, values): + assert len(values) == 1 + return values[0] + + # Move mlir Type out of __init__ to decouple with mlir Context + @property + def mlir_type(self) -> ir.Type: + return _cute_ir.PtrType.get( + self._dtype.mlir_type, self._addr_space, self._assumed_align + ) + + @property + def dtype(self) -> Type[Numeric]: + return self._dtype + + @property + def memspace(self): + return self._addr_space + + def align(self, min_align: int, *, loc=None, ip=None) -> Pointer: + raise NotImplementedError("align is not supported in runtime") + + def verify(self, expected_py_type): + if expected_py_type is Pointer or ( + isinstance(expected_py_type, ir.Value) and expected_py_type.ty is Pointer + ): + return True + + return False + + def __str__(self) -> str: + return f"Ptr<0x{int(self._pointer):016x}@{self._addr_space}>" + + def __repr__(self): + return self.__str__() + + +def make_ptr( + dtype: Type[Numeric], + value: Union[int, ctypes._Pointer], + mem_space: AddressSpace = AddressSpace.generic, + assumed_align=None, +) -> Pointer: + """Creates a pointer from a memory address. + + Args: + dtype (Type[Numeric]): Data type of the pointer elements. + value (Union[int, ctypes._Pointer]): Memory address as an integer or ctypes pointer. + mem_space (AddressSpace, optional): Memory address space. Defaults to AddressSpace.generic. + assumed_align (int, optional): Alignment in bytes. Defaults to None. + + Returns: + Pointer: A pointer object. + + Example: + ```python + import numpy as np + import ctypes + from cutlass import Float32 + from cutlass.cute.runtime import make_ptr + + # Create a numpy array + a = np.random.randn(16, 32).astype(np.float32) + # Get pointer address as ctypes pointer + ptr_address = a.ctypes.data_as(ctypes.POINTER(ctypes.c_float)) + # Create pointer from address + y = make_ptr(cutlass.Float32, ptr_address) + ``` + """ + # check if value is int or ctypes.POINTER + if isinstance(value, int): + address_value = value + elif isinstance(value, ctypes._Pointer): + # get address value + address_value = ctypes.cast(value, ctypes.c_void_p).value + assert address_value is not None, "Pointer address is None" + else: + raise TypeError( + f"Expect int or ctypes.POINTER for value but got {type(value)=}" + ) + + return _Pointer(address_value, dtype, mem_space, assumed_align=assumed_align) + + +def is_power_of_2(x: int) -> bool: + return x > 0 and (x & (x - 1)) == 0 + + +@dsl_user_op +def fmin( + a: Union[float, cutlass.Float32], + b: Union[float, cutlass.Float32], + *, + nan=False, + loc=None, + ip=None, +) -> cutlass.Float32: + return cutlass.Float32( + nvvm.fmin( + cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(b).ir_value(loc=loc, ip=ip), + nan=nan, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def fclip_xorsign( + a: Union[float, cutlass.Float32], + limit: Union[float, cutlass.Float32], + *, + loc=None, + ip=None, +) -> cutlass.Float32: + """Symmetric clip. + + Emits PTX ``min.xorsign.abs.f32`` with semantics + ``d = sign(a XOR limit) * min(|a|, |limit|)``. + When ``limit > 0`` this equals ``clip(a, -limit, +limit)``. + + Requires sm_86+. + """ + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [ + cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(limit).ir_value(loc=loc, ip=ip), + ], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +def sigmoid_f32( + a: Union[float, cutlass.Float32], fastmath: bool = False +) -> Union[float, cutlass.Float32]: + """ + Compute the sigmoid of the input tensor. + """ + return cute.arch.rcp_approx(1.0 + cute.math.exp(-a, fastmath=fastmath)) + + +def silu_f32( + a: Union[float, cutlass.Float32], fastmath: bool = False +) -> Union[float, cutlass.Float32]: + """ + Compute the silu of the input tensor. + """ + return a * sigmoid_f32(a, fastmath=fastmath) + + +def gelu_tanh_f32( + a: Union[float, cutlass.Float32], fastmath: bool = False +) -> Union[float, cutlass.Float32]: + """ + Compute the tanh approximation of GELU (matches F.gelu(approximate="tanh")). + + gelu(a) = 0.5 * a * (1 + tanh(c * (a + 0.044715 * a^3))), c = sqrt(2/pi) + = a * sigmoid(2c * (a + 0.044715 * a^3)) + + using the identity tanh(z) = 2 * sigmoid(2z) - 1, so the activation reuses + the same exp2-based sigmoid as silu_f32. + """ + c2 = 1.5957691216057308 # 2 * sqrt(2/pi) + inner = c2 * (a + 0.044715 * a * a * a) + return a * sigmoid_f32(inner, fastmath=fastmath) + + +# TODO(zhichenj): try to move these to NVVM wrapper or helper functions +@dsl_user_op +def vectorized_atomic_add_bf16x8( + rOut_epi_packed, scatter_out_offset, loc=None, ip=None +): + llvm.inline_asm( + None, + [ + scatter_out_offset.iterator.llvm_ptr, + llvm.bitcast(T.i32(), rOut_epi_packed[0, None].load().ir_value()), + llvm.bitcast(T.i32(), rOut_epi_packed[1, None].load().ir_value()), + llvm.bitcast(T.i32(), rOut_epi_packed[2, None].load().ir_value()), + llvm.bitcast(T.i32(), rOut_epi_packed[3, None].load().ir_value()), + ], + "red.global.v4.bf16x2.add.noftz [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def vectorized_atomic_add_fp16x8( + rOut_epi_packed, scatter_out_offset, loc=None, ip=None +): + llvm.inline_asm( + None, + [ + scatter_out_offset.iterator.llvm_ptr, + llvm.bitcast(T.i32(), rOut_epi_packed[0, None].load().ir_value()), + llvm.bitcast(T.i32(), rOut_epi_packed[1, None].load().ir_value()), + llvm.bitcast(T.i32(), rOut_epi_packed[2, None].load().ir_value()), + llvm.bitcast(T.i32(), rOut_epi_packed[3, None].load().ir_value()), + ], + "red.global.v4.f16x2.add.noftz [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def vectorized_atomic_add_fp32x2( + rOut_epi_packed, scatter_out_offset, loc=None, ip=None +): + llvm.inline_asm( + None, + [ + scatter_out_offset.iterator.llvm_ptr, + rOut_epi_packed[0].ir_value(), + rOut_epi_packed[1].ir_value(), + ], + "red.global.v2.f32.add [$0], {$1, $2};", + "l,f,f", + has_side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def atomic_add_func(rOut_epi_packed, scatter_out_offset, loc=None, ip=None): + if cutlass.const_expr(rOut_epi_packed.dtype == cutlass.Float32): + llvm.inline_asm( + None, + [ + scatter_out_offset.iterator.llvm_ptr, + rOut_epi_packed.ir_value(), + ], + "red.global.add.f32 [$0], $1;", + "l,f", + has_side_effects=True, + loc=loc, + ip=ip, + ) + elif cutlass.const_expr(rOut_epi_packed.dtype == cutlass.BFloat16): + llvm.inline_asm( + None, + [ + scatter_out_offset.iterator.llvm_ptr, + llvm.bitcast(T.i16(), rOut_epi_packed.ir_value()), + ], + "red.add.noftz.bf16 [$0], $1;", + "l,h", + has_side_effects=True, + loc=loc, + ip=ip, + ) + elif cutlass.const_expr(rOut_epi_packed.dtype == cutlass.Float16): + llvm.inline_asm( + None, + [ + scatter_out_offset.iterator.llvm_ptr, + llvm.bitcast(T.i16(), rOut_epi_packed.ir_value()), + ], + "red.add.noftz.f16 [$0], $1;", + "l,h", + has_side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def blk_reduce_bf16(dst_gemm, src_smem, size, loc=None, ip=None): + llvm.inline_asm( + None, + [ + dst_gemm.iterator.llvm_ptr, + src_smem.iterator.llvm_ptr, + size.ir_value(), + ], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.noftz.bf16 [$0], [$1], $2;", + "l,l,r", + has_side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def blk_reduce_fp32(dst_gemm, src_smem, size, loc=None, ip=None): + llvm.inline_asm( + None, + [ + dst_gemm.iterator.llvm_ptr, + src_smem.iterator.llvm_ptr, + size.ir_value(), + ], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [$0], [$1], $2;", + "l,l,r", + has_side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def blk_reduce_fp16(dst_gemm, src_smem, size, loc=None, ip=None): + llvm.inline_asm( + None, + [ + dst_gemm.iterator.llvm_ptr, + src_smem.iterator.llvm_ptr, + size.ir_value(), + ], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.noftz.f16 [$0], [$1], $2;", + "l,l,r", + has_side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def griddepcontrol_wait(*, loc=None, ip=None) -> None: + """ + This instruction is used to wait for the previous kernel's grid ending + (all blocks of the previous kernel have finished and memflushed), i.e., + the instruction after this instruction will not be issued until the previous + grid has finished. + """ + llvm.inline_asm( + res=None, + operands_=[], + asm_string="griddepcontrol.wait;", + constraints="", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def griddepcontrol_launch_dependents(*, loc=None, ip=None) -> None: + """ + Issuing the launch_dependents instruction hints a dependent kernel to launch earlier. + launch_dependents doesn't impact the functionality but the performance: + Launching a dependent kernel too early can compete with current kernels, + while launching too late can lead to a long latency. + """ + llvm.inline_asm( + res=None, + operands_=[], + asm_string="griddepcontrol.launch_dependents;", + constraints="", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index 27b75d17..9b805499 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -124,10 +124,22 @@ "gdn_front.in_proj_ba.rel_l2_vs_stock": 1.05, "gdn_front.in_proj_ba.max_abs_vs_stock": 2.0, }, + "gdn_trt_front": { + "gdn_trt_front.in_proj_qkvz.rel_l2_vs_stock": 1.05, + "gdn_trt_front.in_proj_qkvz.max_abs_vs_stock": 2.0, + "gdn_trt_front.in_proj_ba.rel_l2_vs_stock": 1.05, + "gdn_trt_front.in_proj_ba.max_abs_vs_stock": 2.0, + }, "route_shared": { "route_shared.gate_up.rel_l2_vs_stock": 1.05, "route_shared.gate_up.max_abs_vs_stock": 2.0, }, + # The hidden, norm, FP8 and NVFP4 paths are bitwise; only the router partial's + # FP32 reduction order changes. Its RMS error stayed within 1.05x stock against + # FP64 on the same quantized operands (the TP8 GDN and attention tests). + "route_lin": { + "route_lin.router_logits.rms_vs_fp64_relative_to_stock": 1.05, + }, # SKV2 against SGLang's TRT-LLM MHA decode (test_skv2_is_within_twice_sglangs_error) "skv2": {"skv2.attention.rel_l2_vs_stock": 2.0}, # CP-lite's exact path (SGLang's q35_cp_lite) at one request, as skv2's bound diff --git a/test/kernels/test_qwen3_5_gdn.py b/test/kernels/test_qwen3_5_gdn.py index af274064..6fd12b51 100644 --- a/test/kernels/test_qwen3_5_gdn.py +++ b/test/kernels/test_qwen3_5_gdn.py @@ -24,6 +24,14 @@ on its own projected rows; a padded request leaves the state and windows alone; graph replays equal eager launches; - its gated norm writing static FP8, bit for bit slot gdn_verify's. + +And slot gdn_trt_front (TensorRT-LLM r3's own in_proj front, numeric) at the +TP8 shapes, on the hop's E4M3 rows and on BF16 rows: its q / k / v rows (read back from the windows it +writes) and z against fp64 on the same operands, at most 1.05 x the mean relative L2 and 2 x the max abs error of +SGLang's ops, b / a likewise; its conv output, rolling state and windows bit for bit ``causal_conv1d_update``'s on +those rows; a padded request leaves the pools alone; it reads its rows and state index only after its PDL wait (a +producer that releases it at entry and writes them late; the wait-less build as the positive control); graph +replays equal eager launches, and a launch that was not prepared is refused inside a capture. """ import pytest @@ -458,6 +466,33 @@ def rand(*shape, scale=1.0): ) assert torch.equal(y_cap.view(torch.uint8), y_got.view(torch.uint8)) assert all(torch.equal(c, g) for c, g in zip(capped, got_ring)) + # the committed history loaded before the PDL wait (early_hist) leaves every byte as it was + early = [t.clone() for t in ring] + y_early = torch.empty_like(y_got) + core( + qkv, + z, + a, + b, + A_log, + dt_bias, + state, + *early, + states, + replay, + write_pos, + cache_base, + is_flush, + weight, + scale, + 1e-6, + y_early, + rows_per_request=DRAFT, + num_k_heads=H, + early_hist=True, + ) + assert torch.equal(y_early.view(torch.uint8), y_got.view(torch.uint8)) + assert all(torch.equal(e, g) for e, g in zip(early, got_ring)) live = torch.ones(tokens, dtype=torch.bool, device="cuda") if pad: live[-DRAFT:] = False @@ -479,9 +514,98 @@ def rand(*shape, scale=1.0): assert torch.allclose(got_ring[2], want_ring[2], rtol=0, atol=1e-5) +@pytest.mark.parametrize( + "variant", [{}, {"early_dependents": True}, {"min_blocks": 3}, {"early_hist": True}] +) +@pytest.mark.parametrize("bs", [1, 2, 4, 8, 64]) +@pytest.mark.parametrize("low", [False, True]) +def test_core_flat_stores_leave_every_byte_as_it_was(low, bs, variant): + """Flat stores (predicated, with update vectors stored after the output) write + the FP8 output and this step's ring entries, with and without the low parts, byte for byte as the + branchy stores do at TP8, through the ring's end and with a padded request.""" + from sgl_nvidia_mega.qwen3_5.sm103.gdn.core import core + + h, hv = H, HV + d = 2 * h * K + hv * V + g = torch.Generator(device="cuda").manual_seed(8000 + 10 * bs + low) + ring_len, slots, tokens = 16, bs + 3, bs * DRAFT + + def rand(*shape, scale=1.0): + return (torch.randn(*shape, device="cuda", generator=g) * scale).to( + torch.bfloat16 + ) + + qkvz = rand(tokens, d + hv * V) + qkv, z = qkvz[:, :d], qkvz[:, d:].view(tokens, hv, V) + ba = rand(tokens, 2 * hv) + b, a = ba[:, :hv], ba[:, hv:] + A_log = torch.rand(hv, device="cuda", generator=g) + dt_bias = torch.rand(hv, device="cuda", generator=g) + state = rand(slots, hv, V, K, scale=0.05) + ring = [ + rand(slots, hv, ring_len, V, scale=0.05), + rand(slots, h, ring_len, K, scale=0.09), + -torch.rand(slots, hv, ring_len, device="cuda", generator=g) * 0.5, + rand(slots, hv, ring_len, V, scale=1e-4) if low else None, + rand(slots, h, ring_len, K, scale=1e-4) if low else None, + ] + states = torch.randperm(slots, device="cuda", generator=g)[:bs].to(torch.int32) + replay = torch.randperm(slots, device="cuda", generator=g)[:bs] + if bs > 1: + states[-1] = -1 + write_pos = torch.randint(0, ring_len, (slots,), device="cuda", generator=g).to( + torch.int32 + ) + cache_base = torch.randint(0, ring_len, (slots,), device="cuda", generator=g).to( + torch.int32 + ) + is_flush = torch.zeros(slots, dtype=torch.int8, device="cuda") + weight = rand(V) + scale = torch.tensor([0.02], device="cuda") + got = [] + for flat in (False, True): + r = [None if t is None else t.clone() for t in ring] + y = torch.zeros(tokens, hv * V, dtype=torch.float8_e4m3fn, device="cuda") + core( + qkv, + z, + a, + b, + A_log, + dt_bias, + state, + *r, + states, + replay, + write_pos, + cache_base, + is_flush, + weight, + scale, + 1e-6, + y, + rows_per_request=DRAFT, + num_k_heads=h, + flat=flat, + **variant, + ) + got.append((y.view(torch.uint8), r)) + (y0, r0), (y1, r1) = got + assert torch.equal(y0, y1) + assert all(u is None and v is None or torch.equal(u, v) for u, v in zip(r0, r1)) + assert not torch.equal(r1[0], ring[0]) # the step wrote ring entries + + # ------------------------------------------------------------------- slot gdn_front -FRONT_SLOTS = 5 +FRONT_SLOTS = 12 +# The independent rolling-state, intermediate-window and ring slots used by up +# to eight requests. The TRT front below uses the rolling-state tuple. +FRONT_STATES, FRONT_WINDOWS, FRONT_RINGS = ( + (3, 0, 4, 9, 6, 11, 1, 8), + (1, 5, 0, 2, 10, 3, 7, 4), + (2, 6, 9, 0, 11, 5, 3, 10), +) def _front_operands(seed, *, pad=False, index_dtype=torch.int32): @@ -650,6 +774,47 @@ def test_front_takes_its_weights_by_bulk_copy_with_the_same_bytes(pad): assert torch.equal(t, got[1][n]), n +@pytest.mark.parametrize("pad", [False, True]) +def test_front_holds_its_register_weights_with_the_same_bytes(pad): + """The pinned launch (register weight loads held ahead of the PDL wait) writes every + output and pool as the default bulk-copy launch, eager and from a CUDA graph.""" + from sgl_nvidia_mega.qwen3_5.sm103.gdn import front + + for seed in range(3): + o = _front_operands(30 + seed, pad=pad) + packed = front.pack(o["w"]) + start = {n: o[n].clone() for n in ("state", "phys")} + got = [] + for kw in ({"pin": False}, {"pin": True}): + for n, t in start.items(): + o[n].copy_(t) + outs = _front_outputs() + for t in outs.values(): # a padded request leaves the conv rows unwritten + t.zero_() + _run_front(o, packed, outs, **kw) + torch.cuda.synchronize() + assert outs["error"].item() == 0 + got.append({**outs, "state": o["state"].clone(), "phys": o["phys"].clone()}) + for n, t in got[0].items(): + assert torch.equal(t, got[1][n]), n + # graph replay of the pinned launch + for n, t in start.items(): + o[n].copy_(t) + outs = _front_outputs() + for t in outs.values(): + t.zero_() + _run_front(o, packed, outs, pin=True) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run_front(o, packed, outs, pin=True) + for n, t in start.items(): + o[n].copy_(t) + graph.replay() + torch.cuda.synchronize() + for n, t in got[1].items(): + assert torch.equal(outs[n] if n in outs else o[n], t), n + + def test_front_replays_from_a_cuda_graph(): from sgl_nvidia_mega.qwen3_5.sm103.gdn import front @@ -696,3 +861,293 @@ def test_front_norm_is_slot_gdn_verifys(): norm(core, z, weight, scale, quant, rstd, 1e-6) want = rms_norm_gated_fp8(core, weight, 1e-6, z, scale, activation="silu") assert torch.equal(quant.view(torch.uint8), want.view(torch.uint8).view(-1, V)) + + +# --------------------------------------------- TensorRT-LLM's in_proj front (TRTFRONT) + +TRT_HEADS = {8: (2, 8)} # (H, HV) per rank at TP8 + + +def _trt_operands(seed, tp=8, pad=False, requests=1): + """``requests`` requests of 7 rows at TP ``tp``: the operands of ``trt_front.front`` (the + pools as in ``_front_operands``; ``pad`` gives the last request the pad state index -1).""" + h, hv = TRT_HEADS[tp] + d = 2 * h * K + hv * V + g = torch.Generator(device="cuda").manual_seed(seed) + + def rand(*shape, scale=1.0): + return torch.randn(*shape, device="cuda", generator=g) * scale + + phys = rand(FRONT_SLOTS, d, 9).to(torch.bfloat16) + return dict( + d=d, + x=rand(requests * DRAFT, 4096).to(torch.bfloat16), + w=rand(d + hv * V, 4096, scale=0.05).clamp(-440, 440).to(torch.float8_e4m3fn), + sx=torch.tensor([0.01], device="cuda"), + sw=torch.tensor([0.002], device="cuda"), + ba_w=rand(2 * hv, 4096, scale=0.02).to(torch.bfloat16), + conv_w=rand(d, WIDTH, scale=0.5).to(torch.bfloat16), + state=rand(FRONT_SLOTS, d, WIDTH - 1).to(torch.bfloat16), + phys=phys, + windows=phys.as_strided((FRONT_SLOTS, DRAFT, d, WIDTH - 1), (d * 9, 1, 9, 1)), + states=torch.tensor( + list(FRONT_STATES[: requests - 1]) + + [-1 if pad else FRONT_STATES[requests - 1]], + dtype=torch.int32, + device="cuda", + ), + ) + + +def _trt_outputs(o): + rows = o["x"].shape[0] + return dict( + out=torch.empty(rows, o["w"].shape[0], dtype=torch.bfloat16, device="cuda"), + ba=torch.empty(rows, o["ba_w"].shape[0], dtype=torch.bfloat16, device="cuda"), + ) + + +def _run_trt_front(o, x8, outs, x=None): + from sgl_nvidia_mega.qwen3_5.sm103.gdn import trt_front + + trt_front.prepare( + o["w"].shape[0], o["ba_w"].shape[0], o["conv_w"], o["state"], o["windows"] + ) + trt_front.front( + o["x"] if x is None else x, + x8, + o["w"], + o["sx"], + o["sw"], + o["ba_w"], + o["conv_w"], + o["state"], + o["windows"], + o["states"], + outs["out"], + outs["ba"], + ) + + +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +@pytest.mark.parametrize("act", ["fp8", "quant"]) +@pytest.mark.parametrize("tp", [8]) +def test_trt_front_holds_its_projections_to_fp64_and_its_conv_is_sglangs( + tp, act, requests +): + from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_update + from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 + from sglang.srt.layers.quantization.fp8_utils import flashinfer_bmm_fp8 + + errors = {n: [] for n in ("qkvz", "qkvz_stock", "ba", "ba_stock")} + for seed in range(3): + o = _trt_operands(60 + seed, tp, requests=requests) + d, rows = o["d"], o["x"].shape[0] + state0, phys0 = o["state"].clone(), o["phys"].clone() + xq, _ = static_quant_fp8(o["x"], o["sx"], repeat_scale=False) + outs = _trt_outputs(o) + _run_trt_front(o, xq if act == "fp8" else None, outs) + torch.cuda.synchronize() + # the BF16 q / k / v rows the conv took: window row r holds 2 state columns, then request r's 7 rows + qkv = torch.cat([o["phys"][r, :, 2:].t() for r in range(requests)]) + projected = torch.cat([qkv, outs["out"][:, d:]], dim=1) + ref = (xq.double() @ o["w"].double().t()) * ( + o["sx"].double() * o["sw"].double() + ) + stock = flashinfer_bmm_fp8(xq, o["w"].t(), o["sx"], o["sw"], torch.bfloat16) + errors["qkvz"].append(_errors(projected, ref)) + errors["qkvz_stock"].append(_errors(stock, ref)) + ref = o["x"].double() @ o["ba_w"].double().t() + errors["ba"].append(_errors(outs["ba"], ref)) + errors["ba_stock"].append(_errors(o["x"] @ o["ba_w"].t(), ref)) + # the conv on those rows, as SGLang's update (request r's windows at intermediate row r) + state, phys = state0.clone(), phys0.clone() + windows = phys.as_strided(o["windows"].shape, o["windows"].stride()) + want = causal_conv1d_update( + qkv.contiguous().view(requests, DRAFT, d).transpose(1, 2), + state, + o["conv_w"], + None, + "silu", + conv_state_indices=o["states"], + intermediate_conv_window=windows, + intermediate_state_indices=torch.arange( + requests, dtype=torch.int32, device="cuda" + ), + ) + assert torch.equal(outs["out"][:, :d], want.transpose(1, 2).reshape(rows, d)) + assert torch.equal(o["state"], state) and torch.equal(o["phys"], phys) + e = {n: torch.tensor(v, dtype=torch.float64) for n, v in errors.items()} + print( + f"T0 trt_front tp{tp} {act} requests={requests} qkvz mean rel-L2 {e['qkvz'][:, 0].mean():.4g} vs stock " + f"{e['qkvz_stock'][:, 0].mean():.4g}, max abs {e['qkvz'][:, 1].max():.4g} vs " + f"{e['qkvz_stock'][:, 1].max():.4g}; ba {e['ba'][:, 0].mean():.4g} vs " + f"{e['ba_stock'][:, 0].mean():.4g}, {e['ba'][:, 1].max():.4g} vs " + f"{e['ba_stock'][:, 1].max():.4g}", + flush=True, + ) + for name in ("qkvz", "ba"): + got, stock = e[name], e[f"{name}_stock"] + assert got[:, 0].mean() <= 1.05 * stock[:, 0].mean(), name + assert got[:, 1].max() <= 2 * stock[:, 1].max(), name + + +@pytest.mark.parametrize("tp", [8]) +def test_trt_front_leaves_a_padded_requests_pools_alone(tp): + from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 + + o = _trt_operands(65, tp, pad=True) + state0, phys0 = o["state"].clone(), o["phys"].clone() + _run_trt_front( + o, static_quant_fp8(o["x"], o["sx"], repeat_scale=False)[0], _trt_outputs(o) + ) + torch.cuda.synchronize() + assert torch.equal(o["state"], state0) and torch.equal(o["phys"], phys0) + + +_LATE_PRODUCER = r""" +extern "C" __global__ void late_producer(const uint4* src, uint4* dst, int n16, const int* slot_src, + int* slot_dst, unsigned long long delay_ns) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + unsigned long long t0, t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + do { + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + } while (t - t0 < delay_ns); + for (int i = threadIdx.x; i < n16; i += blockDim.x) dst[i] = src[i]; + if (threadIdx.x == 0) *slot_dst = *slot_src; +} +""" + + +def test_trt_front_reads_its_rows_and_state_index_after_its_pdl_wait(): + """A producer that releases its dependents at entry and writes the front's BF16 / E4M3 + rows and its state index 20 ms later (well past the host's launch of the front): the + front's outputs and pools are those of the written inputs. The build without the wait + (TensorRT-LLM's positive control) reads the stale ones.""" + import ctypes + + from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 + + from sgl_nvidia_mega.common import jit + from sgl_nvidia_mega.qwen3_5.sm103 import trt_lowm + from sgl_nvidia_mega.qwen3_5.sm103.gdn import trt_front + + o = _trt_operands(66) # first: the CUDA context the driver-API load below needs + major, minor = torch.cuda.get_device_capability() + cubin = jit.compile_source( + _LATE_PRODUCER, "late_producer", f"sm_{major}{minor}a", ("-std=c++17",) + ) + producer = jit.load(cubin, "late_producer", (256, 1, 1), 0) + xq = static_quant_fp8(o["x"], o["sx"], repeat_scale=False)[0] + # the inputs as one buffer: BF16 rows, E4M3 rows, and the state index + nx = o["x"].numel() * 2 + src = torch.cat( + [o["x"].view(torch.uint8).flatten(), xq.view(torch.uint8).flatten()] + ) + dst = torch.zeros_like(src) + x, x8 = ( + dst[:nx].view(torch.bfloat16).view(DRAFT, 4096), + dst[nx:].view(torch.float8_e4m3fn).view(DRAFT, 4096), + ) + slot_src, slot = ( + o["states"].clone(), + torch.tensor([FRONT_STATES[1]], dtype=torch.int32, device="cuda"), + ) + o["states"] = slot + start = {n: o[n].clone() for n in ("state", "phys")} + # both builds launched once first (a first launch loads the module: slower than the producer's delay) + _run_trt_front(o, x8, _trt_outputs(o), x=x) + skip = trt_front._SPECS["fp8", o["w"].shape[0], 1]._replace(unsafe_skip_wait=True) + trt_lowm.prepare(skip) + + def run_skip(outs): + trt_lowm.run( + skip, + out=outs["out"], + out2=outs["ba"], + x16=x, + x8=x8, + w_fp8=o["w"], + weight_scale=o["sw"], + input_scale=o["sx"], + w_bf16=o["ba_w"], + conv=(o["conv_w"], o["state"], slot, o["windows"]), + ) + + run_skip(_trt_outputs(o)) + torch.cuda.synchronize() + want = None + for late, skip_wait in ((False, False), (True, False), (True, True)): + for n, t in start.items(): + o[n].copy_(t) + dst.zero_() + slot.fill_(FRONT_STATES[1]) + outs = _trt_outputs(o) + torch.cuda.synchronize() + args = [ + ctypes.c_void_p(src.data_ptr()), + ctypes.c_void_p(dst.data_ptr()), + ctypes.c_int(src.numel() // 16), + ctypes.c_void_p(slot_src.data_ptr()), + ctypes.c_void_p(slot.data_ptr()), + ctypes.c_ulonglong(20_000_000 if late else 0), + ] + jit.launch(producer, (1, 1, 1), args, torch.cuda.current_stream().cuda_stream) + if not late: + torch.cuda.synchronize() + if skip_wait: + run_skip(outs) + else: + _run_trt_front(o, x8, outs, x=x) + torch.cuda.synchronize() + got = {**outs, "state": o["state"].clone(), "phys": o["phys"].clone()} + if want is None: + want = got + elif skip_wait: + assert not all(torch.equal(t, want[n]) for n, t in got.items()) + else: + for n, t in got.items(): + assert torch.equal(t, want[n]), n + + +def test_trt_front_replays_from_a_cuda_graph_and_never_compiles_in_one(): + from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 + + from sgl_nvidia_mega.qwen3_5.sm103 import trt_lowm + + o = _trt_operands(67) + x8 = static_quant_fp8(o["x"], o["sx"], repeat_scale=False)[0] + pools = ("state", "phys") + start = {n: o[n].clone() for n in pools} + outs = _trt_outputs(o) + _run_trt_front(o, x8, outs) # prepare (compile) and run outside capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run_trt_front(o, x8, outs) + for seed in (68, 69): + fresh = _trt_operands(seed) + o["x"].copy_(fresh["x"]) + x8.copy_(static_quant_fp8(o["x"], o["sx"], repeat_scale=False)[0]) + for n in pools: + o[n].copy_(start[n]) + graph.replay() + torch.cuda.synchronize() + got = {n: t.clone() for n, t in outs.items()} + got.update({n: o[n].clone() for n in pools}) + for n in pools: + o[n].copy_(start[n]) + _run_trt_front(o, x8, outs) + torch.cuda.synchronize() + for n, t in outs.items(): + assert torch.equal(got[n], t), n + for n in pools: + assert torch.equal(got[n], o[n]), n + # a launch that was not prepared is refused inside a capture, not compiled there + unprepared = trt_lowm.spec(m=DRAFT, k=4096, n_bf16=16, tactic_bf16=(8, 1)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + with pytest.raises(RuntimeError, match="before graph capture"): + trt_lowm.prepare(unprepared) + with pytest.raises(RuntimeError, match="not prepared"): + trt_lowm.run(unprepared, out=outs["ba"], x16=o["x"], w_bf16=o["ba_w"]) diff --git a/test/kernels/test_qwen3_5_hop.py b/test/kernels/test_qwen3_5_hop.py index dde595ea..7cb914d7 100644 --- a/test/kernels/test_qwen3_5_hop.py +++ b/test/kernels/test_qwen3_5_hop.py @@ -154,6 +154,68 @@ def test_hop_matches_flashinfer_fusion(rank, pair): torch.cuda.synchronize() +def test_finalize_reads_routing_after_the_wait_and_identity_rows(rank, pair): + """The finalize reads routing after its PDL wait; registered identity maps need + no load. Both cases match the fusion with a late rank and freshly written routing.""" + from sgl_nvidia_mega.qwen3_5.sm103 import hop as H + + fusion, hop = pair + world = ps.get_tp_group().world_size + device = torch.device(f"cuda:{torch.cuda.current_device()}") + identity = torch.arange(MAX_M * TOP_K, dtype=torch.int32, device=device) + H.register_identity_rows(identity) + scale = torch.tensor([0.05], device=device) + with torch.inference_mode(): + for i, m in enumerate((1, 2, 5, 7, 15)): + for ident in (False, True): + x = _operands(m, device, rank, seed=400 + i) + if ident: + x["permuted"] = identity[: m * TOP_K].view(m, TOP_K) + want = _run(fusion, "finalize", x) + # the routing rewritten by the launch right before (as the MoE writes it) + w = x["expert_weights"].clone() + x["expert_weights"].zero_() + x["expert_weights"].copy_(w) + if i % (world + 1) == rank: + torch.cuda._sleep(1_000_000) + got = _run(hop, "finalize", x, fp8_scale=scale) + _same(got[:2], want) + torch.cuda.synchronize() + + +def test_identity_rows_registry_keeps_its_buffers(rank): + """A registered identity row map stays allocated after its caller drops it.""" + import gc + + from sgl_nvidia_mega.qwen3_5.sm103 import hop as H + + device = torch.device(f"cuda:{torch.cuda.current_device()}") + n = MAX_M * TOP_K + perm = torch.arange(n, dtype=torch.int32, device=device).flip(0).contiguous() + rows = torch.arange(n, dtype=torch.int32, device=device) + H.register_identity_rows(rows) + ptr = rows.data_ptr() + del rows + gc.collect() + later = [torch.empty_like(perm) for _ in range(8)] + reused = any(t.data_ptr() == ptr for t in later) + alive = any( + block.get("address") == ptr and block["state"] == "active_allocated" + for segment in torch.cuda.memory_snapshot() + for block in segment["blocks"] + ) + if not reused and not alive: + pytest.skip(f"the allocator did not reuse or retain identity buffer {ptr:#x}") + for tensor in later: + tensor.copy_(perm) + identity = ( + tensor.is_contiguous() + and H._IDENTITY_ROWS.get(tensor.data_ptr(), 0) >= tensor.numel() + ) + assert not identity + torch.cuda.synchronize() + + def test_hop_replays_from_a_cuda_graph(rank, pair): fusion, hop = pair device = torch.device(f"cuda:{torch.cuda.current_device()}") diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 6d300d14..8717c03f 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -156,6 +156,29 @@ def test_an_sglang_without_the_eagle_forward_view_is_refused(monkeypatch): {"gdn_verify", "gdn_front"}, "slot 'gdn_front' runs slot 'gdn_pdl''s dependent verify", ), + ( + {"gdn_trt_front"}, + "slot 'gdn_trt_front' replaces slot 'gdn_front' in front of slot " + "'gdn_core'", + ), + ( + {"route_lin"}, + "slot 'route_lin' runs the router before the all-reduce", + ), + ( + { + "route", + "mkm", + "hop", + "gdn_verify", + "gdn_push", + "mka", + "mka_push", + "route_lin", + "route_shared", + }, + "slot 'route_lin' needs 'lowm_se' to replace 'route_shared'", + ), ({"route_shared"}, "slot 'route_shared' runs in slot 'route''s launch"), ({"route", "mkm_wide"}, "slot 'mkm_wide' is slot 'mkm''s 16-token engine"), ( @@ -241,6 +264,24 @@ def _on(slots): return sorted(name for name in plan.Slots.__struct_fields__ if getattr(slots, name)) +def test_the_trt_gdn_front_widens_slot_gdn_front_to_four_requests(): + """The custom front serves one request; the explicit TRT front serves four.""" + front = plan.Facts(gdn_verify=True, gdn_pdl=True, gdn_front=True) + assert plan.slots(7, front).gdn_front + assert not plan.slots(14, front).gdn_front + trt = plan.Facts( + gdn_verify=True, + gdn_pdl=True, + gdn_front=True, + gdn_trt_front=True, + ) + for rows in (7, 14, 21, 28): + slots = plan.slots(rows, trt) + assert slots.gdn_front and slots.gdn_trt_front + assert not plan.slots(35, front).gdn_front + assert not plan.slots(35, trt).gdn_front + + def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): # every slot built except mkm's 16-token engine (mkm_wide, below) every = plan.Facts( @@ -263,7 +304,9 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): assert _on(plan.slots(14, every)) == [ "gdn_conv", "gdn_core", + "gdn_front", "gdn_pdl", + "gdn_trt_front", "gdn_verify", "hop", "hop_fp8", @@ -279,6 +322,7 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): # Three requests keep only the shared expert's validated low-M launch. three = plan.slots(21, every) assert three.lowm_se + assert three.gdn_front and three.gdn_trt_front assert not (three.lowm_in or three.lowm_outproj or three.lowm_oproj) # The 28/56-row low-M paths stay stock; GDN core and MKA serve through 224. for rows in (28, 56, 112, 224): @@ -290,6 +334,8 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): "mka", "route_topk", ] + if rows <= plan.GDN_TRT_FRONT_MAX_ROWS: + expected.extend(("gdn_front", "gdn_trt_front")) if rows <= plan.MKM_M32_MAX_ROWS: expected.append("mkm_m32") if plan.MKM_M32_MAX_ROWS < rows <= plan.MKM_M64_MAX_ROWS: @@ -475,3 +521,9 @@ def test_a_numeric_slot_turns_the_claim_into_its_tolerances(): **spec.NUMERIC_TOLERANCES["skv2"], **spec.NUMERIC_TOLERANCES["mka"], } + small = spec.numerics(frozenset({"gdn_front", "gdn_trt_front", "route_lin"})) + assert dict(small.tolerances) == { + **spec.NUMERIC_TOLERANCES["gdn_front"], + **spec.NUMERIC_TOLERANCES["gdn_trt_front"], + **spec.NUMERIC_TOLERANCES["route_lin"], + } diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index b44466bc..37dd01f9 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -1043,6 +1043,7 @@ def _gdn_front_setup(log): ) st = SimpleNamespace( fb=SimpleNamespace(req_pool_indices=k["req"]), + slots=plan.Slots(gdn_verify=True, gdn_pdl=True, gdn_front=True), view=SimpleNamespace(gdn=meta, mamba=pools), bound={ "gdn": SimpleNamespace( @@ -1129,6 +1130,7 @@ def norm(x, z, weight, scale, quant, rstd, eps): "history_g": "ring_g2(3, 8, 16)", "ring_indices": "req(1,)", "pad_slot": -1, + "prefetch": True, } # q / k / v are the conv rows' views; a / b the front's b8 / a8 halves q, k_, v = seen["qkv"] @@ -1178,6 +1180,241 @@ def norm(x, z, weight, scale, quant, rstd, eps): assert fp8.data_ptr() == quant.data_ptr() and fp8.shape == (7, 1024) +@pytest.mark.parametrize( + "pin, early, rows, prefetch", + [ + (False, False, 7, True), + (True, False, 7, False), + (True, False, 14, True), + (False, True, 14, False), + ], +) +def test_the_gdn_front_leaves_out_its_l2_prefetch_ahead_of_a_pinned_or_early_core( + pin, early, rows, prefetch +): + """In front of slot gdn_core the front's L2 prefetch of the checkpoint / ring is left out + with the core's early history (any request count) and with the pinned one-request front.""" + import msgspec + + from sgl_nvidia_mega.qwen3_5.gdn import front, verify + from sgl_nvidia_mega.qwen3_5.sm103.gdn import core as core_kernel + from sgl_nvidia_mega.qwen3_5.sm103.gdn import front as kernel + + gdn, st = _gdn_front_setup(_log()) + seen = {} + opts = msgspec.structs.replace(verify.SLOT_PDL, core=True) + with ( + patch.object(kernel, "front", lambda *args, **kw: seen.update(kw)), + patch.object(kernel, "PIN", pin), + patch.object(core_kernel, "EARLY", early), + patch.object(verify, "_fused", lambda *args: "fused"), + ): + assert ( + front.step(gdn, torch.zeros(rows, 4096, dtype=torch.bfloat16), st, opts) + == "fused" + ) + assert seen["prefetch"] is prefetch + + +@pytest.mark.parametrize("trt, wants", [(False, False), (True, True)]) +def test_the_hop_quantizes_a_gdn_layers_input_for_the_trt_front_only(trt, wants): + """Slot hop_fp8 writes the input consumed by slot gdn_trt_front.""" + from sglang.srt.models.qwen3_5 import Qwen3_5LinearDecoderLayer + + from sgl_nvidia_mega.qwen3_5 import lowm as lowm_mod + + scale = torch.ones(1) + layer = Qwen3_5LinearDecoderLayer.__new__(Qwen3_5LinearDecoderLayer) + object.__setattr__( + layer, + "linear_attn", + SimpleNamespace(in_proj_qkvz=SimpleNamespace(input_scale=scale)), + ) + st = SimpleNamespace( + slots=plan.Slots( + hop_fp8=True, + gdn_verify=True, + gdn_front=True, + gdn_trt_front=trt, + ) + ) + got = lowm_mod.layer_input_fp8_scale(layer, st, 7) + assert (got is scale) if wants else got is None + + +@pytest.mark.parametrize( + "trt, hop_fp8, rows, core", + [ + (True, True, 7, True), + (True, False, 7, True), + (True, True, 14, True), + (True, True, 21, True), + (True, True, 28, True), + (True, True, 7, False), + (False, True, 7, True), + ], +) +def test_the_trt_gdn_front_runs_ahead_of_slot_gdn_core(trt, hop_fp8, rows, core): + """With slot gdn_trt_front, one to four requests take TensorRT-LLM's + in_proj launch, on the hop's E4M3 rows where it wrote them (BF16 rows otherwise): the raw + [N, K] in_proj_qkvz weight, the layer's conv pools and the step's state index; the core + reads the convolved q / k / v and the raw z of one output, and b / a from the other.""" + import msgspec + + from sgl_nvidia_mega.qwen3_5 import lowm as lowm_mod + from sgl_nvidia_mega.qwen3_5.gdn import front, verify + from sgl_nvidia_mega.qwen3_5.sm103.gdn import front as kernel + from sgl_nvidia_mega.qwen3_5.sm103.gdn import trt_front + + log = _log() + gdn, st = _gdn_front_setup(log) + k = log.known + gdn.in_proj_qkvz.weight = torch.zeros(2560, 4096, dtype=torch.float8_e4m3fn).t() + st.slots = plan.Slots( + hop_fp8=hop_fp8, + gdn_verify=True, + gdn_pdl=True, + gdn_core=core, + gdn_front=True, + gdn_trt_front=trt, + ) + x = k["x"] if rows == 7 else torch.zeros(rows, 4096, dtype=torch.bfloat16) + x8 = torch.zeros(rows, 4096, dtype=torch.float8_e4m3fn) + seen = {} + opts = msgspec.structs.replace(verify.SLOT_PDL, core=core) + with ( + patch.object(trt_front, "front", lambda *args: seen.update(trt=args)), + patch.object(kernel, "front", lambda *args, **kw: seen.update(front=kw)), + patch.object(lowm_mod, "take_fp8", lambda st_, x_: x8 if x_ is x else None), + patch.object( + verify, "_fused", lambda *args: seen.update(fused=args) or "fused" + ), + ): + try: + front.step(gdn, x, st, opts) + except Exception: + assert not core # the stock tail of the mock layer + if not (trt and core): + assert "trt" not in seen and "front" in seen + return + assert "front" not in seen + args = seen["trt"] + assert args[0] is x and args[1] is (x8 if hop_fp8 else None) + w = args[2] + assert w.shape == (2560, 4096) and w.stride() == (4096, 1) + assert tuple(map(log.tag, args[3:10])) == ( + "in_s(1,)", + "in_ws(1,)", + "ba_w(16, 4096)", + "conv_w(1536, 4)", + "conv2(3, 1536, 3)", + "window2(3, 7, 1536, 3)", + "states(1,)", + ) + out, ba = args[10], args[11] + assert out.shape == (rows, 2560) and ba.shape == (rows, 16) + _, qkv, z, a, b, *_ = seen["fused"] + assert qkv.data_ptr() == out.data_ptr() and qkv.shape == (rows, 1536) + assert qkv.stride() == (2560, 1) + assert z.data_ptr() == out[:, 1536:].data_ptr() + assert z.shape == (rows, 8, 128) and z.stride() == (2560, 128, 1) + assert a.data_ptr() == ba[:, 8:].data_ptr() and b.data_ptr() == ba.data_ptr() + assert a.shape == b.shape == (rows, 8) + + +def test_the_trt_gdn_front_prepares_on_the_first_gdn_layers_pools(): + """``front.prepare`` (the executor's first covered step, before any capture) compiles the + TensorRT-LLM front on a GDN layer's weights and its pools' layouts; nothing when + the slot is absent.""" + from sglang.srt.models.qwen3_5 import Qwen3_5LinearDecoderLayer + + from sgl_nvidia_mega.qwen3_5.gdn import front + from sgl_nvidia_mega.qwen3_5.sm103.gdn import trt_front + + log = _log() + gdn, st = _gdn_front_setup(log) + gdn.in_proj_qkvz.weight = torch.zeros(2560, 4096, dtype=torch.float8_e4m3fn).t() + layer = Qwen3_5LinearDecoderLayer.__new__(Qwen3_5LinearDecoderLayer) + object.__setattr__(layer, "linear_attn", gdn) + lm = SimpleNamespace( + model=SimpleNamespace( + start_layer=1, end_layer=3, layers=[None, SimpleNamespace(), layer] + ) + ) + seen = [] + for enabled in (False, True): + with patch.object(trt_front, "prepare", lambda *args: seen.append(args)): + front.prepare(lm, st.view.mamba, trt=enabled) + assert len(seen) == 1 and seen[0][:2] == (2560, 16) + assert tuple(map(log.tag, seen[0][2:])) == ( + "conv_w(1536, 4)", + "conv2(3, 1536, 3)", + "window2(3, 7, 1536, 3)", + ) + + +@pytest.mark.parametrize( + "qkvz, ba, d, requests, tactic_fp8, tactic_quant", + [ + (2560, 16, 1536, 1, (16, 2, 1, True, 2), (16, 4, 1, True, 1)), + (2560, 16, 1536, 2, (8, 4, 1, False, 1), (16, 4, 1, False, 1)), + (2560, 16, 1536, 4, (8, 4, 1, False, 1), (16, 4, 1, False, 1)), + ], +) +def test_trt_lowm_takes_trts_gdn_tactics_and_refuses_an_unprepared_launch( + qkvz, ba, d, requests, tactic_fp8, tactic_quant +): + """The TensorRT-LLM front's TP8 launches use its GDN mixer's tactics + (mixed8 on E4M3 rows, mixed on BF16 rows: k_warps, n_tiles, ksplit, swap_ab, min_blocks; + the BF16 segment one tile a CTA), PDL with the early trigger; ``trt_lowm`` refuses to compile + inside a capture and to launch what was not prepared.""" + from sgl_nvidia_mega.qwen3_5.sm103 import trt_lowm + from sgl_nvidia_mega.qwen3_5.sm103.gdn import trt_front + + conv_w = torch.zeros(d, 4, dtype=torch.bfloat16) + state = torch.zeros(3, d, 3, dtype=torch.bfloat16) + windows = torch.zeros(3, d, 9, dtype=torch.bfloat16).as_strided( + (3, 7, d, 3), (d * 9, 1, 9, 1) + ) + for act, tactic in (("fp8", tactic_fp8), ("quant", tactic_quant)): + s = trt_front.spec(act, qkvz, ba, conv_w, state, windows, requests) + assert s.m == 7 * requests and s.conv_shapes[2] == requests + assert (s.k_warps, s.n_tiles, s.ksplit, s.swap_ab, s.min_blocks) == tactic + assert (s.n_fp8, s.n_bf16, s.n_tiles_bf16, s.fp8_act) == (qkvz, ba, 1, act) + assert ( + s.n_out, + s.n_out2, + s.conv_channels, + s.conv_tokens, + s.conv_write_state, + ) == (qkvz, ba, d, 7, True) + assert (s.a16_stride, s.a8_stride) == (4096, 4096 if act == "fp8" else 0) + assert ( + s.use_pdl + and s.early_trigger + and not s.release_after_wait + and not s.unsafe_skip_wait + ) + assert s.conv_shapes[3] == ((3, 7, d, 3), (d * 9, 1, 9, 1)) + with pytest.raises(RuntimeError, match="not prepared"): + trt_lowm.run(s, out=torch.zeros(7 * requests, qkvz, dtype=torch.bfloat16)) + with patch.object(torch.cuda, "is_current_stream_capturing", lambda: True): + with pytest.raises(RuntimeError, match="before graph capture"): + trt_lowm.prepare(s) + + +def test_sglangs_verify_intermediate_rows_are_the_batch_rows(monkeypatch): + """The TensorRT-LLM front writes a request's conv windows at its batch row, as SGLang's + target verify indexes them (verify_intermediate_state_indices: request r, row r).""" + from sglang.srt.layers.attention.linear import utils + from sglang.srt.utils import common + + monkeypatch.setattr(common, "get_eager_max_batch_size", lambda n: n) + assert utils.build_verify_intermediate_state_indices(8, "cpu").tolist() == list( + range(8) + ) + + @pytest.mark.parametrize("front_slot", [False, True]) def test_the_gdn_step_runs_the_front_where_the_plan_gives_it(monkeypatch, front_slot): from sgl_nvidia_mega.qwen3_5 import gdn @@ -2061,7 +2298,11 @@ def layer(kind): packed, why = [], {} monkeypatch.setattr(runtime, "SLOTS", frozenset({"mka"})) monkeypatch.setattr(mka, "unsupported", lambda l: why.get(id(l), "")) - monkeypatch.setattr(mka, "pack", packed.append) + monkeypatch.setattr( + mka, + "pack", + lambda layer, route_lin=False: packed.append(layer) or "", + ) ctx = SimpleNamespace(tokens_per_request=7) for made in layers[1:]: made.num_heads = 8 From 97426041cc8809c65155b876e8194f7482e309b5 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 16:07:23 -0700 Subject: [PATCH 12/62] feat: add TEP8 routed MoE body --- sgl_nvidia_mega/qwen3_5/lowm.py | 15 +- sgl_nvidia_mega/qwen3_5/moe/__init__.py | 30 +- sgl_nvidia_mega/qwen3_5/moe/body.py | 280 ++++++++ sgl_nvidia_mega/qwen3_5/moe/route.py | 191 +++++- sgl_nvidia_mega/qwen3_5/plan.py | 69 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 4 + sgl_nvidia_mega/qwen3_5/sm103/moe/router.py | 240 ++++++- .../qwen3_5/sm103/moe/src/route_big.cu | 313 +++++++++ .../qwen3_5/sm103/moe/src/route_tables.cuh | 401 +++++++++++ .../qwen3_5/sm103/moe/src/route_topk.cu | 158 +---- test/kernels/test_qwen3_5_moe.py | 627 +++++++++++++++++- test/unit/test_qwen3_5_spec.py | 69 +- 12 files changed, 2224 insertions(+), 173 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/moe/body.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_big.cu create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_tables.cuh diff --git a/sgl_nvidia_mega/qwen3_5/lowm.py b/sgl_nvidia_mega/qwen3_5/lowm.py index 2864aa30..15f02f61 100644 --- a/sgl_nvidia_mega/qwen3_5/lowm.py +++ b/sgl_nvidia_mega/qwen3_5/lowm.py @@ -141,6 +141,7 @@ class Bound(msgspec.Struct, frozen=True): False # every GDN / attention FP8 output projection takes the launch ) se: bool = False # every shared expert takes the two launches + se_pre: bool = False # the launches also cover the MoE body's tier-M rows push: bool = False # the output projections' launches also push into the hop hop_fp8: bool = False # the hop also writes the next projection's E4M3 input # lowm_se: id(mlp) -> (gate_up rows paired for the SwiGLU, the shared gate zero-padded to 8 rows) @@ -229,6 +230,17 @@ def build(model, ctx) -> Bound: "lowm_out" in SLOTS, "lowm_se" in SLOTS, ) + se_rows = plan.LOWM_SE_MAX_ROWS + if want_se and {"route", "route_topk"} <= SLOTS: + from sglang.srt.runtime_context import get_parallel + + par = get_parallel() + tep8_m = ( + par.tp_size == 8 + and par.moe_ep_size > 1 + and {"route", "route_topk"} <= SLOTS + ) + se_rows = plan.lowm_se_pre_rows(par.tp_size, tep8_m) or se_rows want_push = want_out and "lowm_push" in SLOTS and "hop" in SLOTS want_fp8 = "hop_fp8" in SLOTS and "hop" in SLOTS want = want_in or want_out @@ -310,7 +322,7 @@ def build(model, ctx) -> Bound: n, k = gu.shape nd, kd = _w8(se.down_proj).shape for w in ctx.widths: - if w <= plan.LOWM_SE_MAX_ROWS: + if w <= se_rows: launches.add((w, k, n, 8, 1, 1, True)) if want_fp8 and w <= plan.HOP_MAX_ROWS: launches.add((w, k, n, 8, 0, 1, True)) @@ -334,6 +346,7 @@ def build(model, ctx) -> Bound: lowm_in=want_in, lowm_out=want_out, se=want_se, + se_pre=want_se and se_rows > plan.LOWM_SE_MAX_ROWS, push=want_push, hop_fp8=want_fp8, shared=shared, diff --git a/sgl_nvidia_mega/qwen3_5/moe/__init__.py b/sgl_nvidia_mega/qwen3_5/moe/__init__.py index 2457e2b7..8697c1ea 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/moe/__init__.py @@ -74,9 +74,9 @@ def build(model, ctx) -> route.Bound: - """Slots ``route``, ``mkm``, ``route_topk`` and ``mkm_m32`` for every MoE layer; the - result is ``st.bound["moe"]``.""" + """Build the requested MoE slots and TEP8's internal routed-body plan.""" import msgspec + from sglang.srt.runtime_context import get_parallel from sgl_nvidia_mega.qwen3_5 import plan from sgl_nvidia_mega.qwen3_5.runtime import SLOTS @@ -133,6 +133,21 @@ def build(model, ctx) -> route.Bound: m512=m512, wide_bal="mkm_wide_bal" in SLOTS, ) + parallel = get_parallel() + tep8_moe = bound.route and parallel.tp_size == 8 and parallel.moe_ep_size > 1 + if tep8_moe: + reason = route.pre_unsupported(model) + if reason: + return msgspec.structs.replace(bound, reason=f"TEP8 MoE body: {reason}") + bound = msgspec.structs.replace(bound, pre=True) + if ( + plan.widest(ctx.widths, plan.ROUTE_BIG_MAX_ROWS) + >= plan.ROUTE_BIG_MIN_ROWS + ): + mlp = model.model.layers[model.model.start_layer].mlp + router.build_big( + torch.cuda.current_device(), mlp.experts.num_local_experts + ) return bound @@ -239,7 +254,12 @@ def _router_experts(mlp, x, defer: bool, st): if rl is not None and rl[0] == mlp.layer_id and st.slots.mkm: # Slot route_lin produced the logits and NVFP4 input in the preceding hop. return route.routed_lin(mlp, x, st.bound["moe"], *rl[1:]) - return route.routed(mlp, x, st.bound["moe"] if st.slots.mkm else None) + return route.routed( + mlp, + x, + st.bound["moe"] if st.slots.mkm else None, + pre_topk=st.slots.route_pre_topk, + ) if defer and st.slots.mkm_m32: return route.routed_m32(mlp, x, st.bound["moe"], gemm=st.slots.route_gemm) if defer and st.slots.mkm_m64: @@ -251,7 +271,9 @@ def _router_experts(mlp, x, defer: bool, st): if defer and st.slots.mkm_m512: return route.routed_given(mlp, x, st.bound["moe"], "mkm_m512") if defer and st.slots.route_topk: - return route.routed_topk(mlp, x, gemm=st.slots.route_gemm) + return route.routed_topk( + mlp, x, gemm=st.slots.route_gemm, pre=st.slots.route_pre + ) router_logits, _ = mlp.gate(x) topk_output = mlp.topk(x, router_logits) if defer: diff --git a/sgl_nvidia_mega/qwen3_5/moe/body.py b/sgl_nvidia_mega/qwen3_5/moe/body.py new file mode 100644 index 00000000..83930ab7 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/moe/body.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Adapted from SGLang (Apache-2.0, SGLang Team): +# python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +# (fused_experts_none_to_flashinfer_trtllm_fp4, its routed deferred-finalize branch). +"""The trtllm-gen NVFP4 MoE body on routing tables written by TEP8's internal +route_topk or route-big front. It replaces the runner's op, whose own routing kernel +(``routingIndicesDynBlock`` from logits or ``routingIndicesCluster`` from materialized +routing) builds the same tables before the same two GEMMs. Under expert parallelism the +tables cover the rank's local +experts (local_expert_offset = moe_ep_rank * num_local_experts, as SGLang's runner) and +a slot routed elsewhere gets permuted index -1, as the stock routing gives it. + +The body is FlashInfer's DA entry (``MoERunner.forward`` with ``da_routing_metadata``): the +tactic the routed op would run (the autotuner's choice for the same op, inputs and tuning +config), FC1 and FC2 from the tables, and the routed op's deferred-finalize outputs +(FC2 rows, routing weights, expanded -> permuted map). FlashInfer's own table builder +(populate_routing_metadata_multi_tile) serves top-k <= 8 only, so the tables are ours. +Each width keeps one set of tables and body buffers for every layer: a layer's FC2 rows +are finalized at the next boundary, before the next layer's body writes them. The body +buffers (FC1 / FC2 outputs and GEMM scratch, sized by the width's row bound) are views of +one set every width shares, since one step runs at a time (TEP8 serves 96 widths). +""" + +from __future__ import annotations + +import logging +from typing import Optional + +import torch + +logger = logging.getLogger(__name__) + +# The SGLang names this module reads (test/unit/test_contract.py pins them) +SGLANG_READS = { + "srt/layers/moe/moe_runner/flashinfer_trtllm.py": ( + "get_activation_type", + "get_activation_type(is_gated)", + "trtllm_moe_enable_pdl", + ), + "srt/layers/moe/fused_moe_triton/layer.py": ( + "moe_runner_config", + "num_experts", + "num_local_experts", + "moe_ep_rank", + "top_k", + "intermediate_size_per_partition", + ), + "srt/utils/common.py": ("next_power_of_2",), +} +# The FlashInfer names this module reads (its DA body entry, flashinfer 0.7.1) +FLASHINFER_READS = { + "flashinfer/fused_moe/core.py": ( + "get_trtllm_moe_sm100_module", + "MoERunner", + "_make_tuning_config", + "forward(da_routing_metadata)", + "forward(da_body_workspace)", + "forward(do_preparation)", + ), + "flashinfer/fused_moe/shared/inputs.py": ( + "MoeRunnerInputs", + "alloc_trtllm_moe_output", + ), +} +OP = "flashinfer::trtllm_fp4_block_scale_moe" # the routed op's autotuner entry + + +class Width: + """One width's runner, routing tables and body buffers.""" + + def __init__(self): + self.runner = None + self.tables = None # sm103.moe.router.RouteTables of the tactic's tile_n + self.tactic: Optional[list] = None + self.workspace: Optional[tuple] = None + + +_WIDTHS: dict[tuple, Width] = {} +# the body buffers the widths share (FlashInfer's FP4 DA body order), grown to the largest asked for +_SHARED: list[torch.Tensor] = [] + + +def _shared(prepared) -> tuple: + """Views of the shared body buffers in place of a width's ``prepared`` ones (same dtype + and row size, at least as many rows); a larger buffer replaces the shared one.""" + out = [] + for i, t in enumerate(prepared): + if i == len(_SHARED): + _SHARED.append(t) + s = _SHARED[i] + if s.dtype != t.dtype or s.shape[1:] != t.shape[1:] or s.shape[0] < t.shape[0]: + _SHARED[i] = s = t + out.append(s[: t.shape[0]]) + return tuple(out) + + +def local_offset(experts) -> int: + """The rank's first local expert (SGLang's FP4 trtllm quant info).""" + return experts.moe_ep_rank * experts.num_local_experts + + +def _kwargs(experts, m: int) -> dict: + """The routed op's runner arguments for ``experts`` at m rows (SGLang's routed call: + unpacked top-10, deferred finalize, no per-token scale).""" + from flashinfer.tllm_enums import RoutingInputMode + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + get_activation_type, + trtllm_moe_enable_pdl, + ) + + return dict( + routing_input_mode=RoutingInputMode.UnpackedPrecomputed, + num_experts=experts.num_experts, + routing_bias=None, + gemm1_weights=experts.w13_weight.data, + gemm1_weights_scale=experts.w13_weight_scale.data.view(torch.float8_e4m3fn), + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=experts.w2_weight.data, + gemm2_weights_scale=experts.w2_weight_scale.data.view(torch.float8_e4m3fn), + gemm2_bias=None, + output1_scale_scalar=experts.g1_scale_c.data, + output1_scale_gate_scalar=experts.g1_alphas.data, + output2_scale_scalar=experts.g2_alphas.data, + n_group=0, + topk_group=0, + local_expert_offset=local_offset(experts), + routed_scaling_factor=None, + routing_method_type=1, # unused for precomputed routing (SGLang passes 1) + enable_pdl=trtllm_moe_enable_pdl(m), + do_finalize=False, + activation_type=get_activation_type( + experts.moe_runner_config.activation, + is_gated=experts.moe_runner_config.is_gated, + ), + num_fused_shared_experts=0, + norm_topk_prob=True, + routing_replay_out=None, + valid_hidden_size=None, + valid_intermediate_size=None, + ) + + +def _inputs(tables, x_fp4: torch.Tensor, x_sf: torch.Tensor): + """The runner's inputs (MoeRunnerInputs) as the routed op builds them.""" + from flashinfer.fused_moe.shared.inputs import ( + MoeRunnerInputs, + alloc_trtllm_moe_output, + ) + + m = x_fp4.shape[0] + return MoeRunnerInputs( + output=alloc_trtllm_moe_output(m, 2 * x_fp4.shape[1], False, x_fp4.device), + routing_logits=None, + topk_ids=tables.ids, + expert_weights=tables.weights, + hidden_states=x_fp4, + hidden_states_scale=x_sf, + gemm1_lora_delta=None, + per_token_scale=None, + ) + + +# FlashInfer's FP4 tile_n candidates (trtllm_fused_moe_kernel_launcher.cu mSupportedTileNums) +TILES = (8, 16, 32, 64, 128) + + +def default_tile(m: int, top_k: int, local_experts: int) -> int: + """The tile_n the routed op runs without a tuned tactic (the launcher's + selectDefaultTileN: the smallest of the candidates around the first tile that covers + the mean rows an expert gets).""" + mean = m * top_k / local_experts + i = next((i for i, t in enumerate(TILES) if t >= mean), len(TILES) - 1) + return TILES[max(i - 1, 0)] + + +def width(experts, x: torch.Tensor) -> Width: + """The width of ``x`` [m, hidden]: on first use its runner, the routed op's tactic at + m rows (the autotuner's; without one, FlashInfer's default, [-1, -1]) and that tile_n's + tables (the body buffers follow on the first ``run``).""" + from flashinfer.autotuner import AutoTuner + from flashinfer.fused_moe.core import get_trtllm_moe_sm100_module + from flashinfer.tllm_enums import ( + Fp8QuantizationType, + RoutingInputMode, + WeightLayout, + deduce_trtllm_gen_tensor_dtype, + ) + from sglang.srt.utils import next_power_of_2 + + from sgl_nvidia_mega.qwen3_5.sm103.moe import router + + m = x.shape[0] + window = (local_offset(experts), experts.num_local_experts) + key = (m, experts.intermediate_size_per_partition) + window + if key in _WIDTHS: + return _WIDTHS[key] + w = Width() + module = get_trtllm_moe_sm100_module() + kw = _kwargs(experts, m) + x_fp4 = torch.empty((m, x.shape[1] // 2), dtype=torch.uint8, device=x.device) + x_sf = torch.empty( + (m, x.shape[1] // 16), dtype=torch.float8_e4m3fn, device=x.device + ) + w.runner = module.MoERunner( + module.moe_op, + top_k=experts.top_k, + num_local_experts=experts.num_local_experts, + dtype_act=deduce_trtllm_gen_tensor_dtype(x_fp4, x_sf), + dtype_weights=deduce_trtllm_gen_tensor_dtype( + kw["gemm1_weights"], kw["gemm1_weights_scale"] + ), + fp8_quantization_type=Fp8QuantizationType.NoneFp8, + hidden_size=x.shape[1], + hidden_size_output=x.shape[1], + intermediate_size=experts.intermediate_size_per_partition, + activation_type=kw["activation_type"], + weight_layout=WeightLayout.MajorK, + use_shuffled_weight=True, + use_per_token_scaling=False, + num_experts=experts.num_experts, + num_fused_shared_experts=0, + ) + inputs = _inputs(router.RouteTables(m, 8, x.device, *window), x_fp4, x_sf) + config = w.runner._make_tuning_config( + inputs, + tune_max_num_tokens=next_power_of_2(m), + routing_input_mode=RoutingInputMode.UnpackedPrecomputed, + hidden_states_scale_layout=None, + use_cold_l2_cache=True, + use_cuda_graph=True, + ) + _, tactic = AutoTuner.get().choose_one( + OP, [w.runner], config, inputs.to_list(), **kw + ) + if isinstance(tactic, int) or tactic[0] <= 0 or tactic[1] < 0: + # the launcher's fallback: the default tile_n and its default config + w.tactic = [-1, -1] + tile = default_tile(m, experts.top_k, experts.num_local_experts) + logger.info( + "NVIDIA megakernel: TEP8 routed body runs the default tactic " + "(tile_n %d) at %d rows", + tile, + m, + ) + else: + w.tactic = [int(tactic[0]), int(tactic[1])] + tile = w.tactic[0] + w.tables = router.RouteTables(m, tile, x.device, *window) + _WIDTHS[key] = w + return w + + +def run(w: Width, experts, x_fp4: torch.Tensor, x_sf: torch.Tensor): + """The body on ``w``'s tables (written by the routing launch for these rows): the + routed op's deferred-finalize outputs [FC2 rows, routing weights, expanded -> + permuted map].""" + kw = _kwargs(experts, x_fp4.shape[0]) + inputs = _inputs(w.tables, x_fp4, x_sf).to_list() + metadata = w.tables.metadata() + if w.workspace is None: + w.workspace = _shared( + w.runner.forward( + inputs, + tactic=w.tactic, + do_preparation=True, + da_routing_metadata=metadata, + **kw, + ) + ) + return w.runner.forward( + inputs, + tactic=w.tactic, + da_routing_metadata=metadata, + da_body_workspace=w.workspace, + **kw, + ) diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index b330960d..e7ec858e 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -24,6 +24,12 @@ materialized routing (FlashInfer's routed MoE), which skips its own top-k kernel. The logits are SGLang's, so the step's output is the stock one's. +Under TEP8, route_topk or route-big also writes the routing tables of the runner's +trtllm-gen MoE body, which then runs on them (``body.py``) in place of the runner's routed +op and its routing kernel. +The body's FC2 rows equal the op's bit for bit, and under expert parallelism a slot outside +the rank's experts gets permuted index -1 as there. + Slot ``mkm_m32`` takes route_topk's place at 17 to 32 rows: slot mkm's 32-token engine runs the top-10, FC1, SwiGLU, the NVFP4 requantization and FC2 in one launch on the router GEMM's logits and the NVFP4 input (with route_gemm the router front's tiles give both; @@ -74,17 +80,29 @@ "run_moe_core", "run_moe_core(dispatch_output)", "quant_method", + "moe_runner_config", ), "srt/layers/moe/token_dispatcher/standard.py": ( "dispatch(hidden_states)", "dispatch(topk_output)", "combine(combine_input)", "hidden_states_scale", + "StandardCombineInput", ), "srt/layers/moe/moe_runner/flashinfer_trtllm.py": ( "flashinfer_trtllm_deferred_finalize_context", "trtllm_padded_hidden_size", + # TEP8 body: the routed op's deferred-finalize output + "_make_deferred_finalize_output", + "_make_deferred_finalize_output(top_k)", + "_make_deferred_finalize_output(materialized_routing)", + ), + "srt/layers/moe/utils.py": ( + "get_moe_runner_backend", + "is_flashinfer_trtllm", + "is_flashinfer_trtllm_routed", ), + "srt/layers/moe/moe_runner/base.py": ("activation", "is_gated"), "srt/layers/quantization/modelopt_quant.py": ( '"w13_input_scale_quant"', "quant_config", @@ -135,11 +153,16 @@ class Bound(msgspec.Struct, frozen=True): # A caller-specific upper bound for the 16-token engine. The draft uses 12 so # its alternate-stream shared expert can overlap rows 13-16 with the 32-token engine. m16_max_rows: int = 16 + pre: bool = False # TEP8's body on route_topk/route-big routing tables reason: str = "" # why a requested slot is not built def build( - model, max_rows: int, fused: bool, shared: bool = False, wide: bool = False + model, + max_rows: int, + fused: bool, + shared: bool = False, + wide: bool = False, ) -> Bound: """The router front for every MoE layer, compiled for 1..max_rows rows, with ``fused`` slot mkm's kernel and weight maps (with ``wide`` its 16-token engine too, @@ -294,6 +317,39 @@ def topk_unsupported(model) -> str: return "" +def pre_unsupported(model) -> str: + """Why TEP8's routed body cannot serve ("" when it can).""" + from sglang.srt.layers.moe.utils import get_moe_runner_backend + + try: + from flashinfer.fused_moe.core import get_trtllm_moe_sm100_module + from flashinfer.fused_moe.shared.inputs import ( # noqa: F401 + MoeRunnerInputs, + alloc_trtllm_moe_output, + ) + from flashinfer.tllm_enums import RoutingInputMode # noqa: F401 + except ImportError as e: + return f"FlashInfer without the MoE body entry ({e})" + module = get_trtllm_moe_sm100_module() + if not hasattr(module, "MoERunner") or not hasattr(module, "DABodyRunner"): + return "FlashInfer without the MoE body entry (MoERunner / DABodyRunner)" + backend = get_moe_runner_backend() + if not (backend.is_flashinfer_trtllm() or backend.is_flashinfer_trtllm_routed()): + return f"MoE runner backend {backend}" + m = model.model + for i in range(m.start_layer, m.end_layer): + experts = m.layers[i].mlp.experts + cfg = experts.moe_runner_config + if cfg.activation != "silu" or not cfg.is_gated: + return f"layer {i}: activation {cfg.activation} (gated {cfg.is_gated})" + for name in ("gemm1_alpha", "gemm1_beta", "gemm1_clamp_limit"): + if getattr(experts, name, None) is not None: + return f"layer {i}: {name}" + if getattr(experts, "trtllm_padded_hidden_size", None) is not None: + return f"layer {i}: padded hidden size" + return "" + + def shared_unsupported(mlp) -> str: """Why the router front cannot compute ``mlp``'s shared gate/up ("" when it can): a static per-tensor FP8 projection of at most 512 rows, then SiLU-and-mul.""" @@ -337,15 +393,28 @@ def shared_expert_tail(expert, gate_up: torch.Tensor) -> torch.Tensor: return x -def routed(mlp, x: torch.Tensor, fused: Bound | None = None, shared: bool = False): +def routed( + mlp, + x: torch.Tensor, + fused: Bound | None = None, + shared: bool = False, + pre_topk: bool = False, +): """``FusedMoE.forward_deferred_finalize`` of ``mlp``'s routed experts on ``x`` [rows, hidden], with the router front's logits and quantized input (``fused``: on - slot mkm's kernel). With ``shared`` (at most 8 rows) the router launch also computes - the shared expert's gate/up: returns (routed, gate_up, the event recorded on the - current stream right after that launch).""" + slot mkm's kernel). With ``pre_topk``, a routing-only launch writes TEP8's body + tables. With ``shared`` (at most 8 rows) the + router launch also computes the shared expert's gate/up: returns (routed, gate_up, the + event recorded on the current stream right after that launch).""" from sgl_nvidia_mega.qwen3_5.sm103.moe import router experts = mlp.experts + width = tables = None + if pre_topk: + from sgl_nvidia_mega.qwen3_5.moe import body + + width = body.width(experts, x) + tables = width.tables for slot in ("route" if fused is None else "mkm",) + ( ("route_shared",) if shared else () ): @@ -354,6 +423,12 @@ def routed(mlp, x: torch.Tensor, fused: Bound | None = None, shared: bool = Fals logger.info( "NVIDIA megakernel: Qwen3.5 slot %s runs at %d rows", slot, x.shape[0] ) + if pre_topk and ("tep8_body", x.shape[0]) not in _ENGAGED: + _ENGAGED.add(("tep8_body", x.shape[0])) + logger.info( + "NVIDIA megakernel: Qwen3.5 TEP8 routed body runs at %d rows", + x.shape[0], + ) if shared: gu = mlp.shared_expert.gate_up_proj logits, x_fp4, x_sf, gate_up = router.router_gemm_quant_shared( @@ -368,9 +443,22 @@ def routed(mlp, x: torch.Tensor, fused: Bound | None = None, shared: bool = Fals ready.record() out = _routed_experts(mlp, x, fused, logits, x_fp4, x_sf) return out, gate_up, ready - logits, x_fp4, x_sf = router.router_gemm_quant( - x, mlp.gate.weight, experts.w13_input_scale_quant.reshape(1) - ) + scale = experts.w13_input_scale_quant.reshape(1) + if pre_topk: + from sgl_nvidia_mega.qwen3_5 import plan + + rows = x.shape[0] + if ("pre_tier", rows) not in _ENGAGED: + _ENGAGED.add(("pre_tier", rows)) + logger.info( + "NVIDIA megakernel: Qwen3.5 MoE front tier %s at %d rows: slot route's front, then route_topk's tables", + plan.moe_front_tier(rows), + rows, + ) + logits, x_fp4, x_sf = router.router_gemm_quant(x, mlp.gate.weight, scale) + router.route_topk_quant(logits, x, scale, quantize=False, tables=tables) + return _pre_body(mlp, width, x_fp4, x_sf) + logits, x_fp4, x_sf = router.router_gemm_quant(x, mlp.gate.weight, scale) return _routed_experts(mlp, x, fused, logits, x_fp4, x_sf) @@ -389,6 +477,26 @@ def routed_lin(mlp, x: torch.Tensor, fused: Bound, logits, x_fp4, x_sf): return _routed_experts(mlp, x, fused, logits, x_fp4, x_sf) +def _pre_body(mlp, width, x_fp4, x_sf): + """TEP8's MoE body on ``width``'s tables, as a deferred-finalize output.""" + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + _make_deferred_finalize_output, + ) + from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput + + from sgl_nvidia_mega.qwen3_5.moe import body + from sgl_nvidia_mega.qwen3_5.sm103.moe import router + + experts = mlp.experts + result = body.run(width, experts, x_fp4, x_sf.view(torch.float8_e4m3fn)) + deferred = _make_deferred_finalize_output( + result, top_k=router.TOP_K, materialized_routing=True + ) + return experts.dispatcher.combine( + combine_input=StandardCombineInput(hidden_states=deferred) + ) + + def _routed_experts(mlp, x, fused, logits, x_fp4, x_sf): """The routed experts on the router front's outputs (``fused``: slot mkm's kernel), deferred.""" @@ -422,32 +530,89 @@ def _routed_experts(mlp, x, fused, logits, x_fp4, x_sf): return experts.dispatcher.combine(combine_input=combine_input) -def routed_topk(mlp, x: torch.Tensor, gemm: bool = False): +def routed_topk(mlp, x: torch.Tensor, gemm: bool = False, pre: bool = False): """``FusedMoE.forward_deferred_finalize`` of ``mlp``'s routed experts on ``x`` [rows, hidden] with SGLang's router GEMM, then slot route_topk's launch (the top-10 and the NVFP4 input), and the runner on the materialized routing. With ``gemm`` (slot route_gemm, numeric) the router front's GEMM and quantizer in 16-row tiles - stand in for SGLang's router GEMM, and route_topk's launch only routes.""" + stand in for SGLang's router GEMM, and route_topk's launch only routes. With ``pre``, + TEP8's internal front also writes the MoE body's routing tables and the body runs on + them (``body.py``) in place of the runner's routed op.""" from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( flashinfer_trtllm_deferred_finalize_context, ) from sglang.srt.layers.moe.topk import StandardTopKOutput + from sgl_nvidia_mega.qwen3_5 import plan from sgl_nvidia_mega.qwen3_5.sm103.moe import router if ("route_topk", x.shape[0]) not in _ENGAGED: _ENGAGED.add(("route_topk", x.shape[0])) logger.info( - "NVIDIA megakernel: Qwen3.5 slot route_topk runs at %d rows", x.shape[0] + "NVIDIA megakernel: Qwen3.5 slot route_topk runs at %d rows", + x.shape[0], + ) + if pre and ("tep8_body", x.shape[0]) not in _ENGAGED: + _ENGAGED.add(("tep8_body", x.shape[0])) + logger.info( + "NVIDIA megakernel: Qwen3.5 TEP8 routed body runs at %d rows", + x.shape[0], ) experts = mlp.experts scale = experts.w13_input_scale_quant.reshape(1) + width = tables = None + if pre: + from sgl_nvidia_mega.qwen3_5.moe import body + + width = body.width(experts, x) + tables = width.tables if gemm: logits, x_fp4, x_sf = router.router_gemm_quant_tiled(x, mlp.gate.weight, scale) - ids, weights, _, _ = router.route_topk_quant(logits, x, scale, quantize=False) + ids, weights, _, _ = router.route_topk_quant( + logits, x, scale, quantize=False, tables=tables + ) + elif pre and x.shape[0] >= plan.ROUTE_BIG_MIN_ROWS: + rows = x.shape[0] + if ("route_big", rows) not in _ENGAGED: + _ENGAGED.add(("route_big", rows)) + cl, _, words = router.big_shape(rows) + logger.info( + "NVIDIA megakernel: Qwen3.5 MoE front tier %s at %d rows: " + "route-big C%d%s", + plan.moe_front_tier(rows), + rows, + cl, + "W" if words > 1 else "", + ) + logits, _ = mlp.gate(x) + ids, weights, x_fp4, x_sf = router.route_big_quant( + logits, + x, + scale, + tables, + ctas=( + plan.ROUTE_BIG_WIDE_CTAS + if rows >= plan.ROUTE_BIG_WIDE_MIN_ROWS + else None + ), + ) else: logits, _ = mlp.gate(x) - ids, weights, x_fp4, x_sf = router.route_topk_quant(logits, x, scale) + ids, weights, x_fp4, x_sf = router.route_topk_quant( + logits, x, scale, tables=tables + ) + if pre: + rows = x.shape[0] + if rows < plan.ROUTE_BIG_MIN_ROWS and ("pre_tier", rows) not in _ENGAGED: + _ENGAGED.add(("pre_tier", rows)) + logger.info( + "NVIDIA megakernel: Qwen3.5 MoE front tier %s at %d rows: " + "route_topk tables%s", + plan.moe_front_tier(rows), + rows, + " (route_gemm front)" if gemm else "", + ) + return _pre_body(mlp, width, x_fp4, x_sf) topk_output = StandardTopKOutput( topk_weights=weights, topk_ids=ids, router_logits=logits ) diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index f812ba60..135cfa26 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -8,10 +8,12 @@ captures up to 128); a width of more than one request is covered only where a slot runs. ``route`` (the router GEMM and the routed experts' NVFP4 input in one launch) runs up to 16 rows; ``route_topk`` materializes the top-10 and NVFP4 input behind SGLang's router -GEMM at wider rows. ``mkm`` (on route's outputs, the top-10 and both expert GEMMs in one -launch) runs up to 7 rows, and with ``mkm_wide`` (its 16-token engine) up to 16. -``mkm_m32`` runs the same routed-expert arithmetic on route_topk's 17–32 rows. -``gdn_verify`` (the +GEMM at wider rows. Under TEP8, the internal routed-body plan uses route plus a +routing-only top-k launch through 16 rows, the low-M shared expert at 21/28 rows, and +the clustered route-big front from 33 through 896 rows. ``mkm`` (on route's outputs, +the top-10 and both expert GEMMs in one launch) +runs up to 7 rows, and with ``mkm_wide`` (its 16-token engine) up to 16. ``mkm_m32`` runs +the same routed-expert arithmetic on route_topk's 17–32 rows. ``gdn_verify`` (the GatedDeltaNet verify chain with its launch choices changed) covers every width, with ``gdn_core`` (numeric: its ReplaySSM verify, gated norm and FP8 output in one launch) up to 224 rows, and ``hop`` (the boundary's fused all-reduce points on the decode hop) the @@ -81,6 +83,30 @@ LOWM_OUTPROJ_MAX_ROWS = 14 LOWM_OPROJ_MAX_ROWS = 14 LOWM_SE_MAX_ROWS = 21 +LOWM_SE_PRE_MAX_ROWS = 28 +# TEP8's route-big front uses thread-block clusters to build the routed-body tables. +ROUTE_BIG_MIN_ROWS = 33 +# The formal TEP8 route-big path covers every L-tier graph. Past 96 requests it uses +# TRT-LLM r3's 64-CTA launch instead of the 32-CTA launch that overlaps the shared expert. +ROUTE_BIG_MAX_ROWS = MAX_ROWS +ROUTE_BIG_WIDE_MIN_ROWS, ROUTE_BIG_WIDE_CTAS = 673, 64 +# The TEP8 MoE front's row tiers (each with its TRT-LLM r3 path; PLAN.md section 4): name, widest step +MOE_FRONT_TIERS = (("S1", 8), ("S2", 16), ("M", 32), ("L1", 224), ("L2", MAX_ROWS)) +# TEP8 tiers S1/S2: route's front followed by a routing-only top-k launch. +SFRONT_TOPK_MAX_ROWS = ROUTE_MAX_ROWS + + +def moe_front_tier(rows: int) -> str: + """The TEP8 MoE front tier of a step of ``rows`` rows.""" + return next( + (name for name, top in MOE_FRONT_TIERS if rows <= top), MOE_FRONT_TIERS[-1][0] + ) + + +def lowm_se_pre_rows(tp_size: int, tep8_m: bool = False) -> int: + """The shared expert's widest tier-M step under the TEP8 front policy.""" + tep8 = tep8_m and tp_size == 8 + return LOWM_SE_PRE_MAX_ROWS if tep8 else 0 class Slots(msgspec.Struct, frozen=True): @@ -90,6 +116,10 @@ class Slots(msgspec.Struct, frozen=True): # behind SGLang's router GEMM: top-10 + NVFP4 input in one launch (moe/route.py) route_topk: bool = False route_gemm: bool = False # with the router front's tiles in place of that GEMM + # Internal TEP8 path: the MoE body on route_topk/route-big tables (moe/body.py) + route_pre: bool = False + # Internal TEP8 S1/S2 path: route followed by routing-only route_topk tables + route_pre_topk: bool = False mkm: bool = ( False # top-10 + FC1 + SwiGLU + FC2 on route's outputs (sm103/moe/mkm.py) ) @@ -140,6 +170,7 @@ class Facts(msgspec.Struct, frozen=True): False # the top-10 + NVFP4 input launch built for every MoE layer ) route_gemm: bool = False # and the router front's tiles + tep8_moe: bool = False # the internal S/M/L routing-table and routed-body plan mkm: bool = False # the fused routed experts built for every MoE layer mkm_wide: bool = False # with their 16-token engine mkm_m32: bool = False # the 32-token engine built for every MoE layer @@ -166,6 +197,7 @@ class Facts(msgspec.Struct, frozen=True): lowm_in: bool = False # the low-M input projections built for every layer lowm_out: bool = False # the low-M output projections built for every layer lowm_se: bool = False # the low-M shared expert built for every MoE layer + lowm_se_pre: bool = False # also at route_pre's tier-M rows lowm_push: bool = False # with the output projections' push into the hop hop_fp8: bool = False # the hop's FP8 copy of the norm for the next projection @@ -175,12 +207,13 @@ def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> raise ValueError(f"{rows} rows are outside the covered widths 1-{MAX_ROWS}") route = facts.route and rows <= ROUTE_MAX_ROWS route_topk = facts.route_topk and not route and rows <= ROUTE_TOPK_MAX_ROWS + tep8 = facts.tep8_moe hop = facts.hop and rows <= HOP_MAX_ROWS mka = facts.mka and rows <= (MKA_WIDE_MAX_ROWS if facts.mka_wide else MKA_MAX_ROWS) - lowm_se = facts.lowm_se and rows <= LOWM_SE_MAX_ROWS mkm = ( route and facts.mkm + and not tep8 and rows <= (MKM_WIDE_MAX_ROWS if facts.mkm_wide else MKM_MAX_ROWS) ) gdn_front_base = facts.gdn_verify and facts.gdn_pdl and facts.gdn_front @@ -190,29 +223,51 @@ def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> and rows % GDN_FRONT_ROWS == 0 and rows <= GDN_TRT_FRONT_MAX_ROWS ) + pre_front_topk = route and tep8 and not mkm and rows <= SFRONT_TOPK_MAX_ROWS + pre_low = route_topk and tep8 and facts.lowm_se_pre and rows <= LOWM_SE_PRE_MAX_ROWS + lowm_se = facts.lowm_se and (rows <= LOWM_SE_MAX_ROWS or pre_low) # At one request lowm_se replaces both route_shared's shared stage and # SGLang's shared-expert tail. route_shared = ( - route and facts.route_shared and rows <= ROUTE_SHARED_ROWS and not lowm_se + route + and facts.route_shared + and rows <= ROUTE_SHARED_ROWS + and not lowm_se + and not pre_front_topk + ) + pre_topk = pre_low or ( + route_topk and tep8 and ROUTE_BIG_MIN_ROWS <= rows <= ROUTE_BIG_MAX_ROWS ) return Slots( route=route, route_topk=route_topk, - route_gemm=route_topk and facts.route_gemm and rows <= ROUTE_GEMM_MAX_ROWS, + # TEP8's M/L bodies keep SGLang's router GEMM; route_topk or route-big + # handles their routing and quantization. + route_gemm=route_topk + and facts.route_gemm + and rows <= ROUTE_GEMM_MAX_ROWS + and not pre_topk, + route_pre=pre_topk, + route_pre_topk=pre_front_topk, mkm=mkm, mkm_m32=route_topk + and not tep8 and facts.mkm_m32 and MKM_WIDE_MAX_ROWS < rows <= MKM_M32_MAX_ROWS, mkm_m64=route_topk + and not tep8 and facts.mkm_m64 and MKM_M32_MAX_ROWS < rows <= MKM_M64_MAX_ROWS, mkm_m128=route_topk + and not tep8 and facts.mkm_m128 and MKM_M64_MAX_ROWS < rows <= MKM_M128_MAX_ROWS, mkm_m256=route_topk + and not tep8 and facts.mkm_m256 and MKM_M128_MAX_ROWS < rows <= MKM_M256_MAX_ROWS, mkm_m512=route_topk + and not tep8 and facts.mkm_m512 and MKM_M256_MAX_ROWS < rows <= MKM_M512_MAX_ROWS, gdn_verify=facts.gdn_verify, diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index d7ded9f3..2d2aab60 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -355,10 +355,13 @@ def build(model, ctx): "attention": attention.build(model, pctx), "lowm": lowm.build(model, pctx), } + parallel = get_parallel() + tep8_moe = bound["moe"].pre and parallel.tp_size == 8 and parallel.moe_ep_size > 1 facts = slot_plan.Facts( route=bound["moe"].route, route_topk=bound["moe"].topk, route_gemm=bound["moe"].gemm, + tep8_moe=tep8_moe, mkm=bool(bound["moe"].mkm), mkm_wide=bound["moe"].mkm_wide, mkm_m32=bool(bound["moe"].m32) and "mkm_m32" in SLOTS, @@ -386,6 +389,7 @@ def build(model, ctx): lowm_in=bound["lowm"].lowm_in, lowm_out=bound["lowm"].lowm_out, lowm_se=bound["lowm"].se, + lowm_se_pre=bound["lowm"].se_pre, lowm_push=bound["lowm"].push, hop_fp8=bound["lowm"].hop_fp8, ) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py b/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py index 7ee38653..d8194c92 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/router.py @@ -7,9 +7,8 @@ quantization equals the runner's own NVFP4 input quantizer (16-element blocks, e4m3 scales, linear layout). ``router_gemm_quant_shared`` also computes the shared expert's static-FP8 gate/up -projection of up to 8 rows in the same launch (SGLANG_Q35MF_ROUTE_SHARED on SGLang branch -nvidia/qwen35-agentx-gb300); it sums in another order than the FP8 linear's GEMM -(numeric). +projection of up to 8 rows in the same launch; it sums in another order than the FP8 +linear's GEMM (numeric). """ from __future__ import annotations @@ -17,12 +16,14 @@ import ctypes from functools import cache from pathlib import Path +from typing import Optional import torch from sgl_nvidia_mega.common import jit SOURCE = Path(__file__).parent / "src" / "router_front.cu" +TABLES_HEADER = Path(__file__).parent / "src" / "route_tables.cuh" HIDDEN, EXPERTS, MAX_TOKENS = 4096, 512, 16 THREADS, CTAS = 256, 64 # 8 experts a CTA _ROW_BYTES = HIDDEN * 2 + 16 @@ -38,6 +39,32 @@ KSPLIT = 4 +class TablesParams(ctypes.Structure): + """q35_route_tables::Tables (src/route_tables.cuh).""" + + _fields_ = [ + (name, ctypes.c_void_p) + for name in ( + "counter", + "total_padded", + "expanded_to_permuted", + "permuted_to_token", + "tokens_per_expert", + "cta_expert", + "cta_row_limit", + "live_ctas", + ) + ] + [ + ("tile", ctypes.c_int32), + ("local_offset", ctypes.c_int32), + ("local_count", ctypes.c_int32), + ] + + +def _headers() -> dict: + return {"route_tables.cuh": TABLES_HEADER.read_text()} + + def max_m_bucket(m: int) -> int: """The compiled row capacity that serves m rows.""" return 1 if m <= 1 else 8 if m <= 8 else 16 @@ -97,6 +124,7 @@ def _cubin(device_index: int, max_m: int, ksplit: int = 2) -> bytes: f"sm_{major}{minor}a", ("-std=c++17", "-default-device", f"-DQ35_ROUTER_MAX_M={max_m}") + ((f"-DQ35_ROUTER_KSPLIT={ksplit}",) if ksplit != 2 else ()), + headers=_headers(), ) @@ -130,7 +158,9 @@ def build(device_index: int, max_rows: int, shared: bool = False) -> None: def router_gemm_quant( - x: torch.Tensor, weight: torch.Tensor, global_scale: torch.Tensor + x: torch.Tensor, + weight: torch.Tensor, + global_scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """x [M, 4096] bf16 (M <= 16), weight [512, 4096] bf16, global_scale [1] fp32: (logits bf16 [M, 512], x_fp4 uint8 [M, 2048], x_sf uint8 [M, 256]).""" @@ -155,8 +185,12 @@ def router_gemm_quant( num_tokens=m, mode=MODE_QUANT, ) - kernel = _kernel(x.device.index, max_m_bucket(m)) - jit.launch(kernel, (CTAS, 1, 1), [p], torch.cuda.current_stream().cuda_stream) + jit.launch( + _kernel(x.device.index, max_m_bucket(m)), + (CTAS, 1, 1), + [p], + torch.cuda.current_stream().cuda_stream, + ) return logits, x_fp4, x_sf @@ -231,7 +265,84 @@ class TopkParams(ctypes.Structure): "x_fp4", "x_sf", ) - ] + ] + [("tables", TablesParams), ("num_tokens", ctypes.c_int32)] + + +def max_ctas(m: int, tile: int) -> int: + """The MoE body's CTA bound for m tokens (trtllm-gen runner.h getMaxNumCtasInBatchDim).""" + expanded = m * TOP_K + filled = min(EXPERTS, expanded) + return filled + (expanded - filled) // tile + + +class RouteTables: + """Graph-stable trtllm-gen routing tables of one width and tile_n for the local experts + [local_offset, local_offset + local_count) (expert parallelism: the runner's + local_expert_offset / local_num_experts), which route_topk_quant or route-big writes + for the MoE body; ``metadata()`` is + FlashInfer's nine-tensor RoutingMetadataBuffers order. Also the width's top-10 buffers, + which the tables borrow. A width's launches run one after another (the arrival counter + is the width's).""" + + def __init__( + self, + m: int, + tile: int, + device: torch.device, + local_offset: int = 0, + local_count: int = EXPERTS, + ): + i32 = dict(dtype=torch.int32, device=device) + ctas = max_ctas(m, tile) + if not ( + 0 <= local_offset + and 0 < local_count + and local_offset + local_count <= EXPERTS + ): + raise ValueError( + f"route tables: local experts {local_offset} + {local_count}" + ) + self.m, self.tile = m, tile + self.local_offset, self.local_count = local_offset, local_count + self.ids = torch.empty((m, TOP_K), **i32) + self.weights = torch.empty((m, TOP_K), dtype=torch.bfloat16, device=device) + self.counter = torch.zeros(1, **i32) + self.total_padded = torch.zeros(1, **i32) + self.expanded_to_permuted = torch.zeros(m * TOP_K, **i32) + self.permuted_to_token = torch.full((ctas * tile + 1,), -1, **i32) + self.histogram = torch.zeros(2 * EXPERTS, **i32) # FlashInfer's scratch; unused + self.tokens_per_expert = torch.zeros(EXPERTS, **i32) + self.cta_expert = torch.zeros(ctas, **i32) + self.cta_row_limit = torch.zeros(ctas, **i32) + self.live_ctas = torch.zeros(1, **i32) + + def params(self) -> TablesParams: + return TablesParams( + counter=self.counter.data_ptr(), + total_padded=self.total_padded.data_ptr(), + expanded_to_permuted=self.expanded_to_permuted.data_ptr(), + permuted_to_token=self.permuted_to_token.data_ptr(), + tokens_per_expert=self.tokens_per_expert.data_ptr(), + cta_expert=self.cta_expert.data_ptr(), + cta_row_limit=self.cta_row_limit.data_ptr(), + live_ctas=self.live_ctas.data_ptr(), + tile=self.tile, + local_offset=self.local_offset, + local_count=self.local_count, + ) + + def metadata(self) -> list[torch.Tensor]: + return [ + self.total_padded, + self.expanded_to_permuted, + self.permuted_to_token, + self.weights, + self.histogram, + self.tokens_per_expert, + self.cta_expert, + self.cta_row_limit, + self.live_ctas, + ] @cache @@ -242,6 +353,7 @@ def _topk_kernel(device_index: int) -> jit.Kernel: "q35_route_topk", f"sm_{major}{minor}a", ("-std=c++17", "-default-device"), + headers=_headers(), ) return jit.load(cubin, "q35_route_topk_quant", (THREADS, 1, 1), 0) @@ -256,12 +368,14 @@ def route_topk_quant( x: torch.Tensor, global_scale: torch.Tensor, quantize: bool = True, + tables: Optional[RouteTables] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """logits [M, 512] bf16 (the router GEMM's), x [M, 4096] bf16, global_scale [1] fp32: (topk_ids int32 [M, 10], topk_weights bf16 [M, 10], x_fp4 uint8 [M, 2048], x_sf uint8 [M, 256]); the routing equals the trtllm-gen RenormalizeNaive routing's and the input the runner's NVFP4 quantizer's, bit for bit. Without ``quantize`` only - the routing (x_fp4 / x_sf None).""" + the routing (x_fp4 / x_sf None). With ``tables`` (of M rows) the top-10 goes to its + buffers and the launch also writes its routing tables.""" m = x.shape[0] for t, shape in ((logits, (m, EXPERTS)), (x, (m, HIDDEN))): if ( @@ -272,8 +386,13 @@ def route_topk_quant( raise ValueError(f"route_topk: {t.dtype} {tuple(t.shape)}, want {shape}") if not 1 <= m <= TOPK_MAX_ROWS: raise ValueError(f"route_topk: {m} rows (1-{TOPK_MAX_ROWS})") - ids = torch.empty((m, TOP_K), dtype=torch.int32, device=x.device) - weights = torch.empty((m, TOP_K), dtype=torch.bfloat16, device=x.device) + if tables is None: + ids = torch.empty((m, TOP_K), dtype=torch.int32, device=x.device) + weights = torch.empty((m, TOP_K), dtype=torch.bfloat16, device=x.device) + elif tables.m != m: + raise ValueError(f"route_topk: tables of {tables.m} rows for {m}") + else: + ids, weights = tables.ids, tables.weights x_fp4 = x_sf = None if quantize: x_fp4 = torch.empty((m, HIDDEN // 2), dtype=torch.uint8, device=x.device) @@ -286,7 +405,10 @@ def route_topk_quant( topk_weights=weights.data_ptr(), x_fp4=jit.null_or_pointer(x_fp4).value, x_sf=jit.null_or_pointer(x_sf).value, + num_tokens=m, ) + if tables is not None: + p.tables = tables.params() jit.launch( _topk_kernel(x.device.index), (m, 1, 1), @@ -296,6 +418,104 @@ def route_topk_quant( return ids, weights, x_fp4, x_sf +BIG_SOURCE = Path(__file__).parent / "src" / "route_big.cu" +# TRT-LLM r3's route-only big front: one routing cluster of 8 CTAs (up to +# 256 rows) or 16 (non-portable, up to 1024), plus quantization clusters. +BIG_THREADS, BIG_CTAS = 1024, 32 + + +def big_shape(m: int) -> tuple[int, int, int]: + """(CTAs a cluster, tokens a routing CTA, token-mask words an expert) of route_big at m rows.""" + cl = 8 if m <= 256 else 16 + tpc = -(-m // cl) + return cl, tpc, -(-tpc // 32) + + +class BigParams(ctypes.Structure): + """q35_route_big::Params.""" + + _fields_ = TopkParams._fields_ + [("tpc", ctypes.c_int32)] + + +@cache +def _big_kernel(device_index: int, cl: int, mw: int, nl: int) -> jit.Kernel: + major, minor = torch.cuda.get_device_capability(device_index) + defines = f"#define CL {cl}\n#define MW {mw}\n#define NL {nl}\n" + cubin = jit.compile_source( + defines + BIG_SOURCE.read_text(), + "q35_route_big", + f"sm_{major}{minor}a", + ("-std=c++17", "-default-device"), + headers=_headers(), + ) + kernel = jit.load(cubin, "q35_route_big_quant", (BIG_THREADS, 1, 1), 0) + if cl > 8: + cu = jit.driver() + jit.check( + cu.cuFuncSetAttribute( + kernel.function, + cu.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, + 1, + ), + "cuFuncSetAttribute", + ) + return kernel + + +def build_big(device_index: int, local_count: int) -> None: + """Compile route_big's launches for every wide step shape (no compile inside a step).""" + for cl, mw in ((8, 1), (16, 1), (16, 2)): + _big_kernel(device_index, cl, mw, local_count) + + +def route_big_quant( + logits: torch.Tensor, + x: torch.Tensor, + global_scale: torch.Tensor, + tables: RouteTables, + ctas: Optional[int] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """route_topk_quant(logits, x, global_scale, tables=tables) as one launch of thread-block + clusters (TRT-LLM's route-only big front): the same top-10 and NVFP4 input bit for bit, and + the same routing tables up to the order of an expert's rows (token order here).""" + m = x.shape[0] + for t, shape in ((logits, (m, EXPERTS)), (x, (m, HIDDEN))): + if ( + t.dtype != torch.bfloat16 + or tuple(t.shape) != shape + or not t.is_contiguous() + ): + raise ValueError(f"route_big: {t.dtype} {tuple(t.shape)}, want {shape}") + if tables.m != m or not 1 <= m <= 1024: + raise ValueError(f"route_big: tables of {tables.m} rows for {m}") + if tables.local_count % 32 or tables.local_count > BIG_THREADS: + raise ValueError(f"route_big: {tables.local_count} local experts") + cl, tpc, mw = big_shape(m) + x_fp4 = torch.empty((m, HIDDEN // 2), dtype=torch.uint8, device=x.device) + x_sf = torch.empty((m, HIDDEN // 16), dtype=torch.uint8, device=x.device) + p = BigParams( + logits=logits.data_ptr(), + x=x.data_ptr(), + global_scale=global_scale.data_ptr(), + topk_ids=tables.ids.data_ptr(), + topk_weights=tables.weights.data_ptr(), + x_fp4=x_fp4.data_ptr(), + x_sf=x_sf.data_ptr(), + num_tokens=m, + tpc=tpc, + ) + p.tables = tables.params() + launch_ctas = max(2 * cl, (BIG_CTAS if ctas is None else ctas) // cl * cl) + jit.launch( + _big_kernel(x.device.index, cl, mw, tables.local_count), + (launch_ctas, 1, 1), + [p], + torch.cuda.current_stream().cuda_stream, + cluster=(cl, 1, 1), + ) + return tables.ids, tables.weights, x_fp4, x_sf + + GEMM_MAX_ROWS = 56 # slot route_gemm: the router front in tiles of 16 rows up to here diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_big.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_big.cu new file mode 100644 index 00000000..b15e3a1a --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_big.cu @@ -0,0 +1,313 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Adapted from TensorRT-LLM (Apache-2.0, NVIDIA): tensorrt_llm/_torch/cute_dsl_kernels/blackwell/q35_moe_front_big.py +// (Q35MoeRouteBigKernel, the route-only big front of TRTLLM_Q35_MOE_FRONT_BIG: launch shape, token split, per-expert +// token masks, the cluster's counts through distributed shared memory, the table writes). +// +// TEP8's front at wide verify steps: behind SGLang's router GEMM, the top-10 +// of every token's 512 bf16 logits, the trtllm-gen MoE body's routing tables for the rank's window of local experts +// and the NVFP4 quantization of the MoE input, in one launch of thread-block clusters of CL CTAs: +// +// * cluster 0 routes: CTA c takes tokens [c TPC, (c + 1) TPC), one warp a token (two, interleaved, past 32), with +// route_tables.cuh's routing (FlashInfer's trtllm-gen bits, as slot route_topk); a local pick sets the token's bit +// in the CTA's per-expert token mask (shared memory, MW words an expert). After a cluster barrier every routing CTA +// reads the cluster's per-expert counts through distributed shared memory, scans the CTAs per expert (tile) and +// writes the expanded -> permuted / permuted -> token entries of its own picks (an expert's rows in token order), +// the CTA tables and padding rows (-1) of every CL-th expert; CTA 0 also the totals and the tokens per expert. +// * the other clusters quantize x (route_topk.cu's quantizer, bit for bit). +// +// The tables are route_tables.cuh's write_tables' (the same counts, CTA tables and totals; an expert's rows in token +// order where write_tables takes the order of shared-memory atomics, which the body's output does not depend +// on). No CTA waits for another cluster. PDL: every read after griddepcontrol.wait (the global scale, static, before); +// dependents released after a CTA barrier that follows the CTA's last global store. +// +// Defines: CL (CTAs a cluster: 8, or 16, non-portable), MW (token-mask words an expert: a routing CTA takes up to +// 32 MW <= 64 tokens, Params.tpc), NL (local experts, a multiple of 32, <= 1024). Built with NVRTC. + +#include +#include +#include + +// NVRTC has no +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef int int32_t; +typedef unsigned int uint32_t; + +#include "route_tables.cuh" + +namespace q35_route_big { + +using namespace q35_route_tables; + +constexpr uint32_t kThreadsBig = 1024, kWarpsBig = kThreadsBig / 32; +constexpr uint32_t kHidden = 4096, kSfVec = 16, kBlocksPerRow = kHidden / kSfVec; // 256 +constexpr uint32_t kMaxTpc = 32 * MW; // tokens a routing CTA takes at most +static_assert(MW >= 1 && kMaxTpc <= 2 * kWarpsBig, "at most two tokens a warp"); +static_assert(NL % 32 == 0 && NL <= kThreadsBig, "one thread an expert"); +static_assert(NL * MW <= kThreadsBig, "one thread a token-mask word"); + +struct Params { + const __nv_bfloat16* logits; // [M, 512] + const __nv_bfloat16* x; // [M, 4096] + const float* global_scale; // [1] + int32_t* topk_ids; // [M, 10] + uint16_t* topk_weights; // [M, 10] bf16 + uint8_t* x_fp4; // [M, 2048] + uint8_t* x_sf; // [M, 256] + Tables tables; // the counter is unused + int32_t num_tokens; // M + int32_t tpc; // tokens a routing CTA: ceil(M / CL) <= 32 MW +}; + +// ---- the NVFP4 quantizer (route_topk.cu, router_front.cu) +__device__ __forceinline__ float bf16_lo(uint32_t v) { return __uint_as_float(v << 16); } +__device__ __forceinline__ float bf16_hi(uint32_t v) { return __uint_as_float(v & 0xFFFF0000u); } +__device__ __forceinline__ float rcp_approx_ftz(float a) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(a)); + return y; +} +__device__ __forceinline__ uint32_t cvt_e4m3(float a) { + uint32_t out; + asm("{\n.reg .b16 pair;\n.reg .f32 zero;\nmov.f32 zero, 0f00000000;\n" + "cvt.rn.satfinite.e4m3x2.f32 pair, zero, %1;\ncvt.u32.u16 %0, pair;\n}" + : "=r"(out) + : "f"(a)); + return out & 0xFFu; +} +__device__ __forceinline__ float e4m3_to_f32(uint32_t v) { + float out; + asm("{\n.reg .b16 pair;\n.reg .b32 h2;\n.reg .b16 lo, hi;\ncvt.u16.u32 pair, %1;\n" + "cvt.rn.f16x2.e4m3x2 h2, pair;\nmov.b32 {lo, hi}, h2;\ncvt.f32.f16 %0, lo;\n}" + : "=f"(out) + : "r"(v)); + return out; +} +__device__ __forceinline__ uint32_t cvt_e2m1x8(const float (&v)[8]) { + uint32_t out; + asm("{\n.reg .b8 b0, b1, b2, b3;\n" + "cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\ncvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\ncvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" + "mov.b32 %0, {b0, b1, b2, b3};\n}" + : "=r"(out) + : "f"(v[0]), "f"(v[1]), "f"(v[2]), "f"(v[3]), "f"(v[4]), "f"(v[5]), "f"(v[6]), "f"(v[7])); + return out; +} +__device__ __forceinline__ uint32_t bf16x2_abs_max(uint32_t a, uint32_t b) { + uint32_t out; + asm("{\n.reg .b32 ta, tb;\nand.b32 ta, %1, 0x7FFF7FFF;\nand.b32 tb, %2, 0x7FFF7FFF;\nmax.bf16x2 %0, ta, tb;\n}" + : "=r"(out) + : "r"(a), "r"(b)); + return out; +} +__device__ __forceinline__ uint32_t bf16x2_max(uint32_t a, uint32_t b) { + uint32_t out; + asm("max.bf16x2 %0, %1, %2;" : "=r"(out) : "r"(a), "r"(b)); + return out; +} +__device__ __forceinline__ void quantize_block(const uint4& a, const uint4& b, float gs, uint32_t& sf, uint2& packed) { + const uint32_t w[8] = {a.x, a.y, a.z, a.w, b.x, b.y, b.z, b.w}; + const uint32_t m = bf16x2_max(bf16x2_max(bf16x2_abs_max(w[0], w[1]), bf16x2_abs_max(w[2], w[3])), + bf16x2_max(bf16x2_abs_max(w[4], w[5]), bf16x2_abs_max(w[6], w[7]))); + const float amax = fmaxf(bf16_lo(m), bf16_hi(m)); + const float scale = __fmul_rn(gs, __fmul_rn(amax, rcp_approx_ftz(6.0f))); + sf = cvt_e4m3(scale); + const float sf_f = e4m3_to_f32(sf); + const float out_scale = (sf_f == 0.0f) ? 0.0f : rcp_approx_ftz(__fmul_rn(sf_f, rcp_approx_ftz(gs))); + float lo[8], hi[8]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + lo[2 * i] = __fmul_rn(bf16_lo(w[i]), out_scale); + lo[2 * i + 1] = __fmul_rn(bf16_hi(w[i]), out_scale); + hi[2 * i] = __fmul_rn(bf16_lo(w[4 + i]), out_scale); + hi[2 * i + 1] = __fmul_rn(bf16_hi(w[4 + i]), out_scale); + } + packed.x = cvt_e2m1x8(lo); + packed.y = cvt_e2m1x8(hi); +} + +// ---- cluster primitives +__device__ __forceinline__ void cluster_arrive_release() { + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); +} +__device__ __forceinline__ void cluster_arrive_relaxed() { + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} +__device__ __forceinline__ void cluster_wait_acquire() { + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); +} +// the word at this CTA's shared address `p` in cluster CTA `rank` +__device__ __forceinline__ int32_t ld_dsmem(const int32_t* p, uint32_t rank) { + const uint32_t a = static_cast(__cvta_generic_to_shared(p)); + int32_t v; + // no memory clobber: the cluster barrier before (which has one) orders these reads; without it the peers' + // loads issue back to back instead of one round trip each + asm volatile("{\n.reg .b32 ra;\nmapa.shared::cluster.u32 ra, %1, %2;\nld.shared::cluster.b32 %0, [ra];\n}" + : "=r"(v) + : "r"(a), "r"(rank)); + return v; +} + +// one token's logits (lane l: experts 16 l .. 16 l + 15) +__device__ __forceinline__ void load_logits(const Params& p, int32_t token, uint32_t lane, uint4& v0, uint4& v1) { + const uint4* src = reinterpret_cast(p.logits + static_cast(token) * kExperts + lane * 16); + v0 = __ldcg(src); + v1 = __ldcg(src + 1); +} + +__device__ __forceinline__ void route_cta(const Params& p, uint32_t c) { + __shared__ uint32_t s_mask[NL * MW]; + __shared__ int32_t s_cnt[NL], s_pre[NL], s_offs[NL], s_tot[NL], s_wsum[kWarpsBig]; + __shared__ uint16_t s_pick[kMaxTpc * kTopK]; + __shared__ __align__(16) uint32_t s_cand[kWarpsBig][32]; + const uint32_t tid = threadIdx.x, warp = tid / 32, lane = tid % 32; + const Tables& t = p.tables; + const int32_t tile = t.tile; + if (tid < NL * MW) s_mask[tid] = 0u; + const int32_t t_first = static_cast(c) * p.tpc; + const int32_t n_tok = min(p.tpc, p.num_tokens - t_first); + asm volatile("griddepcontrol.wait;" ::: "memory"); + __syncthreads(); + // ---- the routing: warp w takes tokens w and w + 32 of this CTA + const bool two = static_cast(warp + kWarpsBig) < n_tok; + if (static_cast(warp) < n_tok) { + uint32_t expert[2] = {0u, 0u}; + uint16_t wbits[2] = {0, 0}; + // TRT-LLM's candidate top-10 (route_row_fast: the same bits as route_row) + uint4 v[2][2]; + load_logits(p, t_first + static_cast(warp), lane, v[0][0], v[0][1]); + if (two) load_logits(p, t_first + static_cast(warp + kWarpsBig), lane, v[1][0], v[1][1]); + route_row_fast(v[0][0], v[0][1], lane, s_cand[warp], expert[0], wbits[0]); + if (two) route_row_fast(v[1][0], v[1][1], lane, s_cand[warp], expert[1], wbits[1]); +#pragma unroll + for (int j = 0; j < 2; ++j) { + if (j == 1 && !two) break; + const uint32_t tl = warp + j * kWarpsBig; // the token within the CTA + if (lane < kTopK) { + const size_t o = static_cast(t_first + static_cast(tl)) * kTopK + lane; + p.topk_ids[o] = static_cast(expert[j]); + p.topk_weights[o] = wbits[j]; + const uint16_t l = local_slot(static_cast(expert[j]), t); + s_pick[tl * kTopK + lane] = l; + if (l != kRemote) atomicOr(&s_mask[l * MW + tl / 32], 1u << (tl % 32)); + } + } + } + __syncthreads(); + if (tid < NL) { + int32_t cnt = 0; +#pragma unroll + for (uint32_t w = 0; w < MW; ++w) cnt += __popc(s_mask[tid * MW + w]); + s_cnt[tid] = cnt; + } + // ---- the cluster's per-expert counts through distributed shared memory + cluster_arrive_release(); + cluster_wait_acquire(); + int32_t tot = 0, nct = 0; + if (tid < NL) { + int32_t v[CL]; +#pragma unroll + for (uint32_t k = 0; k < CL; ++k) v[k] = ld_dsmem(&s_cnt[tid], k); // every peer's count in flight at once + int32_t pre = 0; +#pragma unroll + for (uint32_t k = 0; k < CL; ++k) { + tot += v[k]; + pre += k < c ? v[k] : 0; + } + s_pre[tid] = pre; + nct = (tot + tile - 1) / tile; + } + cluster_arrive_relaxed(); // this CTA has read its peers' counts (the matching wait is at the end) + // exclusive scan of the CTAs per expert, one expert a thread + int32_t incl = nct; +#pragma unroll + for (uint32_t d = 1; d < 32; d <<= 1) { + const int32_t up = __shfl_up_sync(0xffffffffu, incl, d); + incl += lane >= d ? up : 0; + } + if (lane == 31) s_wsum[warp] = incl; + __syncthreads(); + int32_t wpre = 0, total = 0; +#pragma unroll + for (uint32_t w = 0; w < NL / 32; ++w) { + const int32_t ws = s_wsum[w]; + wpre += w < warp ? ws : 0; + total += ws; + } + const int32_t offs = incl - nct + wpre; + if (tid < NL) { + s_offs[tid] = offs; + s_tot[tid] = tot; + if (c == 0) t.tokens_per_expert[tid] = tot; + if (tid % CL == c) { // expert tid's CTA tables from routing CTA tid % CL + for (int32_t cc = 0; cc < nct; ++cc) { + const int32_t cta = offs + cc; + t.cta_expert[cta] = static_cast(tid); + t.cta_row_limit[cta] = min((cta + 1) * tile, offs * tile + tot); + } + } + } + if (c == 0 && tid == 0) { + t.live_ctas[0] = total; + t.total_padded[0] = total * tile; + } + __syncthreads(); + // the padding rows (-1) past an expert's last row: a warp an expert, spread over the routing CTAs + for (uint32_t e = c + CL * warp; e < NL; e += CL * kWarpsBig) { + const int32_t n = s_tot[e], first = s_offs[e] * tile; + for (int32_t r = first + n + static_cast(lane); r < first + (n + tile - 1) / tile * tile; r += 32) { + t.permuted_to_token[r] = -1; + } + } + // ---- this CTA's picks: expanded -> permuted, permuted -> token (an expert's rows in token order) + for (int32_t i = static_cast(tid); i < n_tok * static_cast(kTopK); i += kThreadsBig) { + const uint32_t tl = static_cast(i) / kTopK; + const int32_t tok = t_first + static_cast(tl); + const uint16_t l = s_pick[i]; + int32_t perm = -1; + if (l != kRemote) { + int32_t rank = s_pre[l]; +#pragma unroll + for (uint32_t w = 0; w < MW; ++w) { + const uint32_t word = s_mask[l * MW + w]; + const uint32_t lo = 32 * w; + rank += tl >= lo + 32 ? __popc(word) : (tl > lo ? __popc(word & ((1u << (tl - lo)) - 1u)) : 0); + } + perm = s_offs[l] * tile + rank; + t.permuted_to_token[perm] = tok; + } + t.expanded_to_permuted[static_cast(tok) * kTopK + (static_cast(i) - tl * kTopK)] = perm; + } + cluster_wait_acquire(); // every peer has read this CTA's counts + __syncthreads(); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +} + +__device__ __forceinline__ void quant_cta(const Params& p, uint32_t qc, uint32_t q_ctas) { + const float gs = p.global_scale[0]; // static (the experts' input scale) + asm volatile("griddepcontrol.wait;" ::: "memory"); + const uint32_t blocks = static_cast(p.num_tokens) * kBlocksPerRow; + for (uint32_t blk = qc * kThreadsBig + threadIdx.x; blk < blocks; blk += q_ctas * kThreadsBig) { + const uint32_t row = blk / kBlocksPerRow, cb = blk % kBlocksPerRow; + const uint4* src = reinterpret_cast(p.x + static_cast(row) * kHidden + cb * kSfVec); + uint32_t sf; + uint2 packed; + quantize_block(__ldcg(src), __ldcg(src + 1), gs, sf, packed); + *reinterpret_cast(p.x_fp4 + static_cast(row) * (kHidden / 2) + cb * 8) = packed; + p.x_sf[static_cast(row) * kBlocksPerRow + cb] = static_cast(sf); + } + __syncthreads(); + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +} + +} // namespace q35_route_big + +extern "C" __global__ __launch_bounds__(q35_route_big::kThreadsBig, 1) void q35_route_big_quant( + const __grid_constant__ q35_route_big::Params p) { + if (blockIdx.x < CL) { + q35_route_big::route_cta(p, blockIdx.x); + } else { + q35_route_big::quant_cta(p, blockIdx.x - CL, gridDim.x - CL); + } +} diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_tables.cuh b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_tables.cuh new file mode 100644 index 00000000..8a6ad2e9 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_tables.cuh @@ -0,0 +1,401 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// The trtllm-gen routing of one token's 512 bf16 router logits (top-10 expert ids and bf16 weights, the bits of +// FlashInfer's single-block routing, RenormalizeNaive) and the routing tables the trtllm-gen MoE body reads (what its +// routing kernels write, for the rank's window of local experts; a slot outside it gets permuted index -1). Shared by +// route_topk.cu and route_big.cu; the includer provides uint16_t, uint32_t, int32_t, cuda_bf16.h and +// cooperative_groups (with its reduce header). + +#pragma once + +#ifndef INFINITY +#define INFINITY __int_as_float(0x7f800000) +#endif + +namespace q35_route_tables { + +constexpr uint32_t kExperts = 512, kTopK = 10, kThreads = 256; + +// ---- the routing (mkm_moe_m16.cu): key = bf16 order above 65535 - expert (higher score first, smaller expert on ties) +__device__ __forceinline__ uint32_t topk_key(uint32_t bits, uint32_t expert) { + const uint32_t twiddled = bits ^ ((bits & 0x8000u) ? 0xFFFFu : 0x8000u); + return (twiddled << 16) | (65535u - expert); +} +__device__ __forceinline__ void topk_unkey(uint32_t key, uint32_t& bits, uint32_t& expert) { + expert = 65535u - (key & 0xFFFFu); + const uint32_t twiddled = key >> 16; + bits = twiddled ^ ((twiddled & 0x8000u) ? 0x8000u : 0xFFFFu); +} +__device__ __forceinline__ uint32_t max16(const uint32_t (&k)[16]) { + const uint32_t a = max(max(max(k[0], k[1]), max(k[2], k[3])), max(max(k[4], k[5]), max(k[6], k[7]))); + const uint32_t b = max(max(max(k[8], k[9]), max(k[10], k[11])), max(max(k[12], k[13]), max(k[14], k[15]))); + return max(a, b); +} +__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_sum_u32(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.add.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); + return r; +} +__device__ __forceinline__ uint32_t warp_topk_of_lanes(uint32_t v, uint32_t lane) { + uint32_t mine = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(v); + v = v == top ? 0u : v; + mine = lane == k ? top : mine; + } + return mine; +} +// Top-10 keys of 512 (16 per lane), lane k holding the k-th largest (keys unique and nonzero). +__device__ __forceinline__ uint32_t warp_top10_redux(uint32_t (&keys)[16], uint32_t lane, uint32_t* s_cand) { + uint32_t v = max16(keys); + uint32_t bound = 0; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + bound = warp_max_u32(v); + v = v == bound ? 0u : v; + } + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) count += keys[i] >= bound; + const uint32_t total = warp_sum_u32(count); + if (total <= 32) { + uint32_t offset = count; +#pragma unroll + for (uint32_t d = 1; d < 32; d <<= 1) { + const uint32_t up = __shfl_up_sync(0xffffffffu, offset, d); + offset += lane >= d ? up : 0u; + } + offset -= count; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) { + if (keys[i] >= bound) s_cand[offset++] = keys[i]; + } + __syncwarp(); + const uint32_t cand = lane < total ? s_cand[lane] : 0u; + __syncwarp(); + return warp_topk_of_lanes(cand, lane); + } + uint32_t mine = 0; +#pragma unroll 1 + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(max16(keys)); +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) keys[i] = keys[i] == top ? 0u : keys[i]; + mine = lane == k ? top : mine; + } + return mine; +} +// FlashInfer calcSoftmax (bf16 scores) as built with -use_fast_math. +__device__ __forceinline__ float fast_expf(float x) { + float y; + asm("{\n .reg .f32 t;\n mul.ftz.f32 t, %1, 0f3FB8AA3B;\n ex2.approx.ftz.f32 %0, t;\n}" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float fast_div(float a, float b) { + float y; + asm("div.full.ftz.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ uint16_t softmax_topk(uint16_t score_bits, uint32_t lane) { + namespace cg = cooperative_groups; + const auto warp = cg::tiled_partition<32>(cg::this_thread_block()); + const float si = __uint_as_float((uint32_t)score_bits << 16); + float max_score = -INFINITY; + if (lane < kTopK) max_score = si >= max_score ? si : max_score; + max_score = cg::reduce(warp, max_score, cg::greater()); + float sum_score = 0.f, new_score = 0.f; + if (lane < kTopK) { + new_score = fast_expf(__fsub_rn(si, max_score)); + sum_score = __fadd_rn(sum_score, new_score); + } + sum_score = cg::reduce(warp, sum_score, cg::plus()); + uint16_t out = score_bits; + if (lane < kTopK) { + const __nv_bfloat16 b = __float2bfloat16_rn(fast_div(new_score, sum_score)); + out = *reinterpret_cast(&b); + } + return out; +} + + +// One token's routing by one warp from its logits (lane l holds experts [16 l, 16 l + 16), loaded by the caller): lane +// k < 10 gets the k-th expert and its bf16 weight bits. +__device__ __forceinline__ void route_row( + const uint4& v0, const uint4& v1, uint32_t lane, uint32_t* s_cand, uint32_t& expert, uint16_t& wbits) { + uint32_t keys[16]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t w = i < 4 ? (&v0.x)[i] : (&v1.x)[i - 4]; + keys[2 * i] = topk_key(w & 0xFFFFu, lane * 16 + 2 * i); + keys[2 * i + 1] = topk_key(w >> 16, lane * 16 + 2 * i + 1); + } + const uint32_t key = warp_top10_redux(keys, lane, s_cand); + uint32_t bits; + topk_unkey(key, bits, expert); + wbits = softmax_topk(static_cast(bits), lane); +} + +// Two tokens' routing by one warp, their steps interleaved (each token's operations as route_row's); with more than +// a warp of candidates at or past either token's 10th key both take route_row. s_cand: [2][32]. +__device__ __forceinline__ void route_row2( + const uint4 (&v)[2][2], uint32_t lane, uint32_t* s_cand, uint32_t (&expert)[2], uint16_t (&wbits)[2]) { + uint32_t keys[2][16], mx[2], bound[2] = {0u, 0u}, count[2] = {0u, 0u}, total[2]; +#pragma unroll + for (int t = 0; t < 2; ++t) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t w = i < 4 ? (&v[t][0].x)[i] : (&v[t][1].x)[i - 4]; + keys[t][2 * i] = topk_key(w & 0xFFFFu, lane * 16 + 2 * i); + keys[t][2 * i + 1] = topk_key(w >> 16, lane * 16 + 2 * i + 1); + } + mx[t] = max16(keys[t]); + } +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { +#pragma unroll + for (int t = 0; t < 2; ++t) { + bound[t] = warp_max_u32(mx[t]); + mx[t] = mx[t] == bound[t] ? 0u : mx[t]; + } + } +#pragma unroll + for (int t = 0; t < 2; ++t) { +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) count[t] += keys[t][i] >= bound[t]; + total[t] = warp_sum_u32(count[t]); + } + if (total[0] > 32 || total[1] > 32) { +#pragma unroll + for (int t = 0; t < 2; ++t) route_row(v[t][0], v[t][1], lane, s_cand, expert[t], wbits[t]); + return; + } + uint32_t offset[2] = {count[0], count[1]}; +#pragma unroll + for (uint32_t d = 1; d < 32; d <<= 1) { +#pragma unroll + for (int t = 0; t < 2; ++t) { + const uint32_t up = __shfl_up_sync(0xffffffffu, offset[t], d); + offset[t] += lane >= d ? up : 0u; + } + } +#pragma unroll + for (int t = 0; t < 2; ++t) { + offset[t] -= count[t]; +#pragma unroll + for (uint32_t i = 0; i < 16; ++i) { + if (keys[t][i] >= bound[t]) s_cand[t * 32 + offset[t]++] = keys[t][i]; + } + } + __syncwarp(); + uint32_t cand[2], mine[2] = {0u, 0u}; +#pragma unroll + for (int t = 0; t < 2; ++t) cand[t] = lane < total[t] ? s_cand[t * 32 + lane] : 0u; + __syncwarp(); +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { +#pragma unroll + for (int t = 0; t < 2; ++t) { + const uint32_t top = warp_max_u32(cand[t]); + cand[t] = cand[t] == top ? 0u : cand[t]; + mine[t] = lane == k ? top : mine[t]; + } + } +#pragma unroll + for (int t = 0; t < 2; ++t) { + uint32_t bits; + topk_unkey(mine[t], bits, expert[t]); + wbits[t] = softmax_topk(static_cast(bits), lane); + } +} + +// ---- TRT-LLM r3's candidate top-10 (q35_moe_front_big._route_warp_fast): each lane keeps the 4 largest of its 16 +// keys (a selection network), the 10 rounds pop from those; when a lane runs dry by the ninth round (it holds 5 or more +// of the top 10) the warp takes the full path. The same 10 keys in the same order as warp_top10_redux. +__device__ __forceinline__ void cas_u32(uint32_t& hi, uint32_t& lo) { + const uint32_t a = max(hi, lo), b = min(hi, lo); + hi = a; + lo = b; +} +// the 4 largest of two descending 4s, descending +__device__ __forceinline__ void top4of8(const uint32_t (&a)[4], const uint32_t (&b)[4], uint32_t (&o)[4]) { + uint32_t t0 = max(a[0], b[3]), t1 = max(a[1], b[2]), t2 = max(a[2], b[1]), t3 = max(a[3], b[0]); + cas_u32(t0, t2); + cas_u32(t1, t3); + cas_u32(t0, t1); + cas_u32(t2, t3); + o[0] = t0; + o[1] = t1; + o[2] = t2; + o[3] = t3; +} +// this lane's 4 largest of 16 keys, descending +__device__ __forceinline__ void top4of16(const uint32_t (&k)[16], uint32_t (&c)[4]) { + uint32_t q[4][4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + uint32_t a0 = k[4 * i], a1 = k[4 * i + 1], b0 = k[4 * i + 2], b1 = k[4 * i + 3]; + cas_u32(a0, a1); + cas_u32(b0, b1); + // merge two descending pairs into the descending 4 + uint32_t c0 = a0, m1 = b0, m2 = a1, c3 = b1; + cas_u32(c0, m1); + cas_u32(m2, c3); + cas_u32(m1, m2); + q[i][0] = c0; + q[i][1] = m1; + q[i][2] = m2; + q[i][3] = c3; + } + uint32_t h0[4], h1[4]; + top4of8(q[0], q[1], h0); + top4of8(q[2], q[3], h1); + top4of8(h0, h1, c); +} +// warp_top10_redux by candidates: false when a lane ran dry (the caller takes the full path) +__device__ __forceinline__ bool warp_top10_fast(const uint32_t (&keys)[16], uint32_t lane, uint32_t& mine) { + uint32_t c[4]; + top4of16(keys, c); + mine = 0; + bool dry = false; +#pragma unroll + for (uint32_t k = 0; k < kTopK; ++k) { + const uint32_t top = warp_max_u32(c[0]); + mine = lane == k ? top : mine; + if (k + 1 < kTopK) { + const bool own = c[0] == top; + c[0] = own ? c[1] : c[0]; + c[1] = own ? c[2] : c[1]; + c[2] = own ? c[3] : c[2]; + c[3] = own ? 0u : c[3]; + } + if (k == kTopK - 2) dry = __any_sync(0xffffffffu, c[0] == 0u); + } + return !dry; +} +__device__ __forceinline__ void row_keys(const uint4& v0, const uint4& v1, uint32_t lane, uint32_t (&keys)[16]) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t w = i < 4 ? (&v0.x)[i] : (&v1.x)[i - 4]; + keys[2 * i] = topk_key(w & 0xFFFFu, lane * 16 + 2 * i); + keys[2 * i + 1] = topk_key(w >> 16, lane * 16 + 2 * i + 1); + } +} +// route_row with the candidate top-10 (bit for bit route_row) +__device__ __forceinline__ void route_row_fast( + const uint4& v0, const uint4& v1, uint32_t lane, uint32_t* s_cand, uint32_t& expert, uint16_t& wbits) { + uint32_t keys[16]; + row_keys(v0, v1, lane, keys); + uint32_t key; + if (!warp_top10_fast(keys, lane, key)) key = warp_top10_redux(keys, lane, s_cand); + uint32_t bits; + topk_unkey(key, bits, expert); + wbits = softmax_topk(static_cast(bits), lane); +} + +// ---- the routing tables +struct Tables { + int32_t* counter; // [1] the launch's CTA arrivals + int32_t* total_padded; // [1] live CTAs * tile + int32_t* expanded_to_permuted; // [M * 10], -1 for a slot outside the window + int32_t* permuted_to_token; // [max CTAs * tile + 1], -1 in a tile's padding + int32_t* tokens_per_expert; // [local_count] + int32_t* cta_expert; // [max CTAs] the CTA's local expert + int32_t* cta_row_limit; // [max CTAs] + int32_t* live_ctas; // [1] + int32_t tile; // tile_n + int32_t local_offset; // the rank's first expert + int32_t local_count; // and its number of experts +}; + +constexpr uint16_t kRemote = 0xFFFFu; + +// The slot's local expert, or kRemote outside the window. +__device__ __forceinline__ uint16_t local_slot(int32_t expert, const Tables& t) { + const int32_t l = expert - t.local_offset; + return (l >= 0 && l < t.local_count) ? static_cast(l) : kRemote; +} + +// The tables of n expanded slots (token i / 10) whose local experts are in s_local (visible to the CTA), by the whole +// CTA of kThreads: per-expert counts, the tile-padded offsets (local experts in order), the CTA -> local expert and +// CTA -> row-limit maps, the live CTA count and padded size, and the expanded <-> permuted maps. kTokenOrder (at most +// 32 tokens): an expert's rows in token order, as the single-block routing writes them; otherwise in the order of +// shared-memory atomics, as the cluster routing does. Scratch: s_count / s_next [512], s_warp [8]. +template +__device__ __forceinline__ void write_tables( + const Tables& t, uint32_t n, const uint16_t* s_local, int32_t* s_count, int32_t* s_next, int32_t* s_warp) { + constexpr uint32_t kWarps = kThreads / 32; + static_assert(2 * kThreads == kExperts, "the scan takes two experts a thread"); + const uint32_t tid = threadIdx.x; + const int32_t tile = t.tile; + for (uint32_t e = tid; e < kExperts; e += kThreads) s_count[e] = 0; + __syncthreads(); + for (uint32_t i = tid; i < n; i += kThreads) { + const uint32_t l = s_local[i]; + if (l == kRemote) continue; + if constexpr (kTokenOrder) { + atomicOr(reinterpret_cast(s_count) + l, 1u << (i / kTopK)); + } else { + atomicAdd(s_count + l, 1); + } + } + __syncthreads(); + // exclusive scan of the experts' CTA counts, two experts a thread + const uint32_t e0 = 2 * tid, lane = tid % 32, warp = tid / 32; + const int32_t c0 = kTokenOrder ? __popc(s_count[e0]) : s_count[e0]; + const int32_t c1 = kTokenOrder ? __popc(s_count[e0 + 1]) : s_count[e0 + 1]; + const int32_t n0 = (c0 + tile - 1) / tile, n1 = (c1 + tile - 1) / tile; + int32_t incl = n0 + n1; +#pragma unroll + for (uint32_t d = 1; d < 32; d <<= 1) { + const int32_t up = __shfl_up_sync(0xffffffffu, incl, d); + incl += lane >= d ? up : 0; + } + if (lane == 31) s_warp[warp] = incl; + __syncthreads(); + int32_t base = 0; +#pragma unroll + for (uint32_t w = 0; w < kWarps; ++w) base += w < warp ? s_warp[w] : 0; + const int32_t off0 = base + incl - n0 - n1, off1 = off0 + n0; + const int32_t offs[2] = {off0, off1}, counts[2] = {c0, c1}, ctas[2] = {n0, n1}; +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int32_t e = static_cast(e0) + j, off = offs[j], count = counts[j]; + if (e < t.local_count) t.tokens_per_expert[e] = count; + s_next[e] = off * tile; + for (int32_t c = 0; c < ctas[j]; ++c) { + t.cta_expert[off + c] = e; + const int32_t limit = min((off + c + 1) * tile, off * tile + count); + t.cta_row_limit[off + c] = limit; + for (int32_t r = limit; r < (off + c + 1) * tile; ++r) t.permuted_to_token[r] = -1; + } + } + if (tid == kThreads - 1) { + t.live_ctas[0] = off1 + n1; + t.total_padded[0] = (off1 + n1) * tile; + } + __syncthreads(); + for (uint32_t i = tid; i < n; i += kThreads) { + const uint32_t l = s_local[i]; + if (l == kRemote) { + t.expanded_to_permuted[i] = -1; + continue; + } + const uint32_t token = i / kTopK; + int32_t pos; + if constexpr (kTokenOrder) { + pos = s_next[l] + __popc(static_cast(s_count[l]) & ((1u << token) - 1u)); + } else { + pos = atomicAdd(s_next + l, 1); + } + t.expanded_to_permuted[i] = pos; + t.permuted_to_token[pos] = static_cast(token); + } +} + +} // namespace q35_route_tables diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_topk.cu b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_topk.cu index 12f1f9f6..72b840f3 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_topk.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/src/route_topk.cu @@ -10,6 +10,10 @@ // other 7 warps quantize the row (router_front.cu's quantize_block). Launched as a programmatic dependent; it waits // before reading anything and releases its dependents at once (they read its outputs after their own wait). // +// With routing tables the last CTA to finish its token also writes the trtllm-gen MoE body's routing +// tables for tile_n and the rank's window of local experts (what its routingIndicesCluster kernel writes, which the body +// then skips; route_tables.cuh), rows of one expert in the order of shared-memory atomics, as there. +// // Built with NVRTC. #include @@ -21,13 +25,14 @@ typedef unsigned char uint8_t; typedef unsigned short uint16_t; typedef int int32_t; typedef unsigned int uint32_t; -#ifndef INFINITY -#define INFINITY __int_as_float(0x7f800000) -#endif + +#include "route_tables.cuh" namespace q35_route_topk { -constexpr uint32_t kHidden = 4096, kExperts = 512, kTopK = 10, kThreads = 256, kSfVec = 16; +using namespace q35_route_tables; + +constexpr uint32_t kHidden = 4096, kSfVec = 16; constexpr uint32_t kBlocksPerRow = kHidden / kSfVec; // 256 struct Params { @@ -38,8 +43,21 @@ struct Params { uint16_t* topk_weights; // [M, 10] bf16 uint8_t* x_fp4; // [M, 2048] uint8_t* x_sf; // [M, 256] + // routing tables (null counter: none); the counter is 0 between launches, the last CTA resets it + Tables tables; + int32_t num_tokens; // M }; +// The routing tables from every token's top-10 (in global memory, visible to this CTA), by the whole CTA. +constexpr uint32_t kMaxRows = 896; // router.py TOPK_MAX_ROWS +__device__ __forceinline__ void write_tables(const Params& p) { + __shared__ uint16_t s_local[kMaxRows * kTopK]; + __shared__ int32_t s_count[kExperts], s_next[kExperts], s_warp[kThreads / 32]; + const uint32_t n = static_cast(p.num_tokens) * kTopK; + for (uint32_t i = threadIdx.x; i < n; i += kThreads) s_local[i] = local_slot(__ldcg(p.topk_ids + i), p.tables); + q35_route_tables::write_tables(p.tables, n, s_local, s_count, s_next, s_warp); +} + // ---- the NVFP4 quantizer (router_front.cu) __device__ __forceinline__ float bf16_lo(uint32_t v) { return __uint_as_float(v << 16); } __device__ __forceinline__ float bf16_hi(uint32_t v) { return __uint_as_float(v & 0xFFFF0000u); } @@ -107,136 +125,22 @@ __device__ __forceinline__ void quantize_block(const uint4& a, const uint4& b, f packed.y = cvt_e2m1x8(hi); } -// ---- the routing (mkm_moe_m16.cu): key = bf16 order above 65535 - expert (higher score first, smaller expert on ties) -__device__ __forceinline__ uint32_t topk_key(uint32_t bits, uint32_t expert) { - const uint32_t twiddled = bits ^ ((bits & 0x8000u) ? 0xFFFFu : 0x8000u); - return (twiddled << 16) | (65535u - expert); -} -__device__ __forceinline__ void topk_unkey(uint32_t key, uint32_t& bits, uint32_t& expert) { - expert = 65535u - (key & 0xFFFFu); - const uint32_t twiddled = key >> 16; - bits = twiddled ^ ((twiddled & 0x8000u) ? 0x8000u : 0xFFFFu); -} -__device__ __forceinline__ uint32_t max16(const uint32_t (&k)[16]) { - const uint32_t a = max(max(max(k[0], k[1]), max(k[2], k[3])), max(max(k[4], k[5]), max(k[6], k[7]))); - const uint32_t b = max(max(max(k[8], k[9]), max(k[10], k[11])), max(max(k[12], k[13]), max(k[14], k[15]))); - return max(a, b); -} -__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) { - uint32_t r; - asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); - return r; -} -__device__ __forceinline__ uint32_t warp_sum_u32(uint32_t v) { - uint32_t r; - asm volatile("redux.sync.add.u32 %0, %1, 0xffffffff;" : "=r"(r) : "r"(v)); - return r; -} -__device__ __forceinline__ uint32_t warp_topk_of_lanes(uint32_t v, uint32_t lane) { - uint32_t mine = 0; -#pragma unroll - for (uint32_t k = 0; k < kTopK; ++k) { - const uint32_t top = warp_max_u32(v); - v = v == top ? 0u : v; - mine = lane == k ? top : mine; - } - return mine; -} -// Top-10 keys of 512 (16 per lane), lane k holding the k-th largest (keys unique and nonzero). -__device__ __forceinline__ uint32_t warp_top10_redux(uint32_t (&keys)[16], uint32_t lane, uint32_t* s_cand) { - uint32_t v = max16(keys); - uint32_t bound = 0; -#pragma unroll - for (uint32_t k = 0; k < kTopK; ++k) { - bound = warp_max_u32(v); - v = v == bound ? 0u : v; - } - uint32_t count = 0; -#pragma unroll - for (uint32_t i = 0; i < 16; ++i) count += keys[i] >= bound; - const uint32_t total = warp_sum_u32(count); - if (total <= 32) { - uint32_t offset = count; -#pragma unroll - for (uint32_t d = 1; d < 32; d <<= 1) { - const uint32_t up = __shfl_up_sync(0xffffffffu, offset, d); - offset += lane >= d ? up : 0u; - } - offset -= count; -#pragma unroll - for (uint32_t i = 0; i < 16; ++i) { - if (keys[i] >= bound) s_cand[offset++] = keys[i]; - } - __syncwarp(); - const uint32_t cand = lane < total ? s_cand[lane] : 0u; - __syncwarp(); - return warp_topk_of_lanes(cand, lane); - } - uint32_t mine = 0; -#pragma unroll 1 - for (uint32_t k = 0; k < kTopK; ++k) { - const uint32_t top = warp_max_u32(max16(keys)); -#pragma unroll - for (uint32_t i = 0; i < 16; ++i) keys[i] = keys[i] == top ? 0u : keys[i]; - mine = lane == k ? top : mine; - } - return mine; -} -// FlashInfer calcSoftmax (bf16 scores) as built with -use_fast_math. -__device__ __forceinline__ float fast_expf(float x) { - float y; - asm("{\n .reg .f32 t;\n mul.ftz.f32 t, %1, 0f3FB8AA3B;\n ex2.approx.ftz.f32 %0, t;\n}" : "=f"(y) : "f"(x)); - return y; -} -__device__ __forceinline__ float fast_div(float a, float b) { - float y; - asm("div.full.ftz.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); - return y; -} -__device__ __forceinline__ uint16_t softmax_topk(uint16_t score_bits, uint32_t lane) { - namespace cg = cooperative_groups; - const auto warp = cg::tiled_partition<32>(cg::this_thread_block()); - const float si = __uint_as_float((uint32_t)score_bits << 16); - float max_score = -INFINITY; - if (lane < kTopK) max_score = si >= max_score ? si : max_score; - max_score = cg::reduce(warp, max_score, cg::greater()); - float sum_score = 0.f, new_score = 0.f; - if (lane < kTopK) { - new_score = fast_expf(__fsub_rn(si, max_score)); - sum_score = __fadd_rn(sum_score, new_score); - } - sum_score = cg::reduce(warp, sum_score, cg::plus()); - uint16_t out = score_bits; - if (lane < kTopK) { - const __nv_bfloat16 b = __float2bfloat16_rn(fast_div(new_score, sum_score)); - out = *reinterpret_cast(&b); - } - return out; -} - __device__ __forceinline__ void route_topk_quant(const Params& p) { __shared__ __align__(16) uint32_t s_cand[32]; + __shared__ int32_t s_last; const uint32_t t = blockIdx.x, warp = threadIdx.x / 32, lane = threadIdx.x % 32; asm volatile("griddepcontrol.wait;" ::: "memory"); asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); if (warp == 0) { const uint4* src = reinterpret_cast(p.logits + static_cast(t) * kExperts + lane * 16); - const uint4 v0 = src[0], v1 = src[1]; - uint32_t keys[16]; -#pragma unroll - for (int i = 0; i < 8; ++i) { - const uint32_t w = i < 4 ? (&v0.x)[i] : (&v1.x)[i - 4]; - keys[2 * i] = topk_key(w & 0xFFFFu, lane * 16 + 2 * i); - keys[2 * i + 1] = topk_key(w >> 16, lane * 16 + 2 * i + 1); - } - const uint32_t key = warp_top10_redux(keys, lane, s_cand); - uint32_t bits, expert; - topk_unkey(key, bits, expert); - const uint16_t wbits = softmax_topk(static_cast(bits), lane); + uint32_t expert; + uint16_t wbits; + route_row(src[0], src[1], lane, s_cand, expert, wbits); if (lane < kTopK) { p.topk_ids[t * kTopK + lane] = static_cast(expert); p.topk_weights[t * kTopK + lane] = wbits; } + if (p.tables.counter != nullptr) __threadfence(); } else if (p.x_fp4 != nullptr) { // null: the router front already quantized the rows const float gs = p.global_scale[0]; const __nv_bfloat16* row = p.x + static_cast(t) * kHidden; @@ -249,6 +153,14 @@ __device__ __forceinline__ void route_topk_quant(const Params& p) { p.x_sf[static_cast(t) * kBlocksPerRow + blk] = static_cast(sf); } } + if (p.tables.counter == nullptr) return; + __syncthreads(); + if (threadIdx.x == 0) s_last = atomicAdd(p.tables.counter, 1) == p.num_tokens - 1; + __syncthreads(); + if (!s_last) return; + __threadfence(); + write_tables(p); + if (threadIdx.x == 0) *p.tables.counter = 0; } } // namespace q35_route_topk diff --git a/test/kernels/test_qwen3_5_moe.py b/test/kernels/test_qwen3_5_moe.py index fb9a817d..46fef75e 100644 --- a/test/kernels/test_qwen3_5_moe.py +++ b/test/kernels/test_qwen3_5_moe.py @@ -17,7 +17,14 @@ - the router front with the shared gate/up stage (slot route_shared, numeric): its logits and NVFP4 input bit for bit the plain launch's, and its gate/up rows against fp64 on SGLang's static FP8 quantization of the rows, at most 1.05 x the mean relative - L2 and 2 x the max abs error of FlashInfer's FP8 BMM, at 1-8 rows; graph replay.""" + L2 and 2 x the max abs error of FlashInfer's FP8 BMM, at 1-8 rows and the TP8 / TP16 + widths (256 / 128 rows); graph replay; +- the TEP8 routed path: the routing tables route_topk and route-big write for a rank's + window of local experts, and the trtllm-gen MoE body on them against FlashInfer's op + as SGLang calls it (on + route_topk's materialized routing, from the logits up to 16 rows) at TP8 and at TEP16 + (MoE TP2 x EP8: 64 local experts, intermediate 512 a rank, ranks 0 and 5): the weights, + the slots routed outside the rank (-1) and every local slot's FC2 row bit for bit.""" import pytest import torch @@ -43,15 +50,15 @@ def _gpu_ready() -> bool: pytestmark = pytest.mark.skipif(not _gpu_ready(), reason="needs an SM100-family GPU") -def _experts(g, rows, cols, std): - """NVFP4 experts [E, rows, cols]: gaussian weights with row and block spread, +def _experts(g, rows, cols, std, n=E): + """NVFP4 experts [n, rows, cols]: gaussian weights with row and block spread, quantized by FlashInfer with a global scale per expert.""" from flashinfer import fp4_quantize - packed = torch.empty(E, rows, cols // 2, dtype=torch.uint8, device="cuda") - sf = torch.empty(E, rows, cols // 16, dtype=torch.uint8, device="cuda") - scale = torch.empty(E, dtype=torch.float32, device="cuda") - for e in range(E): + packed = torch.empty(n, rows, cols // 2, dtype=torch.uint8, device="cuda") + sf = torch.empty(n, rows, cols // 16, dtype=torch.uint8, device="cuda") + scale = torch.empty(n, dtype=torch.float32, device="cuda") + for e in range(n): rs = torch.exp(torch.randn(rows, 1, device="cuda", generator=g) * 0.5) cs = torch.exp(torch.randn(1, cols // 16, device="cuda", generator=g) * 0.5) w = torch.randn(rows, cols, device="cuda", generator=g) * std @@ -425,9 +432,11 @@ def test_router_gemm_quant_shared_replays_from_a_cuda_graph(): # ---------------------------------------------------------------- slot route_topk -def _routed_stock(layer, routing, x_fp4, x_sf): +def _routed_stock(layer, routing, x_fp4, x_sf, tune=None, offset=0, raw=False): """``_stock`` on a materialized routing (ids int32, weights bf16 [M, 10]): the - runner's routed path (FlashInfer's routed MoE, as SGLang calls it).""" + runner's routed path (FlashInfer's routed MoE, as SGLang calls it; ``tune`` its + tune_max_num_tokens, SGLang's next power of 2); for the local experts from ``offset`` + of a layer with ``count`` of them, with ``raw`` FlashInfer's three outputs.""" from flashinfer.fused_moe import trtllm_fp4_block_scale_routed_moe from flashinfer.fused_moe.core import ActivationType @@ -455,15 +464,17 @@ def _routed_stock(layer, routing, x_fp4, x_sf): n_group=0, topk_group=0, intermediate_size=layer["inter"], - local_expert_offset=0, - local_num_experts=E, + local_expert_offset=offset, + local_num_experts=layer.get("count", E), routed_scaling_factor=None, routing_method_type=1, do_finalize=False, activation_type=ActivationType.Swiglu, - tune_max_num_tokens=max(8, m), + tune_max_num_tokens=tune or max(8, m), enable_pdl=True, )[:3] + if raw: + return rows, weights, permuted return rows[permuted.reshape(-1).long()[: m * TOP_K]], weights @@ -511,7 +522,367 @@ def test_route_topk_replays_from_a_cuda_graph(): assert all(torch.equal(a, b) for a, b in zip(eager, captured)) -@pytest.mark.parametrize("m", [17, 21, 28, 35, 49, 56]) +# -------------------------------------------------------- TEP8 routed-body primitives + + +def _tables_ref(ids, tile, window=(0, E)): + """routingIndicesCluster's tables for top-10 ids [M, 10] and tile_n over the local + experts ``window`` (first, count): per-local-expert counts, each one's first padded row + (its CTA offset * tile), the CTA -> local expert and CTA -> row-limit maps of the live + CTAs, and their count.""" + lo, count = window + local = ids.flatten().long() - lo + local = local[(local >= 0) & (local < count)] + counts = torch.bincount(local, minlength=count) + ctas = (counts + tile - 1) // tile + offset = torch.cumsum(ctas, 0) - ctas + live = int(ctas.sum()) + expert = torch.repeat_interleave(torch.arange(count, device=ids.device), ctas) + first = torch.repeat_interleave(offset, ctas) + cta = torch.arange(live, device=ids.device) + limit = torch.minimum( + (cta + 1) * tile, first * tile + torch.repeat_interleave(counts, ctas) + ) + return counts, offset * tile, expert, limit, live + + +def _check_tables(tables, window, token_order=False): + """``tables`` (written for its ids) against ``_tables_ref``: counts, CTA maps, live + count, each local slot in its expert's rows once (in token order with + ``token_order``), mapping back to its token, -1 for every other slot and padding row.""" + m, tile = tables.m, tables.tile + lo, count = window + counts, first, expert, limit, live = _tables_ref(tables.ids, tile, window) + assert tables.live_ctas.item() == live + assert tables.total_padded.item() == live * tile + assert torch.equal(tables.tokens_per_expert[:count], counts.int()) + assert torch.equal(tables.cta_expert[:live], expert.int()) + assert torch.equal(tables.cta_row_limit[:live], limit.int()) + flat = tables.ids.flatten().long() - lo + on = (flat >= 0) & (flat < count) + pos = tables.expanded_to_permuted.long() + assert (pos[~on] == -1).all() + flat, pos = flat[on], pos[on] + assert ((pos >= first[flat]) & (pos < first[flat] + counts[flat])).all() + assert torch.unique(pos).numel() == pos.numel() + token = (torch.arange(m * TOP_K, device=pos.device) // TOP_K)[on] + assert torch.equal(tables.permuted_to_token[pos].long(), token) + if token_order: # an expert's rows ascend with the token (the single-block routing) + rank = pos - first[flat] + same = flat[:, None] == flat[None, :] + before = same & (token[None, :] < token[:, None]) + assert torch.equal(rank, before.sum(1)) + rows = torch.zeros(live * tile, dtype=torch.bool, device=pos.device) + rows[pos] = True + assert (tables.permuted_to_token[: live * tile][~rows] == -1).all() + + +@pytest.mark.parametrize("window", [(0, E), (0, 64), (320, 64)], ids=str) +@pytest.mark.parametrize("tile", [8, 16, 32]) +@pytest.mark.parametrize("m", [7, 14, 17, 28, 56, 112, 224]) +def test_route_topk_writes_the_moe_bodys_routing_tables(m, tile, window): + """With tables, route_topk routes as without them and writes + the trtllm-gen tables of its top-10 for tile_n over the local experts ``window`` (all + 512, or TEP16's 64 of ranks 0 and 5), a second launch alike (the counter resets); each + expert's rows are its own, in any order (atomics, as stock).""" + if window[1] < E and m > 28: + pytest.skip("the expert-parallel steps are up to four requests") + logits, x, scale = _topk_inputs(m, seed=700 + m + tile) + router.build_topk(x.device.index) + want = router.route_topk_quant(logits, x, scale) + tables = router.RouteTables(m, tile, x.device, *window) + for _ in range(2): + got = router.route_topk_quant(logits, x, scale, tables=tables) + torch.cuda.synchronize() + for a, b in zip(got, want): + assert torch.equal(a, b) + assert got[0].data_ptr() == tables.ids.data_ptr() + assert tables.counter.item() == 0 + _check_tables(tables, window) + + +def _experts_of(layer, rank=0): + """The FusedMoE attributes moe/body.py reads, from the layer fixture (EP rank + ``rank`` of a layer with ``count`` local experts).""" + from types import SimpleNamespace + + return SimpleNamespace( + num_experts=E, + num_local_experts=layer.get("count", E), + moe_ep_rank=rank, + top_k=TOP_K, + intermediate_size_per_partition=layer["inter"], + w13_weight=layer["w13"], + w13_weight_scale=layer["w13_sf"], + w2_weight=layer["w2"], + w2_weight_scale=layer["w2_sf"], + g1_scale_c=layer["g1c"], + g1_alphas=layer["g1"], + g2_alphas=layer["g2"], + moe_runner_config=SimpleNamespace(activation="silu", is_gated=True), + ) + + +@pytest.mark.parametrize("tuned", [True, False], ids=["tuned", "default"]) +@pytest.mark.parametrize("m", [17, 28, 56, 112, 224]) +def test_moe_body_on_route_topks_tables_is_the_routed_op(layer, m, tuned, monkeypatch): + """The MoE body (FlashInfer's DA entry, the routed op's tactic: the + autotuner's, or without one FlashInfer's default) on the tables route_topk's launch + writes gives the routed op's FC2 rows and weights, row for row by expanded index, bit + for bit; replayed from a CUDA graph too.""" + from flashinfer.autotuner import AutoTuner, autotune + from sglang.srt.utils import next_power_of_2 + + from sgl_nvidia_mega.qwen3_5.moe import body + + logits, x, scale = _topk_inputs(m, seed=1300 + m) + router.build_topk(x.device.index) + ids, weights, x_fp4, x_sf = router.route_topk_quant(logits, x, scale) + tune = next_power_of_2(m) + if tuned: + with autotune(True): + _routed_stock(layer, (ids, weights), x_fp4, x_sf, tune) + else: # no tactic for the routed op or the body: both run FlashInfer's default + monkeypatch.setattr( + AutoTuner, "choose_one", lambda self, op, runners, *a, **k: (runners[0], -1) + ) + want_rows, want_w = _routed_stock(layer, (ids, weights), x_fp4, x_sf, tune) + experts = _experts_of(layer) + body._WIDTHS.clear() + width = body.width(experts, x) + assert (width.tactic == [-1, -1]) != tuned + + def step(): + got = router.route_topk_quant(logits, x, scale, tables=width.tables) + rows, w, permuted = body.run( + width, experts, got[2], got[3].view(torch.float8_e4m3fn) + )[:3] + return rows[permuted.reshape(-1).long()], w + + got_rows, got_w = step() + torch.cuda.synchronize() + assert torch.equal(got_w.view(torch.int16), want_w.view(torch.int16)) + assert torch.equal(got_rows.view(torch.int16), want_rows.view(torch.int16)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = step() + for _ in range(2): + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured[0].view(torch.int16), want_rows.view(torch.int16)) + body._WIDTHS.clear() + + +# ------------------------------------------------------ body under expert parallelism + +TEP16_EXPERTS, TEP16_INTER = ( + 64, + 512, +) # MoE TP2 x EP8: local experts and intermediate a rank + + +@pytest.fixture(scope="module") +def ep_layer(): + """One TEP16 Qwen3.5 MoE layer's local experts (64, intermediate 512 a rank) on the + trtllm-gen weight layout, as SGLang loads them; any rank's window holds them.""" + return _local_layer(TEP16_EXPERTS, TEP16_INTER, 13) + + +def _local_layer(n, inter, seed): + """A layer's ``n`` local experts (intermediate ``inter`` a rank) on the trtllm-gen + weight layout, as SGLang loads them.""" + from sglang.srt.layers.quantization.utils import ( + prepare_static_weights_for_trtllm_fp4_moe, + ) + + g = torch.Generator(device="cuda").manual_seed(seed) + w13, sf13, s13 = _experts(g, 2 * inter, H, 0.02, n) + w2, sf2, s2 = _experts(g, H, inter, 0.02, n) + w13, sf13, w2, sf2 = prepare_static_weights_for_trtllm_fp4_moe( + w13, w2, sf13, sf2, H, inter, n, is_gated=True + ) + g1 = (1.0 / (X_SCALE * s13)).float() + return { + "w13": w13, + "w13_sf": sf13.view(torch.float8_e4m3fn).reshape(n, 2 * inter, H // 16), + "w2": w2, + "w2_sf": sf2.view(torch.float8_e4m3fn).reshape(n, H, inter // 16), + "g1": g1, + "g1c": (g1 / H_SCALE).float(), + "g2": (H_SCALE / s2).float(), + "inter": inter, + "count": n, + } + + +def _stock_ep(layer, logits, x_fp4, x_sf, offset): + """FlashInfer's op from the logits as SGLang calls it at up to 16 rows (deferred + finalize) for the local experts from ``offset``: (FC2 rows, weights, expanded -> + permuted index, -1 for a slot routed elsewhere).""" + from flashinfer.fused_moe import trtllm_fp4_block_scale_moe + from flashinfer.fused_moe.core import ActivationType + from sglang.srt.layers.moe.utils import RoutingMethodType + from sglang.srt.utils import next_power_of_2 + + return trtllm_fp4_block_scale_moe( + routing_logits=logits, + routing_bias=None, + hidden_states=x_fp4, + hidden_states_scale=x_sf.view(torch.float8_e4m3fn), + gemm1_weights=layer["w13"], + gemm1_weights_scale=layer["w13_sf"], + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=layer["w2"], + gemm2_weights_scale=layer["w2_sf"], + gemm2_bias=None, + output1_scale_scalar=layer["g1c"], + output1_scale_gate_scalar=layer["g1"], + output2_scale_scalar=layer["g2"], + per_token_scale=None, + num_experts=E, + top_k=TOP_K, + n_group=None, + topk_group=None, + intermediate_size=layer["inter"], + local_expert_offset=offset, + local_num_experts=layer.get("count", E), + routed_scaling_factor=None, + routing_method_type=RoutingMethodType.RenormalizeNaive, + do_finalize=False, + activation_type=ActivationType.Swiglu, + tune_max_num_tokens=next_power_of_2(logits.shape[0]), + enable_pdl=True, + )[:3] + + +def _same_moe(got, want, m): + """Two deferred-finalize outputs (FC2 rows, weights, expanded -> permuted) agree: the + weights bit for bit, the same slots routed outside the rank (-1), and every local + slot's FC2 row bit for bit. Returns the two permutations.""" + rows, w, perm = got[:3] + want_rows, want_w, want_perm = want[:3] + perm = perm.reshape(-1)[: m * TOP_K].long() + want_perm = want_perm.reshape(-1)[: m * TOP_K].long() + assert torch.equal(w.view(torch.int16), want_w.view(torch.int16)) + assert torch.equal(perm < 0, want_perm < 0) + on = want_perm >= 0 + assert torch.equal( + rows[perm[on]].view(torch.int16), want_rows[want_perm[on]].view(torch.int16) + ) + return perm, want_perm + + +def _ep_router_inputs(m, seed, kind, window): + """Rows, router weights and the routed input scale; ``kind`` "local" lifts the + window's experts by ~1 sd of a logit (about half the slots local), "same" sends every + token to the window's first 10 experts (several CTAs an expert past 8 rows), "ties" + gives every third expert the same top logit (the lowest ids win, past a warp of + candidates).""" + lo, count = window + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(m, H, device="cuda", generator=g) + x[0, :32] = 0 + w = torch.randn(E, H, device="cuda", generator=g) * 0.02 + mean = x.mean(0) + up = mean / mean.pow(2).sum() # x_t . up ~ 1 for every token t + if kind == "local": + w[lo : lo + count] += 1.5 * up + elif kind == "same": + w[lo : lo + TOP_K] += 12.0 * up + elif kind == "ties": + w[::3] = w[1] + 3.0 * up + return ( + x.to(torch.bfloat16), + w.to(torch.bfloat16), + torch.tensor([X_SCALE], device="cuda"), + ) + + +EP_CASES = [ # (layout, rank): TP8 (all 512 experts local) and TEP16 ranks 0 and 5 + ("tp8", 0), + ("tep16", 0), + ("tep16", 5), +] + + +def _ep_case(request, layout, rank): + layer = request.getfixturevalue("ep_layer" if layout == "tep16" else "tp8_layer") + count = layer.get("count", E) + return layer, (rank * count, count) + + +@pytest.fixture(scope="module") +def tp8_layer(): + """The TP8 layer of the ``layer`` fixture (intermediate 128 a rank, 512 experts).""" + from sglang.srt.layers.quantization.utils import ( + prepare_static_weights_for_trtllm_fp4_moe, + ) + + inter = 128 + g = torch.Generator(device="cuda").manual_seed(11) + w13, sf13, s13 = _experts(g, 2 * inter, H, 0.02) + w2, sf2, s2 = _experts(g, H, inter, 0.02) + w13, sf13, w2, sf2 = prepare_static_weights_for_trtllm_fp4_moe( + w13, w2, sf13, sf2, H, inter, E, is_gated=True + ) + g1 = (1.0 / (X_SCALE * s13)).float() + return { + "w13": w13, + "w13_sf": sf13.view(torch.float8_e4m3fn).reshape(E, 2 * inter, H // 16), + "w2": w2, + "w2_sf": sf2.view(torch.float8_e4m3fn).reshape(E, H, inter // 16), + "g1": g1, + "g1c": (g1 / H_SCALE).float(), + "g2": (H_SCALE / s2).float(), + "inter": inter, + } + + +@pytest.mark.parametrize("kind", ["random", "local", "same"]) +@pytest.mark.parametrize("m", [7, 14, 28]) +@pytest.mark.parametrize("layout,rank", EP_CASES) +def test_moe_body_on_route_topks_ep_tables_is_the_routed_op( + request, layout, rank, m, kind +): + """For a rank's local experts, the MoE body on route_topk's + tables is FlashInfer's routed op on route_topk's routing (the runner's path past 16 + rows) and its op from the logits: weights, the slots routed elsewhere (-1) and every + local slot's FC2 row bit for bit; replayed from a CUDA graph too.""" + from sgl_nvidia_mega.qwen3_5.moe import body + + layer, window = _ep_case(request, layout, rank) + x, w, scale = _ep_router_inputs(m, 3000 + 31 * m + rank, kind, window) + logits = torch.nn.functional.linear(x, w) + router.build_topk(x.device.index) + ids, wts, x_fp4, x_sf = router.route_topk_quant(logits, x, scale) + experts = _experts_of(layer, rank) + body._WIDTHS.clear() + width = body.width(experts, x) + + def step(): + got = router.route_topk_quant(logits, x, scale, tables=width.tables) + return body.run(width, experts, got[2], got[3].view(torch.float8_e4m3fn))[:3] + + got = step() + torch.cuda.synchronize() + _check_tables(width.tables, window) + routed = _routed_stock(layer, (ids, wts), x_fp4, x_sf, None, window[0], raw=True) + _same_moe(got, routed, m) + _same_moe(got, _stock_ep(layer, logits, x_fp4, x_sf, window[0]), m) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = step() + for _ in range(2): + graph.replay() + torch.cuda.synchronize() + _same_moe(captured, routed, m) + body._WIDTHS.clear() + + +@pytest.mark.parametrize("m", [17, 21, 28, 32]) def test_router_front_tiles_are_its_16_row_launches(m): """Slot route_gemm (numeric past 16 rows): the router front over 16-row token tiles in one launch equals the 16-row launch on each tile, bit for bit; its NVFP4 input is the @@ -536,3 +907,233 @@ def test_router_front_tiles_are_its_16_row_launches(m): only = router.route_topk_quant(logits, x, scale, quantize=False) assert only[2] is None and only[3] is None assert torch.equal(routed[0], only[0]) and torch.equal(routed[1], only[1]) + + +# ------------------------------------------- routed body at TEP8's wide request counts + +TEP8_EXPERTS, TEP8_INTER = ( + 128, + 512, +) # MoE TP2 x EP4: local experts and intermediate a rank +TEP8_ROWS = [231, 448, 896] # 33, 64 and 128 requests of 7 verify rows +# route_big's TEP8 tiers L1 / L2: 5, 16, 32, 33, 64 and 128 requests +TEP8_BIG_ROWS = [35, 112, 224, 231, 448, 518, 672, 784, 896] + + +@pytest.fixture(scope="module") +def tep8_layer(): + """One TEP8 Qwen3.5 MoE layer's local experts (128, intermediate 512 a rank).""" + return _local_layer(TEP8_EXPERTS, TEP8_INTER, 17) + + +@pytest.mark.parametrize( + "window", [(0, TEP8_EXPERTS), (384, TEP8_EXPERTS), (0, E)], ids=str +) +@pytest.mark.parametrize("tile", [8, 32, 128]) +@pytest.mark.parametrize("m", TEP8_ROWS) +def test_route_topk_writes_the_tep8_routing_tables(m, tile, window): + """At TEP8's widths, route_topk with tables routes as without + them and writes the tables of the rank's 128 local experts (ranks 0 and 3; all 512) + for tile_n, a second launch alike.""" + logits, x, scale = _topk_inputs(m, seed=4100 + m + tile) + router.build_topk(x.device.index) + want = router.route_topk_quant(logits, x, scale) + tables = router.RouteTables(m, tile, x.device, *window) + for _ in range(2): + got = router.route_topk_quant(logits, x, scale, tables=tables) + torch.cuda.synchronize() + for a, b in zip(got, want): + assert torch.equal(a, b) + assert tables.counter.item() == 0 + _check_tables(tables, window) + + +@pytest.mark.parametrize("tuned", [True, False], ids=["tuned", "default"]) +@pytest.mark.parametrize("kind", ["random", "local", "same"]) +@pytest.mark.parametrize("m", TEP8_ROWS) +@pytest.mark.parametrize("rank", [0, 3]) +def test_moe_body_on_route_topks_tep8_tables_is_the_routed_op( + tep8_layer, rank, m, kind, tuned, monkeypatch +): + """At TEP8's widths, the MoE body on route_topk's tables for a rank's + 128 local experts is FlashInfer's routed op on route_topk's routing (the runner's + path there; the autotuner's tactic or FlashInfer's default for both): weights, the + slots routed elsewhere (-1) and every local slot's FC2 row bit for bit; replayed from + a CUDA graph too.""" + from flashinfer.autotuner import AutoTuner, autotune + from sglang.srt.utils import next_power_of_2 + + from sgl_nvidia_mega.qwen3_5.moe import body + + window = (rank * TEP8_EXPERTS, TEP8_EXPERTS) + x, w, scale = _ep_router_inputs(m, 5000 + 31 * m + rank, kind, window) + logits = torch.nn.functional.linear(x, w) + router.build_topk(x.device.index) + ids, wts, x_fp4, x_sf = router.route_topk_quant(logits, x, scale) + tune = next_power_of_2(m) + if tuned: + with autotune(True): + _routed_stock(tep8_layer, (ids, wts), x_fp4, x_sf, tune, window[0]) + else: # no tactic for the routed op or the body: both run FlashInfer's default + monkeypatch.setattr( + AutoTuner, "choose_one", lambda self, op, runners, *a, **k: (runners[0], -1) + ) + routed = _routed_stock( + tep8_layer, (ids, wts), x_fp4, x_sf, tune, window[0], raw=True + ) + experts = _experts_of(tep8_layer, rank) + body._WIDTHS.clear() + width = body.width(experts, x) + assert (width.tactic == [-1, -1]) != tuned + + def step(): + got = router.route_topk_quant(logits, x, scale, tables=width.tables) + return body.run(width, experts, got[2], got[3].view(torch.float8_e4m3fn))[:3] + + got = step() + torch.cuda.synchronize() + _check_tables(width.tables, window) + _same_moe(got, routed, m) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = step() + for _ in range(2): + graph.replay() + torch.cuda.synchronize() + _same_moe(captured, routed, m) + body._WIDTHS.clear() + + +@pytest.mark.parametrize( + "window", [(0, TEP8_EXPERTS), (384, TEP8_EXPERTS), (0, E)], ids=str +) +@pytest.mark.parametrize("tile", [8, 32, 128]) +@pytest.mark.parametrize("m", [35, 112, 168, 224, 231, 252, 259, 448, 512, 518, 896]) +def test_route_big_is_route_topk_with_its_tables_in_token_order(m, tile, window): + """Route-big's launch (one routing cluster of 8 / 16 CTAs, the rest + quantizing) gives route_topk's top-10 and NVFP4 input bit for bit and the same routing + tables, an expert's rows in token order; a second launch on the same tables alike; and from + a CUDA graph.""" + logits, x, scale = _topk_inputs(m, seed=4300 + m + tile) + router.build_topk(x.device.index) + router.build_big(x.device.index, window[1]) + want = router.route_topk_quant(logits, x, scale) + tables = router.RouteTables(m, tile, x.device, *window) + for _ in range(2): + got = router.route_big_quant(logits, x, scale, tables) + torch.cuda.synchronize() + for a, b in zip(got, want): + assert torch.equal(a, b) + _check_tables(tables, window, token_order=True) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = router.route_big_quant(logits, x, scale, tables) + tables.expanded_to_permuted.fill_(7) + tables.permuted_to_token.fill_(7) + graph.replay() + torch.cuda.synchronize() + for a, b in zip(captured, want): + assert torch.equal(a, b) + _check_tables(tables, window, token_order=True) + + +@pytest.mark.parametrize("window", [(0, TEP8_EXPERTS), (384, TEP8_EXPERTS)], ids=str) +@pytest.mark.parametrize("tile", [32, 64]) +@pytest.mark.parametrize("ctas", [32, 64]) +@pytest.mark.parametrize("m", [672, 679, 728, 784, 840, 896]) +def test_route_big_wide_launches_are_route_topk(m, ctas, tile, window): + """Both validated L2 launch widths reproduce route_topk and its padded tables.""" + logits, x, scale = _topk_inputs(m, seed=7100 + m + ctas + tile) + router.build_topk(x.device.index) + router.build_big(x.device.index, window[1]) + want = router.route_topk_quant(logits, x, scale) + tables = router.RouteTables(m, tile, x.device, *window) + got = router.route_big_quant(logits, x, scale, tables, ctas=ctas) + torch.cuda.synchronize() + for actual, expected in zip(got, want): + assert torch.equal(actual, expected) + _check_tables(tables, window, token_order=True) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = router.route_big_quant(logits, x, scale, tables, ctas=ctas) + tables.expanded_to_permuted.fill_(7) + tables.permuted_to_token.fill_(7) + graph.replay() + torch.cuda.synchronize() + for actual, expected in zip(captured, want): + assert torch.equal(actual, expected) + _check_tables(tables, window, token_order=True) + + +@pytest.mark.parametrize("kind", ["random", "local", "same"]) +@pytest.mark.parametrize("m", TEP8_BIG_ROWS) +@pytest.mark.parametrize("rank", [0, 3]) +def test_moe_body_on_route_bigs_tep8_tables_is_the_routed_op(tep8_layer, rank, m, kind): + """The MoE body on route_big's tables (a rank's 128 local experts) is + FlashInfer's routed op on route_topk's routing: weights, the slots routed elsewhere (-1) and + every local slot's FC2 row bit for bit; replayed from a CUDA graph too.""" + from sgl_nvidia_mega.qwen3_5.moe import body + + window = (rank * TEP8_EXPERTS, TEP8_EXPERTS) + x, w, scale = _ep_router_inputs(m, 5200 + 31 * m + rank, kind, window) + logits = torch.nn.functional.linear(x, w) + router.build_topk(x.device.index) + router.build_big(x.device.index, TEP8_EXPERTS) + ids, wts, x_fp4, x_sf = router.route_topk_quant(logits, x, scale) + routed = _routed_stock( + tep8_layer, (ids, wts), x_fp4, x_sf, None, window[0], raw=True + ) + experts = _experts_of(tep8_layer, rank) + body._WIDTHS.clear() + width = body.width(experts, x) + + def step(): + got = router.route_big_quant(logits, x, scale, width.tables) + return body.run(width, experts, got[2], got[3].view(torch.float8_e4m3fn))[:3] + + got = step() + torch.cuda.synchronize() + _check_tables(width.tables, window, token_order=True) + _same_moe(got, routed, m) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = step() + for _ in range(2): + graph.replay() + torch.cuda.synchronize() + _same_moe(captured, routed, m) + body._WIDTHS.clear() + + +def test_moe_body_widths_share_their_buffers(tep8_layer): + """TEP8's widths run their bodies on views of one set of buffers (the widest asked + for first, as SGLang captures its graphs), each still the routed op's, alternately.""" + from sgl_nvidia_mega.qwen3_5.moe import body + + window = (TEP8_EXPERTS, TEP8_EXPERTS) + router.build_topk(torch.cuda.current_device()) + experts = _experts_of(tep8_layer, 1) + body._WIDTHS.clear() + body._SHARED.clear() + cases = [] + for m in (896, 231): + x, w, scale = _ep_router_inputs(m, 6000 + m, "local", window) + logits = torch.nn.functional.linear(x, w) + ids, wts, x_fp4, x_sf = router.route_topk_quant(logits, x, scale) + want = _routed_stock( + tep8_layer, (ids, wts), x_fp4, x_sf, None, window[0], raw=True + ) + cases.append((m, x, logits, scale, body.width(experts, x), want)) + for _ in range(2): + for m, x, logits, scale, width, want in cases: + got = router.route_topk_quant(logits, x, scale, tables=width.tables) + got = body.run(width, experts, got[2], got[3].view(torch.float8_e4m3fn))[:3] + torch.cuda.synchronize() + _same_moe(got, want, m) + wide, narrow = cases[0][4].workspace, cases[1][4].workspace + for a, b in zip( + wide, narrow + ): # the narrow width's buffers that fit are the wide one's + if b.numel() and b.shape[0] <= a.shape[0]: + assert a.data_ptr() == b.data_ptr() + body._WIDTHS.clear() diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 8717c03f..e0e04aba 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -186,6 +186,8 @@ def test_an_sglang_without_the_eagle_forward_view_is_refused(monkeypatch): "slot 'mkm_wide' serves the steps of more than one request ('multi_request')", ), ({"mkm_m32"}, "slot 'mkm_m32' takes slot 'route_topk''s rows"), + ({"route_pre"}, "unknown slot 'route_pre'"), + ({"route_pre_front"}, "unknown slot 'route_pre_front'"), ( {"route_topk", "mkm_m32"}, "slot 'mkm_m32' serves the steps of more than two requests ('multi_request')", @@ -283,9 +285,13 @@ def test_the_trt_gdn_front_widens_slot_gdn_front_to_four_requests(): def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): - # every slot built except mkm's 16-token engine (mkm_wide, below) + # Every public slot is built except mkm's 16-token engine (mkm_wide, below). every = plan.Facts( - **{**dict.fromkeys(plan.Facts.__struct_fields__, True), "mkm_wide": False} + **{ + **dict.fromkeys(plan.Facts.__struct_fields__, True), + "mkm_wide": False, + "tep8_moe": False, + } ) assert _on(plan.slots(7, every)) == sorted( set(plan.Slots.__struct_fields__) @@ -296,6 +302,8 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): "mkm_m256", "mkm_m512", "route_gemm", + "route_pre", + "route_pre_topk", "route_shared", "route_topk", } @@ -361,6 +369,63 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): assert _on(plan.slots(28, plan.Facts(route_topk=True))) == ["route_topk"] assert _on(plan.slots(7, plan.Facts(route_topk=True))) == ["route_topk"] assert not plan.slots(28, plan.Facts(route_gemm=True)).route_gemm + # TEP8's routed-body plan is internal: S1/S2 use route plus a routing-only + # top-k launch, M uses the low-M shared expert, and L uses route-big. It + # supersedes the all-expert MKM engines and the numeric tiled router front. + tep8 = dict(route=True, route_topk=True, tep8_moe=True) + for rows in (7, 14): + assert _on(plan.slots(rows, plan.Facts(**tep8))) == [ + "route", + "route_pre_topk", + ] + for rows in (21, 28): + assert _on(plan.slots(rows, plan.Facts(**tep8))) == ["route_topk"] + tier_m = plan.Facts( + **tep8, + route_gemm=True, + lowm_se=True, + lowm_se_pre=True, + ) + for rows in (21, 28): + assert _on(plan.slots(rows, tier_m)) == [ + "lowm_se", + "route_pre", + "route_topk", + ] + tier_l = plan.Facts( + **tep8, + route_gemm=True, + mkm=True, + mkm_wide=True, + mkm_m32=True, + mkm_m64=True, + mkm_m128=True, + mkm_m256=True, + mkm_m512=True, + ) + for rows in (35, 112, 224, 448, 672, 896): + assert _on(plan.slots(rows, tier_l)) == ["route_pre", "route_topk"] + assert plan.lowm_se_pre_rows(8, tep8_m=True) == 28 + assert plan.lowm_se_pre_rows(4, tep8_m=True) == 0 + assert [plan.moe_front_tier(r) for r in (7, 14, 28, 35, 224, 231, 896)] == [ + "S1", + "S2", + "M", + "L1", + "L1", + "L2", + "L2", + ] + # The S tiers run beside the low-M shared expert and replace route_shared. + s_tier = plan.Facts( + **tep8, + route_shared=True, + lowm_se=True, + mkm=True, + ) + for rows in (7, 14): + assert _on(plan.slots(rows, s_tier)) == ["lowm_se", "route", "route_pre_topk"] + assert _on(plan.slots(21, s_tier)) == ["route_topk"] # SGLang offers every multiple of 7 up to MAX_ROWS; every request count is a candidate offered = tuple(range(7, plan.MAX_ROWS + 1, 7)) assert plan.candidates(offered, 7) == offered From 9559e9949c79b94d46f166ec01621abbb81fe861 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 16:26:54 -0700 Subject: [PATCH 13/62] feat: fuse TEP8 draft boundaries --- sgl_nvidia_mega/qwen3_5/draft/layer.py | 115 ++++++++++++++++++-- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 4 +- test/kernels/test_qwen3_5_draft_front.py | 30 +++++ test/unit/test_qwen3_5_runtime.py | 91 +++++++++++++++- 4 files changed, 230 insertions(+), 10 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/draft/layer.py b/sgl_nvidia_mega/qwen3_5/draft/layer.py index ff305147..a37f08a3 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/layer.py +++ b/sgl_nvidia_mega/qwen3_5/draft/layer.py @@ -14,9 +14,12 @@ - ``top1``: the LM head on this rank's vocabulary shard, the all-gather of the logits and their arg-max (``torch.argmax``'s choice, as SGLang's top-1). -Stock: the ops SGLang's draft runs for these rows, in its order. Slot ``draft_moe`` -replaces the MoE, slot ``draft_hop`` the two all-reduce points and ``o_proj``; slot -``draft_fin`` (numeric) hands a deferred MoE finalize to the hop's final norm launch. +Stock: the ops SGLang's draft runs for these rows, in its order. The pre-fc norms write +directly into the two halves of ``fc``'s input, bit for bit the concatenated result. +Slot ``draft_moe`` replaces the MoE, slot ``draft_hop`` the two all-reduce points and +``o_proj``; slot ``draft_fin`` (numeric) hands a deferred MoE finalize to a fused final +norm launch. At TEP8 it also uses SGLang's MNNVL CuTe DSL workspace for steps wider +than the draft hop. """ from __future__ import annotations @@ -57,7 +60,18 @@ "m", "complete", ), - "srt/layers/layernorm.py": ('"gemma_weight"',), + "srt/layers/layernorm.py": ( + '"gemma_weight"', + "gemma_rmsnorm", + "variance_epsilon", + ), + "srt/layers/flashinfer_mnnvl_cutedsl.py": ( + "_WORKSPACE", + "supports", + "all_reduce_residual_rms_norm", + "moe_finalize_all_reduce_rms_norm", + ), + "srt/runtime_context.py": ("get_parallel", '"tp_size"', '"moe_ep_size"'), } # The SGLang functions whose body this module mirrors, by the digest of their code @@ -72,9 +86,51 @@ def fc_input(parts, token_ids, hidden): """The fc output [rows, hidden] of the tokens ``token_ids`` [rows] and the hidden rows ``hidden`` [rows, hidden] they follow.""" lm = parts.lm - embeds = lm.pre_fc_norm_embedding(parts.embed(token_ids)) - hidden = lm.pre_fc_norm_hidden(hidden) - return lm.fc(torch.cat([embeds, hidden], dim=-1)) + _note_fc_cat() + return lm.fc( + norms_into_cat( + parts.embed(token_ids), + hidden, + lm.pre_fc_norm_embedding, + lm.pre_fc_norm_hidden, + ) + ) + + +def norms_into_cat(embeds, hidden, norm_embeds, norm_hidden): + """The concatenated pre-fc Gemma RMSNorm outputs, written in place.""" + from sglang.srt.layers.layernorm import gemma_rmsnorm + + rows, hidden_size = hidden.shape + out = torch.empty((rows, 2 * hidden_size), dtype=hidden.dtype, device=hidden.device) + gemma_rmsnorm( + embeds, + norm_embeds.weight.data, + norm_embeds.variance_epsilon, + out=out[:, :hidden_size], + ) + gemma_rmsnorm( + hidden, + norm_hidden.weight.data, + norm_hidden.variance_epsilon, + out=out[:, hidden_size:], + ) + return out + + +def _mnnvl_workspace(parts, rows: int): + """The target's TEP8 MNNVL workspace for ``draft_fin`` at this width.""" + if "draft_fin" not in parts.slots: + return None + from sglang.srt.runtime_context import get_parallel + + parallel = get_parallel() + if parallel.tp_size != 8 or parallel.moe_ep_size <= 1: + return None + from sglang.srt.layers import flashinfer_mnnvl_cutedsl + + workspace = flashinfer_mnnvl_cutedsl._WORKSPACE + return workspace if workspace is not None and workspace.supports(rows) else None # Slot draft_hop on the draft extend's rows too: their FFN input all-reduce + residual + @@ -110,6 +166,9 @@ def layer(parts, x, positions, cache_loc, meta, kv, *, rows_per_request: int): and parts.hop.hop.supports(rows) else None ) + cfuse = ( + _mnnvl_workspace(parts, rows) if hop is None and extend_hop is None else None + ) attn_out = parts.attention( lay.input_layernorm(x), positions, @@ -150,6 +209,12 @@ def layer(parts, x, positions, cache_loc, meta, kv, *, rows_per_request: int): moe_in, residual = extend_hop.ffn_input( attn_out, residual, lay.post_attention_layernorm, moe_single_launch=False ) + elif cfuse is not None: + moe_in, residual = cfuse.all_reduce_residual_rms_norm( + local_contribution=attn_out, + residual=residual, + gamma=lay.post_attention_layernorm.gemma_weight, + ) else: moe_in, residual = lay.post_attention_layernorm( tensor_model_parallel_all_reduce(attn_out), residual @@ -186,6 +251,23 @@ def layer(parts, x, positions, cache_loc, meta, kv, *, rows_per_request: int): # SGLang's deferred finish already applies its configured reduction. out, _ = parts.model.norm(moe_out, residual) return out + elif ( + cfuse is not None + and mkm is None + and moe is None + and lay.mlp.supports_deferred_finalize + ): + owed = lay.mlp(moe_in, defer_finalize=True) + _note_fin(rows, "in SGLang's MNNVL CuTe DSL final norm", True) + out, _ = cfuse.moe_finalize_all_reduce_rms_norm( + routed_output=owed.routed_output, + expert_weights=owed.expert_weights, + permuted_indices=owed.permuted_indices, + gated_shared_output=owed.gated_shared_output, + residual=residual, + gamma=parts.model.norm.gemma_weight, + ) + return out elif mkm is not None: moe_out = mkm(lay, moe_in).complete() else: @@ -200,11 +282,30 @@ def layer(parts, x, positions, cache_loc, meta, kv, *, rows_per_request: int): and parts.mkd.head_follows(), ) return out + if cfuse is not None: + _note_fin(rows, "after SGLang's MoE in the MNNVL CuTe DSL final norm", True) + out, _ = cfuse.all_reduce_residual_rms_norm( + local_contribution=moe_out, + residual=residual, + gamma=parts.model.norm.gemma_weight, + ) + return out out, _ = parts.model.norm(tensor_model_parallel_all_reduce(moe_out), residual) return out _FIN_NOTED: set[tuple[bool, int]] = set() +_FC_CAT_NOTED = False + + +def _note_fc_cat() -> None: + global _FC_CAT_NOTED + if not _FC_CAT_NOTED: + _FC_CAT_NOTED = True + logger.info( + "NVIDIA megakernel: Qwen3.5 MTP draft pre-fc norms write " + "directly into fc's input" + ) def _note_fin(rows: int, where: str, extend: bool = False) -> None: diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 2d2aab60..dcf5f42f 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -65,7 +65,7 @@ "draft_q1", "draft_xchg", "draft_hop", - # numeric: hand deferred MoE rows to the draft hop's fused final norm + # numeric: hand deferred MoE rows to a fused final norm "draft_fin", "draft_front", "draft_skv2", @@ -234,7 +234,7 @@ def slot_reasons() -> list[str]: reasons.append(f"{SLOTS_ENV}: slot 'draft_hop' runs on slot 'hop''s mailbox") if "draft_fin" in SLOTS and "draft_hop" not in SLOTS: reasons.append( - f"{SLOTS_ENV}: slot 'draft_fin' runs in slot 'draft_hop''s final norm" + f"{SLOTS_ENV}: slot 'draft_fin' extends slot 'draft_hop''s fused final norm" ) if "draft_multi" in SLOTS and "multi_request" not in SLOTS: reasons.append( diff --git a/test/kernels/test_qwen3_5_draft_front.py b/test/kernels/test_qwen3_5_draft_front.py index f5268172..05354a98 100644 --- a/test/kernels/test_qwen3_5_draft_front.py +++ b/test/kernels/test_qwen3_5_draft_front.py @@ -142,3 +142,33 @@ def test_the_front_replays_in_a_graph_with_new_positions_and_slots(heads): q_s, gate_s = _stock(fresh, kv_s) assert _same(q_s, q_p) and _same(gate_s, gate_p) assert _same(kv_s[0], kv_p[0]) and _same(kv_s[1], kv_p[1]) + + +@pytest.mark.parametrize("rows", [1, 7, 16, 32, 96, 128, 672]) +def test_the_pre_fc_norms_write_fcs_input_as_the_concatenation(rows): + """The in-place pre-fc norms equal the concatenated GemmaRMSNorm outputs.""" + from sglang.srt.layers.layernorm import GemmaRMSNorm + + from sgl_nvidia_mega.qwen3_5.draft.layer import norms_into_cat + + torch.manual_seed(9100 + rows) + hidden_size = 4096 + # BF16 weights, as the server loads them (the kernel takes the input's dtype). + norm_embeds, norm_hidden = ( + GemmaRMSNorm(hidden_size, 1e-6).cuda().to(torch.bfloat16) for _ in range(2) + ) + with torch.no_grad(): + norm_embeds.weight.normal_(0.0, 0.2) + norm_hidden.weight.normal_(0.0, 0.2) + embeds = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16) + hidden = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16) * 3 + expected = torch.cat([norm_embeds(embeds), norm_hidden(hidden)], dim=-1) + actual = norms_into_cat(embeds, hidden, norm_embeds, norm_hidden) + torch.cuda.synchronize() + assert torch.equal(actual, expected) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = norms_into_cat(embeds, hidden, norm_embeds, norm_hidden) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, expected) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 28e016fd..d6da40d7 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -402,7 +402,7 @@ def attend( ({"draft_loop", "draft_hop"}, "slot 'draft_hop' runs on slot 'hop''s mailbox"), ( {"draft_loop", "draft_fin"}, - "slot 'draft_fin' runs in slot 'draft_hop''s final norm", + "slot 'draft_fin' extends slot 'draft_hop''s fused final norm", ), ( {"draft_loop", "draft_skv2"}, @@ -535,6 +535,95 @@ def all_reduce(h): assert torch.equal(out, x + 100.0 if "fin" in log else 8 * x + x) +def test_only_tep8_draft_fin_uses_the_target_mnnvl_workspace(monkeypatch): + from sglang.srt.layers import flashinfer_mnnvl_cutedsl + + from sgl_nvidia_mega.qwen3_5.draft import layer + + workspace = SimpleNamespace(supports=lambda rows: rows <= 896) + monkeypatch.setattr(flashinfer_mnnvl_cutedsl, "_WORKSPACE", workspace) + parts = SimpleNamespace(slots=frozenset({"draft_fin"})) + + monkeypatch.setattr( + "sglang.srt.runtime_context.get_parallel", + lambda: SimpleNamespace(tp_size=8, moe_ep_size=4), + ) + assert layer._mnnvl_workspace(parts, 896) is workspace + assert layer._mnnvl_workspace(parts, 897) is None + assert layer._mnnvl_workspace(SimpleNamespace(slots=frozenset()), 32) is None + + monkeypatch.setattr( + "sglang.srt.runtime_context.get_parallel", + lambda: SimpleNamespace(tp_size=8, moe_ep_size=1), + ) + assert layer._mnnvl_workspace(parts, 32) is None + + +@pytest.mark.parametrize("deferred", [True, False]) +def test_tep8_draft_fin_uses_the_target_workspace_past_the_hop(monkeypatch, deferred): + from sgl_nvidia_mega.qwen3_5.draft import layer + + rows, calls = 32, [] + x = torch.ones(rows, 4) + + def fused_ar(*, local_contribution, residual, gamma): + calls.append("fused ar") + return local_contribution + gamma, residual + + def fused_fin(*, residual, gamma, **_): + calls.append("fused finalize") + return residual + gamma, residual + + workspace = SimpleNamespace( + all_reduce_residual_rms_norm=fused_ar, + moe_finalize_all_reduce_rms_norm=fused_fin, + ) + monkeypatch.setattr(layer, "_mnnvl_workspace", lambda parts, rows: workspace) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + def mlp(hidden, defer_finalize=False): + calls.append("mlp deferred" if defer_finalize else "mlp") + if not defer_finalize: + return 2 * hidden + return SimpleNamespace( + routed_output=torch.zeros(1, dtype=torch.bfloat16), + expert_weights=torch.zeros(1, dtype=torch.bfloat16), + permuted_indices=torch.zeros(1, dtype=torch.int32), + gated_shared_output=None, + ) + + mlp.supports_deferred_finalize = deferred + post_norm = SimpleNamespace(gemma_weight=10.0) + final_norm = SimpleNamespace(gemma_weight=100.0) + parts = SimpleNamespace( + layer=SimpleNamespace( + input_layernorm=lambda hidden: hidden, + post_attention_layernorm=post_norm, + mlp=mlp, + ), + attention=lambda hidden, *args, rows_per_request, o_proj: hidden, + hop=None, + moe=None, + moe_mkm=None, + slots=frozenset({"draft_hop", "draft_fin"}), + model=SimpleNamespace(norm=final_norm), + ) + + def stock_all_reduce(_): + raise AssertionError("the fused workspace owns both all-reduce points") + + monkeypatch.setattr( + "sglang.srt.distributed.tensor_model_parallel_all_reduce", stock_all_reduce + ) + out = layer.layer(parts, x, None, None, None, None, rows_per_request=1) + if deferred: + assert calls == ["fused ar", "mlp deferred", "fused finalize"] + assert torch.equal(out, x + 100.0) + else: + assert calls == ["fused ar", "mlp", "fused ar"] + assert torch.equal(out, 2 * (x + 10.0) + 100.0) + + class _DraftParts: """Parts with no slot built: the stock model step.""" From 4858568a381385a3569f9b2af26bf7c04e493498 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 16:37:25 -0700 Subject: [PATCH 14/62] fix: align plugin contracts with runtime follow-up --- sgl_nvidia_mega/qwen3_5/attention/split_kv.py | 3 ++ sgl_nvidia_mega/qwen3_5/draft/layer.py | 8 +++-- sgl_nvidia_mega/qwen3_5/gdn/verify.py | 7 ++++- sgl_nvidia_mega/qwen3_5/moe/__init__.py | 1 + sgl_nvidia_mega/qwen3_5/moe/route.py | 21 ++++++++++---- sgl_nvidia_mega/qwen3_5/spec.py | 29 ++++++++++++++++++- test/unit/test_qwen3_5_runtime.py | 3 +- test/unit/test_qwen3_5_spec.py | 9 ++++-- 8 files changed, 68 insertions(+), 13 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py index 7aa16cd7..f3946d2c 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py +++ b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py @@ -80,9 +80,12 @@ # SGLang aa3e34aa routes KV access through kv_layers. This slot already reads # the same caches through EagleBinding; its FP8 target-verify math is unchanged. SGLANG_MIRRORS = { + # Stock, SGLang aa3e34aa, and the SGLang follow-up's CP-lite dispatch; the metadata + # this module mirrors is unchanged by the dispatch around the fixed-q-length call. "srt/layers/attention/trtllm_mha_backend.py:TRTLLMHAAttnBackend.forward_extend": ( "8e36076dafabff66", "1445b85387928caf", + "d7e7c1fe125739e6", ), "srt/layers/attention/trtllm_mha_backend.py:TRTLLMHAAttnBackend._should_use_fused_fp8_path": "e0b14791d523a019", "srt/layers/attention/trtllm_mha_backend.py:TRTLLMHAAttnBackend._fused_fp8_qkv_kv_cache": ( diff --git a/sgl_nvidia_mega/qwen3_5/draft/layer.py b/sgl_nvidia_mega/qwen3_5/draft/layer.py index a37f08a3..f117cc36 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/layer.py +++ b/sgl_nvidia_mega/qwen3_5/draft/layer.py @@ -24,6 +24,7 @@ from __future__ import annotations +import importlib import logging from functools import partial @@ -127,9 +128,10 @@ def _mnnvl_workspace(parts, rows: int): parallel = get_parallel() if parallel.tp_size != 8 or parallel.moe_ep_size <= 1: return None - from sglang.srt.layers import flashinfer_mnnvl_cutedsl - - workspace = flashinfer_mnnvl_cutedsl._WORKSPACE + # Read the live module global after the target initialized its workspace. + workspace = importlib.import_module( + "sglang.srt.layers.flashinfer_mnnvl_cutedsl" + )._WORKSPACE return workspace if workspace is not None and workspace.supports(rows) else None diff --git a/sgl_nvidia_mega/qwen3_5/gdn/verify.py b/sgl_nvidia_mega/qwen3_5/gdn/verify.py index 4a220508..ddf6c970 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/verify.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/verify.py @@ -147,7 +147,12 @@ "b291a3c452e428a5", "1d5fb1eabbb87ec4", ), - "srt/layers/quantization/fp8_utils.py:apply_fp8_linear_bmm_flashinfer": "c66562e1de0008db", + # The follow-up also accepts an already-quantized FP8 input, which is the path + # _out_proj mirrors; the BF16 path remains the stock alternative. + "srt/layers/quantization/fp8_utils.py:apply_fp8_linear_bmm_flashinfer": ( + "c66562e1de0008db", + "c94a4f1a622c2fe0", + ), "srt/layers/quantization/modelopt_quant.py:ModelOptFp8LinearMethod.apply": "91b563cf074bff42", } diff --git a/sgl_nvidia_mega/qwen3_5/moe/__init__.py b/sgl_nvidia_mega/qwen3_5/moe/__init__.py index 8697c1ea..c2c1955d 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/moe/__init__.py @@ -62,6 +62,7 @@ ), "kernels/ops/elementwise/elementwise.py": ("fused_gate_sigmoid_mul_add",), "srt/model_executor/runner/__init__.py": ("get_is_capture_mode",), + "srt/runtime_context.py": ("get_parallel", '"tp_size"', '"moe_ep_size"'), } # The SGLang functions whose body this module mirrors, by the digest of their code diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index e7ec858e..9abe20c4 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -68,6 +68,8 @@ "act_fn", "_enable_silu_fp4_quant_fusion", "_silu_fp4_quant_fused", + "_silu_fp8_quant_fusion", + "_down_takes_static_fp8", ), "srt/models/qwen3_5.py": ("layers", "start_layer", "end_layer", "mlp"), "srt/layers/moe/topk.py": ( @@ -95,7 +97,6 @@ # TEP8 body: the routed op's deferred-finalize output "_make_deferred_finalize_output", "_make_deferred_finalize_output(top_k)", - "_make_deferred_finalize_output(materialized_routing)", ), "srt/layers/moe/utils.py": ( "get_moe_runner_backend", @@ -114,6 +115,7 @@ '"weight_scale"', ), "srt/layers/activation.py": ("SiluAndMul",), + "kernels/ops/quantization/fp8_kernel.py": ("silu_and_mul_static_quant_fp8",), } # The SGLang functions whose body this module mirrors, by the digest of their code @@ -121,7 +123,7 @@ SGLANG_MIRRORS = { "srt/layers/moe/fused_moe_triton/layer.py:FusedMoE.forward_deferred_finalize": "af4b474cbaf33493", "srt/layers/moe/fused_moe_triton/layer.py:FusedMoE._dispatch_with_pre_quant": "dcdefc28cf643081", - "srt/models/qwen2_moe.py:Qwen2MoeMLP.forward": "6091ffb613568a10", + "srt/models/qwen2_moe.py:Qwen2MoeMLP.forward": "bb4c9fb2d7994e97", } @@ -388,6 +390,17 @@ def shared_expert_tail(expert, gate_up: torch.Tensor) -> torch.Tensor: if expert._enable_silu_fp4_quant_fusion and not isinstance(gate_up, tuple): x, _ = expert.down_proj(expert._silu_fp4_quant_fused(gate_up)) return x + if expert._silu_fp8_quant_fusion is None: + expert._silu_fp8_quant_fusion = expert._down_takes_static_fp8() + if expert._silu_fp8_quant_fusion and gate_up.dtype == torch.bfloat16: + from sglang.kernels.ops.quantization.fp8_kernel import ( + silu_and_mul_static_quant_fp8, + ) + + x, _ = expert.down_proj( + silu_and_mul_static_quant_fp8(gate_up, expert.down_proj.input_scale) + ) + return x x = expert.act_fn(gate_up) x, _ = expert.down_proj(x) return x @@ -489,9 +502,7 @@ def _pre_body(mlp, width, x_fp4, x_sf): experts = mlp.experts result = body.run(width, experts, x_fp4, x_sf.view(torch.float8_e4m3fn)) - deferred = _make_deferred_finalize_output( - result, top_k=router.TOP_K, materialized_routing=True - ) + deferred = _make_deferred_finalize_output(result, top_k=router.TOP_K) return experts.dispatcher.combine( combine_input=StandardCombineInput(hidden_states=deferred) ) diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index 9b805499..dbd5e833 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -37,6 +37,7 @@ "SPECULATIVE_DECODING", "MULTI_NODE_TP", "KV_OFFLOAD", + "PD_DISAGGREGATION", "MegakernelSpec", "Numerics", "NumericsKind", @@ -61,6 +62,7 @@ "enable_linear_replayssm_spec", "mamba_ssm_dtype", "startup_weight_load_mode", + "disaggregation_mode", ), "srt/configs/qwen3_5.py": ("text_config",), "srt/kernel_plugins/nvidia/eagle_binding.py|srt/megakernels/eagle_binding.py": ( @@ -130,6 +132,15 @@ "gdn_trt_front.in_proj_ba.rel_l2_vs_stock": 1.05, "gdn_trt_front.in_proj_ba.max_abs_vs_stock": 2.0, }, + # The fused ReplaySSM, gated norm and FP8 output against gdn_verify on the same + # inputs (test_core_is_gdn_verify_numerically). + "gdn_core": { + "gdn_core.replay.rel_l2_vs_gdn_verify": 2e-3, + "gdn_core.norm_fp8.rel_l2_vs_gdn_verify": 2e-2, + "gdn_core.ring_d.rel_l2_vs_gdn_verify": 5e-3, + "gdn_core.ring_k.rel_l2_vs_gdn_verify": 1e-5, + "gdn_core.ring_g.max_abs_vs_gdn_verify": 1e-5, + }, "route_shared": { "route_shared.gate_up.rel_l2_vs_stock": 1.05, "route_shared.gate_up.max_abs_vs_stock": 2.0, @@ -213,9 +224,23 @@ def binding_reasons() -> list[str]: return [] +def pd_reasons(cfg) -> list[str]: + """Why this PD role has no covered step ([] on the decode side).""" + if (cfg.disaggregation_mode or "null") == "prefill": + return [ + "a PD prefill server runs no covered step; enable it on the decode side" + ] + return [] + + def _config_disable_reasons(cfg, hf_config) -> list[str]: runtime.apply_expert_parallel(cfg.ep_size or 1) - return text_config_reasons(hf_config) + binding_reasons() + runtime.slot_reasons() + return ( + text_config_reasons(hf_config) + + binding_reasons() + + runtime.slot_reasons() + + pd_reasons(cfg) + ) def _build(model, hf_config, ctx): @@ -236,6 +261,8 @@ def _build(model, hf_config, ctx): Feature.KV_OFFLOAD, # TEP (MoE TP x EP): the slots in runtime.EP_STOCK_SLOTS stay off Feature.EXPERT_PARALLEL, + # The decode side of PD; pd_reasons refuses the prefill side. + Feature.PD_DISAGGREGATION, } ), numerics=NUMERICS, diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index d6da40d7..dbce5fba 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -360,6 +360,7 @@ def attend( monkeypatch.setattr(split_kv, "cp_lite", CpLite) monkeypatch.setattr(draft_attention, "_q1_label", lambda: label) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) monkeypatch.setattr( "sglang.srt.runtime_context.get_buffer", lambda name, make: None ) @@ -536,7 +537,7 @@ def all_reduce(h): def test_only_tep8_draft_fin_uses_the_target_mnnvl_workspace(monkeypatch): - from sglang.srt.layers import flashinfer_mnnvl_cutedsl + import sglang.srt.layers.flashinfer_mnnvl_cutedsl as flashinfer_mnnvl_cutedsl from sgl_nvidia_mega.qwen3_5.draft import layer diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index e0e04aba..f0f14493 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -425,7 +425,7 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): ) for rows in (7, 14): assert _on(plan.slots(rows, s_tier)) == ["lowm_se", "route", "route_pre_topk"] - assert _on(plan.slots(21, s_tier)) == ["route_topk"] + assert _on(plan.slots(21, s_tier)) == ["lowm_se", "route_topk"] # SGLang offers every multiple of 7 up to MAX_ROWS; every request count is a candidate offered = tuple(range(7, plan.MAX_ROWS + 1, 7)) assert plan.candidates(offered, 7) == offered @@ -574,12 +574,17 @@ def test_a_numeric_slot_turns_the_claim_into_its_tolerances(): assert claim.kind is NumericsKind.TOLERANCE assert dict(claim.tolerances) == spec.NUMERIC_TOLERANCES["gdn_push"] shared = spec.numerics(frozenset({"route", "route_shared"})) - assert dict(shared.tolerances) == spec.NUMERIC_TOLERANCES["route_shared"] + assert dict(shared.tolerances) == { + **spec.NUMERIC_TOLERANCES["route"], + **spec.NUMERIC_TOLERANCES["route_shared"], + } both = spec.numerics(frozenset({"gdn_verify", "gdn_pdl", "gdn_front", "gdn_push"})) assert dict(both.tolerances) == { **spec.NUMERIC_TOLERANCES["gdn_push"], **spec.NUMERIC_TOLERANCES["gdn_front"], } + core = spec.numerics(frozenset({"gdn_verify", "gdn_core"})) + assert dict(core.tolerances) == spec.NUMERIC_TOLERANCES["gdn_core"] attention = spec.numerics(frozenset({"skv2", "mka", "mka_push", "hop"})) assert attention.kind is NumericsKind.TOLERANCE assert dict(attention.tolerances) == { From f39f9364583fa09783b77134a575ffa2d169a98f Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 16:37:29 -0700 Subject: [PATCH 15/62] feat: ship validated bitwise paths by default --- README.md | 22 +++++---- sgl_nvidia_mega/qwen3_5/draft/__init__.py | 8 ++-- sgl_nvidia_mega/qwen3_5/plan.py | 4 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 52 ++++++++++++--------- sgl_nvidia_mega/qwen3_5/spec.py | 6 +-- test/unit/test_qwen3_5_runtime.py | 49 +++++++++++++++++-- test/unit/test_qwen3_5_spec.py | 10 ++++ 7 files changed, 106 insertions(+), 45 deletions(-) diff --git a/README.md b/README.md index 1c7f6eb3..f6a48dc4 100644 --- a/README.md +++ b/README.md @@ -50,23 +50,28 @@ The step is the plugin's own code over the weights SGLang loaded (`deepseek_v41/ With the megakernel on, the default 39 slots run: the ones above, the attention chain, hop B's triggers, and the DSpark draft step's 16 slots. `SGLANG_NVIDIA_MEGAKERNEL_SLOTS` (comma-separated) overrides that list for tests and bisecting; empty runs every step function's stock code. The opt-in slots run only when listed there: `target_lm_head` (the target's LM head on the low-M kernel at 1-8 rows), `lm_head_pinned_tactic` (both low-M LM heads on one fixed tactic per row count, untuned), `hop_b_quad_trigger` (hop B's finalize-publish launches its consumer right after its own wait, batch 1 only), `hop_b_cap16` and `draft_embed`; listing `sm148` there admits the 148 SMs of a B300 instead of the 152 required otherwise. -Qwen3.5 (`qwen3_5/runtime`) runs the same way: per layer the stage boundary (`boundary/`), the GatedDeltaNet or the attention (`gdn/`, `attention/`) and the MoE (`moe/`), each with a stock implementation that reproduces SGLang's verify step bit for bit. It reads the running forward's metadata and pools through SGLang's EAGLE binding (`sglang.srt.kernel_plugins.nvidia.eagle_binding`). Its slots run only when listed in `SGLANG_NVIDIA_MEGAKERNEL_SLOTS`, and each is bit for bit SGLang's: +Qwen3.5 (`qwen3_5/runtime`) runs the same way: per layer the stage boundary (`boundary/`), the GatedDeltaNet or the attention (`gdn/`, `attention/`) and the MoE (`moe/`), each with a stock implementation that reproduces SGLang's verify step bit for bit. It reads the running forward's metadata and pools through SGLang's EAGLE binding (`sglang.srt.kernel_plugins.nvidia.eagle_binding`). With the megakernel enabled, its shape-bounded bitwise paths run by default. `SGLANG_NVIDIA_MEGAKERNEL_SLOTS` overrides that policy for testing and bisecting; an empty value runs the stock paths. | Slot | Rows | Kernel | | --- | --- | --- | -| `route` | 7 | The router GEMM and the NVFP4 quantization of the MoE input in one launch | -| `mkm` (with `route`) | 7 | On `route`'s outputs, the top-10, FC1, SwiGLU, the NVFP4 requantization and FC2 in one launch | -| `gdn_verify` | 7 | The GatedDeltaNet verify chain (conv update, split, ReplaySSM verify, gated norm with the FP8 `out_proj` input) with its launch choices changed | -| `gdn_pdl` (with `gdn_verify`) | 7 | Its ReplaySSM verify launched as a programmatic dependent of the conv | -| `hop` | 7 | The layer boundary's fused all-reduce + residual add + RMSNorm points on an NVLink multicast hop | +| `route_topk`, `mkm_m32` … `mkm_m512` | 1–896 | Materialized stock routing and the shape-selected fused wide routed-expert engines | +| `gdn_verify`, `gdn_pdl`, `gdn_conv` | 7–896 | The opened GatedDeltaNet verify chain, dependent launch, and bitwise conv update | +| `hop`, `hop_fp8` | 1–15 | The layer boundary's fused all-reduce + residual add + RMSNorm, also writing the next static-FP8 input | +| `multi_request` | 2–128 requests | The target graphs keep every offered request count instead of padding to a coarser width | +| `draft_loop`, `draft_extend`, `draft_multi`, `draft_embed` | 1–128 requests | EAGLE draft entries, exact sharded top-1 and embedding exchange; the pre-FC norms write directly into the FC input | -These slots change the arithmetic (numeric: each sums in another order than SGLang's op, within the tolerances `spec.py` declares): +The remaining kernel families are opt-in through that one slot policy. These change the arithmetic within the tolerances `spec.py` declares; `mkm_wide_bal` also stays opt-in because its serving evidence covers TP8 only: | Slot | Rows | Kernel | | --- | --- | --- | +| `route`, `route_gemm`, `mkm`, `mkm_wide` | 1–32 | Fused router fronts and the small routed-expert engines | +| `gdn_core` | 7–224 | ReplaySSM verify, gated norm and FP8 output in one launch | | `gdn_push` (with `gdn_verify`, `hop`) | 7 | The GatedDeltaNet `out_proj` as the hop's GEMM, whose epilogue pushes the partial into the hop's mailbox | -| `gdn_front` (with `gdn_verify`, `gdn_pdl`) | 7 (one request) | The GatedDeltaNet `in_proj_qkvz` (static FP8) and `in_proj_ba` (BF16) projections, the conv update and an L2 prefetch of the verify's state in one launch, then the ReplaySSM verify and the gated norm writing `out_proj`'s FP8 input, each launched dependent on the one before | +| `gdn_front`, `gdn_trt_front` | 7–28 | Fused GatedDeltaNet projections and conv before the dependent verify | | `route_shared` (with `route`) | 1–8 | The shared expert's static-FP8 gate/up in the router front's launch, so the shared expert on the alternate stream waits for that launch only | +| `route_lin`, `skv2`, `mka`, `mka_push` | shape-bounded | Router/attention overlap, split-KV attention and fused attention boundaries | +| `lowm_in`, `lowm_out`, `lowm_se`, `lowm_push` | shape-bounded | Validated low-row projection and shared-expert launches | +| `draft_moe`, `draft_q1`, `draft_xchg`, `draft_hop`, `draft_fin`, `draft_front`, `draft_skv2` | shape-bounded | Fused MTP draft kernels and boundaries | The sampling and acceptance after the verify, and the GatedDeltaNet extend of prefill, stay SGLang's own code. In an SGLang that has the switches, `SGLANG_ENABLE_EAGLE_POST_VERIFY_CUDA_GRAPH=1` runs the sampling and acceptance in CUDA graphs, and `SGLANG_ENABLE_GDN_PREFILL_GRAPH_EXTEND=1` captures the GatedDeltaNet extend inside the breakable prefill CUDA graphs. @@ -124,7 +129,6 @@ The three environment variables are required by the DeepSeek-V4.1 spec; without Qwen3.5, on each of the two trays (`--node-rank 0` and `1`): ```bash -SGLANG_NVIDIA_MEGAKERNEL_SLOTS=route,mkm,gdn_verify,gdn_pdl,hop \ python -m sglang.launch_server \ --model-path nvidia/Qwen3.5-397B-A17B-NVFP4-V2 \ --tp 8 --nnodes 2 --node-rank 0 --dist-init-addr : \ diff --git a/sgl_nvidia_mega/qwen3_5/draft/__init__.py b/sgl_nvidia_mega/qwen3_5/draft/__init__.py index d9461613..712cf21e 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/draft/__init__.py @@ -3,9 +3,9 @@ """Qwen3.5's MTP draft on SM103 (GB300): the megakernel draft spec and its built object. -The draft's forward stays SGLang's (the built object covers no step width). Requested -in ``SGLANG_NVIDIA_MEGAKERNEL_SLOTS``, the built object offers SGLang's EAGLE draft -entries for one request (bs 1), which run inside SGLang's draft graphs: +The draft's forward stays SGLang's (the built object covers no step width). Selected by +the Qwen3.5 runtime policy, the built object offers SGLang's EAGLE draft entries, which +run inside SGLang's draft graphs: - ``draft_loop`` (with ``begin_draft_round``): the whole draft loop of a served draft graph (``loop``); @@ -202,7 +202,7 @@ def _draft_extend( def _requested() -> frozenset: - """The draft's entries and slots in SGLANG_NVIDIA_MEGAKERNEL_SLOTS.""" + """The draft entries and slots selected by the Qwen3.5 runtime policy.""" from sgl_nvidia_mega.qwen3_5 import runtime return runtime.SLOTS & (set(runtime.DRAFT_ENTRIES) | set(runtime.DRAFT_SLOTS)) diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index 135cfa26..b56ee1fd 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -2,8 +2,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Which kernel fills each slot of a covered step, per width (rows), fixed at build. -A slot runs only where it is requested (``SGLANG_NVIDIA_MEGAKERNEL_SLOTS``), built and -inside its row bounds; every other op is SGLang's own. The covered steps are the +A slot runs only where the runtime's default or explicit slot policy requests it, it is +built and the rows are inside its bounds; every other op is SGLang's own. The covered steps are the target verify of 1 to 128 requests (``REQUESTS``: every decode graph batch size SGLang captures up to 128); a width of more than one request is covered only where a slot runs. ``route`` (the router GEMM and the routed experts' NVFP4 input in one launch) runs up to diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index dcf5f42f..2eb27dfe 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -14,15 +14,10 @@ logger = logging.getLogger(__name__) -# The kernel slots that run by default, and the ones that run only when requested -# (until step-level timing and accuracy runs move them to the default list) -SHIPPED_SLOTS: tuple[str, ...] = () -OPT_IN_SLOTS: tuple[str, ...] = ( - "route", +# Bitwise, shape-bounded paths run whenever the Qwen3.5 megakernel is enabled. Numeric +# paths stay behind the one comma-separated slot policy below; BAL is opt-in too. +SHIPPED_SLOTS: tuple[str, ...] = ( "route_topk", - "route_gemm", - "mkm", - "mkm_wide", "mkm_m32", "mkm_m64", "mkm_m128", @@ -30,10 +25,21 @@ "mkm_m512", "gdn_verify", "gdn_pdl", - "gdn_core", - # the verify conv with every x row loaded first and the windows in one store a channel (bitwise) "gdn_conv", "hop", + "hop_fp8", + "multi_request", + "draft_loop", + "draft_extend", + "draft_multi", + "draft_embed", +) +OPT_IN_SLOTS: tuple[str, ...] = ( + "route", + "route_gemm", + "mkm", + "mkm_wide", + "gdn_core", "gdn_push", "gdn_front", "gdn_trt_front", @@ -46,12 +52,10 @@ "lowm_out", "lowm_se", "lowm_push", - "hop_fp8", ) -# Steps of more than one request (TP8's decode graphs of batch sizes 2 to 128), covered only -# when requested in the same variable: covering them costs the one-request step about -# 0.3% at TP8 -OPT_IN_COVERAGE: tuple[str, ...] = ("multi_request",) +# Optional coverage names requested in the same variable. Multi-request coverage is +# shipped because every covered width keeps its own shape instead of padding. +OPT_IN_COVERAGE: tuple[str, ...] = () # Graph or build variants of a slot, requested in the same variable: no width of their # own. BAL stays opt-in because its serving evidence covers TP8 only. # skv2_cp: one long request past the measured crossover runs CP-lite in place of SKV2, @@ -123,13 +127,17 @@ def apply_expert_parallel(ep_size: int) -> frozenset: def slot_reasons() -> list[str]: """Why the requested slots cannot be served ([] when they can).""" - known = ( - *SHIPPED_SLOTS, - *OPT_IN_SLOTS, - *OPT_IN_COVERAGE, - *OPT_IN_VARIANTS, - *DRAFT_ENTRIES, - *DRAFT_SLOTS, + known = tuple( + dict.fromkeys( + ( + *SHIPPED_SLOTS, + *OPT_IN_SLOTS, + *OPT_IN_COVERAGE, + *OPT_IN_VARIANTS, + *DRAFT_ENTRIES, + *DRAFT_SLOTS, + ) + ) ) reasons = [ f"{SLOTS_ENV}: unknown slot {name!r} (known: {', '.join(known) or 'none'})" diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index dbd5e833..f531c525 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -108,9 +108,9 @@ ) # Against SGLang's forward on the same GPU (test/e2e/test_qwen3_5_stock_parity.py): the -# stock step and the bitwise slots (mkm, gdn_verify, gdn_pdl, hop) reproduce it bit for -# bit. A numeric slot's outputs are held to stock or fp64 on the same quantized operands -# (test/kernels), at most these multiples of the error of SGLang's own op. +# stock step and the shipped slots reproduce it bit for bit. A numeric opt-in slot's +# outputs are held to stock or fp64 on the same quantized operands (test/kernels), at +# most these multiples of the error of SGLang's own op. NUMERIC_TOLERANCES = { # Four router GEMV chains sum in another order than SGLang's cuBLAS GEMM. "route": {"route.router_logits.bf16_ulp_vs_stock": 1.0}, diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index dbce5fba..4d37f5a3 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -46,6 +46,37 @@ def _batch(verify=True): ) +def test_the_default_policy_ships_only_bitwise_shape_bounded_paths(monkeypatch): + monkeypatch.delenv(runtime.SLOTS_ENV, raising=False) + assert runtime._requested_slots() == frozenset( + { + "route_topk", + "mkm_m32", + "mkm_m64", + "mkm_m128", + "mkm_m256", + "mkm_m512", + "gdn_verify", + "gdn_pdl", + "gdn_conv", + "hop", + "hop_fp8", + "multi_request", + "draft_loop", + "draft_extend", + "draft_multi", + "draft_embed", + } + ) + monkeypatch.setattr(runtime, "SLOTS", runtime._requested_slots()) + assert runtime.slot_reasons() == [] + + +def test_an_explicit_empty_policy_keeps_every_slot_off(monkeypatch): + monkeypatch.setenv(runtime.SLOTS_ENV, "") + assert runtime._requested_slots() == frozenset() + + def test_only_target_verify_steps_of_built_widths_are_covered(): built = _verify() ids = torch.zeros(7, dtype=torch.int64) @@ -92,7 +123,7 @@ def test_the_build_covers_more_requests_only_where_a_slot_runs( assert built.widths == want and built.summary() == "slots none (stock)" -def test_the_build_hands_over_the_binding_and_every_entry_stays_sglangs(): +def test_the_build_hands_over_the_binding_and_ships_the_bitwise_draft_entries(): from sglang.kernels.plugin_kernels.nvidia.spec import ENTRIES from sgl_nvidia_mega.qwen3_5 import draft @@ -104,7 +135,11 @@ def test_the_build_hands_over_the_binding_and_every_entry_stays_sglangs(): assert built.plan(ids, _batch()) == Step(mode="target_verify", rows=7) target, mtp = built, draft.SPEC.build("lm", None, _draft_ctx()) assert [n for n in ENTRIES if getattr(target, n, None) is not None] == [] - assert [n for n in ENTRIES if getattr(mtp, n, None) is not None] == [] + assert [n for n in ENTRIES if getattr(mtp, n, None) is not None] == [ + "draft_loop", + "draft_extend", + "begin_draft_round", + ] def test_verify_sample_is_refused_with_the_switch_that_replaces_it(monkeypatch): @@ -201,11 +236,15 @@ def test_the_mtp_draft_covers_no_step(): ) assert built.plan(torch.zeros(7, dtype=torch.int64), _batch()) is None assert (draft.ARCHITECTURE, "sgl_nvidia_mega.qwen3_5.draft:SPEC") in plugin.SPECS - # without a request every draft phase stays SGLang's - assert [n for n in ENTRIES if getattr(built, n, None) is not None] == [] + # Entries replace SGLang's draft graphs, not a model-forward width. + assert [n for n in ENTRIES if getattr(built, n, None) is not None] == [ + "draft_loop", + "draft_extend", + "begin_draft_round", + ] -def test_the_draft_entries_are_offered_only_on_request(monkeypatch): +def test_an_explicit_policy_controls_the_draft_entries(monkeypatch): from sglang.kernels.plugin_kernels.nvidia.spec import ENTRIES from sgl_nvidia_mega.qwen3_5 import draft diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index f0f14493..349f9c2e 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -597,3 +597,13 @@ def test_a_numeric_slot_turns_the_claim_into_its_tolerances(): **spec.NUMERIC_TOLERANCES["gdn_trt_front"], **spec.NUMERIC_TOLERANCES["route_lin"], } + + +def test_the_shipped_policy_claims_bitwise_numerics(): + from sglang.kernels.plugin_kernels.nvidia.spec import NumericsKind + + from sgl_nvidia_mega.qwen3_5 import draft + + shipped = frozenset(runtime.SHIPPED_SLOTS) + assert spec.numerics(shipped).kind is NumericsKind.BITWISE + assert draft.numerics(shipped).kind is NumericsKind.BITWISE From 0fbc03c6582aa1a301cc309c64acb6a9462c1f20 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 16:42:58 -0700 Subject: [PATCH 16/62] fix: complete migrated kernel path guards --- sgl_nvidia_mega/qwen3_5/boundary/hop.py | 2 +- sgl_nvidia_mega/qwen3_5/gdn/front.py | 2 +- sgl_nvidia_mega/qwen3_5/moe/__init__.py | 2 +- test/unit/test_qwen3_5_steps.py | 25 ++++++++++++++++++------- 4 files changed, 21 insertions(+), 10 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/boundary/hop.py b/sgl_nvidia_mega/qwen3_5/boundary/hop.py index b12a5251..81e3eef5 100644 --- a/sgl_nvidia_mega/qwen3_5/boundary/hop.py +++ b/sgl_nvidia_mega/qwen3_5/boundary/hop.py @@ -373,7 +373,7 @@ def moe_input(hop, layer, hidden, st): _engage("MoE input", rows) scale = lowm.moe_input_fp8_scale(layer, st) - if hop._pushed_rl: + if st.slots.route_lin and hop._pushed_rl: # Slot route_lin: the mixer pushed its router partial too; this hop also gives # the MoE its logits and NVFP4 input (moe/route.py routed_lin). lid, h0 = st.scratch.get("rlin.h0", (None, None)) diff --git a/sgl_nvidia_mega/qwen3_5/gdn/front.py b/sgl_nvidia_mega/qwen3_5/gdn/front.py index 2453d7b9..99f57a9d 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/front.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/front.py @@ -314,7 +314,7 @@ def step(gdn, x: torch.Tensor, st, opts) -> torch.Tensor: return _fused( gdn, conv, - projected[:, CONV_CHANNELS:].view(ROWS, attn.num_v_heads, attn.head_v_dim), + projected[:, CONV_CHANNELS:].view(rows, attn.num_v_heads, attn.head_v_dim), ba[:, 8:], ba[:, :8], st, diff --git a/sgl_nvidia_mega/qwen3_5/moe/__init__.py b/sgl_nvidia_mega/qwen3_5/moe/__init__.py index c2c1955d..ef1b8df1 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/moe/__init__.py @@ -251,7 +251,7 @@ def _router_experts(mlp, x, defer: bool, st): from sglang.srt.layers.moe.topk import TopKOutputChecker if defer and st.slots.route: - rl = st.scratch.pop("rlin.moe", None) + rl = st.scratch.pop("rlin.moe", None) if st.slots.route_lin else None if rl is not None and rl[0] == mlp.layer_id and st.slots.mkm: # Slot route_lin produced the logits and NVFP4 input in the preceding hop. return route.routed_lin(mlp, x, st.bound["moe"], *rl[1:]) diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 37dd01f9..8c78cbbc 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -338,6 +338,7 @@ def test_the_route_shared_slot_starts_the_shared_expert_after_the_router_launch( weight=k["sh_w"], input_scale=k["sh_in"], weight_scale=k["sh_ws"] ), _enable_silu_fp4_quant_fusion=False, + _silu_fp8_quant_fusion=False, act_fn=log.op("act", "act_out"), down_proj=log.op("down_proj", ("down_out", None)), ) @@ -481,7 +482,9 @@ def combine(combine_input, log=log, combined=combined): mlp.experts.dispatcher.combine = combine bound = SimpleNamespace( - mkm={3: "maps"}, identity=torch.arange(70, dtype=torch.int32) + mkm={3: "maps"}, + identity=torch.arange(70, dtype=torch.int32), + m16_max_rows=16, ) st = SimpleNamespace(slots=slots, bound={"moe": bound}) fused = patch.object( @@ -513,7 +516,11 @@ def combine(combine_input, log=log, combined=combined): "c(512,)", "g2(512,)", ), - (("routing", None),), + ( + ("m16_max_rows", 16), + ("routing", None), + ("wide_bal", False), + ), ), ("combine", "main"), ] @@ -1372,9 +1379,10 @@ def test_trt_lowm_takes_trts_gdn_tactics_and_refuses_an_unprepared_launch( from sgl_nvidia_mega.qwen3_5.sm103.gdn import trt_front conv_w = torch.zeros(d, 4, dtype=torch.bfloat16) - state = torch.zeros(3, d, 3, dtype=torch.bfloat16) - windows = torch.zeros(3, d, 9, dtype=torch.bfloat16).as_strided( - (3, 7, d, 3), (d * 9, 1, 9, 1) + pool_requests = max(3, requests) + state = torch.zeros(pool_requests, d, 3, dtype=torch.bfloat16) + windows = torch.zeros(pool_requests, d, 9, dtype=torch.bfloat16).as_strided( + (pool_requests, 7, d, 3), (d * 9, 1, 9, 1) ) for act, tactic in (("fp8", tactic_fp8), ("quant", tactic_quant)): s = trt_front.spec(act, qkvz, ba, conv_w, state, windows, requests) @@ -1395,7 +1403,10 @@ def test_trt_lowm_takes_trts_gdn_tactics_and_refuses_an_unprepared_launch( and not s.release_after_wait and not s.unsafe_skip_wait ) - assert s.conv_shapes[3] == ((3, 7, d, 3), (d * 9, 1, 9, 1)) + assert s.conv_shapes[3] == ( + (pool_requests, 7, d, 3), + (d * 9, 1, 9, 1), + ) with pytest.raises(RuntimeError, match="not prepared"): trt_lowm.run(s, out=torch.zeros(7 * requests, qkvz, dtype=torch.bfloat16)) with patch.object(torch.cuda, "is_current_stream_capturing", lambda: True): @@ -2617,7 +2628,7 @@ def _pushing_hop(log): from sgl_nvidia_mega.qwen3_5.sm103 import hop as sm103_hop hop = object.__new__(sm103_hop.Hop) - hop.max_m, hop._pushed = 8, None + hop.max_m, hop._pushed, hop._pushed_rl = 8, None, False hop._args = lambda m, residual, gamma, scale: (SimpleNamespace(m=m), ("out",)) hop._launch = lambda name, args, ctas, cluster=None: log.append(name) hop._consume = lambda args: log.append("consume") From 5685a5717e7bff07b87107bded7e1a7acebb8f2e Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 17:37:54 -0700 Subject: [PATCH 17/62] fix: accept runtime-owned megakernel slots --- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 5 +++++ test/unit/test_qwen3_5_runtime.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 2eb27dfe..8e91b682 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -83,6 +83,9 @@ "draft_q1_cp", "draft_skv2_cp", ) +# SGLang consumes these names from the shared slot policy itself. They are valid +# requests, but are intentionally not built or dispatched by this plugin. +RUNTIME_ONLY_SLOTS: tuple[str, ...] = ("gdn_commit", "sharded_gumbel") SLOTS_ENV = "SGLANG_NVIDIA_MEGAKERNEL_SLOTS" # These slots run all 512 experts themselves. Under expert parallelism, keep the # route-only and boundary slots but let SGLang's EP path run the routed experts. @@ -136,6 +139,7 @@ def slot_reasons() -> list[str]: *OPT_IN_VARIANTS, *DRAFT_ENTRIES, *DRAFT_SLOTS, + *RUNTIME_ONLY_SLOTS, ) ) ) @@ -417,6 +421,7 @@ def build(model, ctx): - set(DRAFT_SLOTS) - set(OPT_IN_COVERAGE) - set(OPT_IN_VARIANTS) + - set(RUNTIME_ONLY_SLOTS) ) # mkm's 16-token engine has nothing to serve where no covered step is wider than 7 rows vacuous = ( diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 4d37f5a3..5761d220 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -77,6 +77,20 @@ def test_an_explicit_empty_policy_keeps_every_slot_off(monkeypatch): assert runtime._requested_slots() == frozenset() +def test_runtime_only_slots_are_accepted_but_not_built_by_the_plugin(monkeypatch): + requested = frozenset({"gdn_commit", "sharded_gumbel"}) + monkeypatch.setattr(runtime, "SLOTS", requested) + assert runtime.slot_reasons() == [] + + model = SimpleNamespace( + model=SimpleNamespace(start_layer=0, end_layer=0, layers=[]) + ) + ctx = SimpleNamespace( + widths=(7,), tokens_per_request=7, deadline_s=60.0, binding="EagleBinding" + ) + assert runtime.build(model, ctx).summary() == "slots none (stock)" + + def test_only_target_verify_steps_of_built_widths_are_covered(): built = _verify() ids = torch.zeros(7, dtype=torch.int64) From 6708f924bd25e5e4a1956701c681fb974f9676b8 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 17:52:13 -0700 Subject: [PATCH 18/62] fix: compile the pinned GDN front --- sgl_nvidia_mega/qwen3_5/sm103/gdn/src/front.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/front.cu b/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/front.cu index 267401a9..3efba9a2 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/front.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/gdn/src/front.cu @@ -411,7 +411,7 @@ __device__ __forceinline__ void front_kernel(const Params& p) { #pragma unroll for (int i = 0; i < 2; ++i) if (g < p.m) { - const auto* x = static_cast(p.x) + (req * p.m + g) * 4096 + rank * 512 + + const auto* x = static_cast(p.x) + g * 4096 + rank * 512 + (2 * warp + i) * 32 + 4 * t; raw[i][0] = *reinterpret_cast(x); raw[i][1] = *reinterpret_cast(x + 16); From 9b226e219b4bed40a01f42086dda8571b54f0d0c Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 18:16:21 -0700 Subject: [PATCH 19/62] fix: treat multi-request as coverage --- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 7 ++++--- test/unit/test_qwen3_5_runtime.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 8e91b682..c5f205d1 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -53,9 +53,10 @@ "lowm_se", "lowm_push", ) -# Optional coverage names requested in the same variable. Multi-request coverage is -# shipped because every covered width keeps its own shape instead of padding. -OPT_IN_COVERAGE: tuple[str, ...] = () +# Coverage names requested in the same variable have no kernel build result of their +# own. Multi-request coverage is shipped because every covered width keeps its own +# shape instead of padding. +OPT_IN_COVERAGE: tuple[str, ...] = ("multi_request",) # Graph or build variants of a slot, requested in the same variable: no width of their # own. BAL stays opt-in because its serving evidence covers TP8 only. # skv2_cp: one long request past the measured crossover runs CP-lite in place of SKV2, diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 5761d220..f24fb3fe 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -91,6 +91,19 @@ def test_runtime_only_slots_are_accepted_but_not_built_by_the_plugin(monkeypatch assert runtime.build(model, ctx).summary() == "slots none (stock)" +def test_multi_request_coverage_has_no_kernel_build_result(monkeypatch): + monkeypatch.setattr(runtime, "SLOTS", frozenset({"multi_request"})) + model = SimpleNamespace( + model=SimpleNamespace(start_layer=0, end_layer=0, layers=[]) + ) + ctx = SimpleNamespace( + widths=(7, 14), tokens_per_request=7, deadline_s=60.0, binding="EagleBinding" + ) + + built = runtime.build(model, ctx) + assert built.widths == (7,) + + def test_only_target_verify_steps_of_built_widths_are_covered(): built = _verify() ids = torch.zeros(7, dtype=torch.int64) From 45b7e22bc8565d59a59009e98a754ed21e2f698b Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 20:43:20 -0700 Subject: [PATCH 20/62] fix: give draft MoE the complete slot schema --- sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py | 2 +- test/unit/test_qwen3_5_steps.py | 11 +++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py b/sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py index b0080813..4a976391 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py +++ b/sgl_nvidia_mega/qwen3_5/draft/moe_mkm.py @@ -118,7 +118,7 @@ def state(self, rows: int): [k for k, v in on.items() if v], ) return SimpleNamespace( - slots=SimpleNamespace(**on), bound={"moe": self.bound}, scratch={} + slots=plan.Slots(**on), bound={"moe": self.bound}, scratch={} ) def __call__(self, layer, x: torch.Tensor): diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 8c78cbbc..86a66049 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -586,6 +586,17 @@ def meta(dtype, *shape): assert mkm.unsupported(mlp).startswith("an expert activation") +def test_the_draft_mkm_uses_the_complete_slot_schema(monkeypatch): + """Draft widths leave unrelated target slots off instead of omitting them.""" + from sgl_nvidia_mega.qwen3_5.draft.moe_mkm import DraftMkm + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + slots = DraftMkm(bound="bound", wide=False).state(7).slots + assert isinstance(slots, plan.Slots) + assert slots.route and slots.mkm + assert not slots.route_lin and not slots.route_pre_topk + + def test_the_route_slot_is_not_built_for_a_layer_it_cannot_serve(monkeypatch): from sgl_nvidia_mega.qwen3_5.moe import route From 7cc7235551158ffb9d04a8c46a11f78914241f78 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 20:43:44 -0700 Subject: [PATCH 21/62] fix: bound GDN front prep by pool capacity --- .../qwen3_5/sm103/gdn/trt_front.py | 6 ++-- test/unit/test_qwen3_5_steps.py | 28 +++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/gdn/trt_front.py b/sgl_nvidia_mega/qwen3_5/sm103/gdn/trt_front.py index c5fc2f5d..33b4a813 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/gdn/trt_front.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/gdn/trt_front.py @@ -73,10 +73,12 @@ def spec( def prepare(qkvz: int, ba: int, conv_weight, conv_state, windows) -> None: - """Compile both activation forms for 1..MAX_REQUESTS requests on these layouts (not inside a graph capture).""" + """Compile both activation forms through the request-window pool's capacity, + capped at ``MAX_REQUESTS`` (not inside a graph capture).""" from sgl_nvidia_mega.qwen3_5.sm103 import trt_lowm - for requests in range(1, MAX_REQUESTS + 1): + max_requests = min(MAX_REQUESTS, int(windows.shape[0])) + for requests in range(1, max_requests + 1): for act in MODES: s = spec(act, qkvz, ba, conv_weight, conv_state, windows, requests) trt_lowm.prepare(s) diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 86a66049..91303b37 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -1371,6 +1371,34 @@ def test_the_trt_gdn_front_prepares_on_the_first_gdn_layers_pools(): ) +def test_the_trt_gdn_front_stops_at_the_request_window_capacity(monkeypatch): + """A small max-running-requests pool does not prepare impossible widths.""" + from sgl_nvidia_mega.qwen3_5.sm103 import trt_lowm + from sgl_nvidia_mega.qwen3_5.sm103.gdn import trt_front + + compiled = [] + monkeypatch.setattr(trt_front, "_SPECS", {}) + monkeypatch.setattr( + trt_front, + "spec", + lambda act, qkvz, ba, conv_weight, conv_state, windows, requests: ( + act, + requests, + ), + ) + monkeypatch.setattr(trt_lowm, "prepare", compiled.append) + windows = SimpleNamespace(shape=(2, 7, 1536, 3)) + trt_front.prepare(2560, 16, "weight", "state", windows) + + assert compiled == [("fp8", 1), ("quant", 1), ("fp8", 2), ("quant", 2)] + assert set(trt_front._SPECS) == { + ("fp8", 2560, 1), + ("quant", 2560, 1), + ("fp8", 2560, 2), + ("quant", 2560, 2), + } + + @pytest.mark.parametrize( "qkvz, ba, d, requests, tactic_fp8, tactic_quant", [ From b24bfe217233569577b34cd56844eadf37b8fa5e Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 21:26:10 -0700 Subject: [PATCH 22/62] fix: validate routed MoE under expert parallelism --- sgl_nvidia_mega/qwen3_5/moe/route.py | 8 ++++++-- test/unit/test_qwen3_5_steps.py | 25 +++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index 9abe20c4..2431657c 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -307,13 +307,17 @@ def unsupported(mlp) -> str: def topk_unsupported(model) -> str: """Why slot route_topk cannot serve every MoE layer of ``model`` ("" when it can): the router front's gate and input quantizer, and the trtllm-gen weights with a - softmax top-10 renormalized over 512 local experts (mkm.unsupported).""" + softmax top-10 renormalized over 512 experts. Under expert parallelism only the + materialized-routing contract applies; the plugin's full-expert engines do not.""" from sgl_nvidia_mega.qwen3_5.moe import mkm m = model.model for i in range(m.start_layer, m.end_layer): mlp = m.layers[i].mlp - reason = unsupported(mlp) or mkm.unsupported(mlp) + expert_parallel = mlp.experts.num_local_experts != mlp.experts.num_experts + reason = unsupported(mlp) or ( + mkm.routing_unsupported(mlp) if expert_parallel else mkm.unsupported(mlp) + ) if reason: return f"layer {i}: {reason}" return "" diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 91303b37..dfb14a44 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -586,6 +586,31 @@ def meta(dtype, *shape): assert mkm.unsupported(mlp).startswith("an expert activation") +def test_route_topk_checks_only_materialized_routing_under_expert_parallelism( + monkeypatch, +): + from sgl_nvidia_mega.qwen3_5.moe import mkm, route + + experts = SimpleNamespace(num_experts=512, num_local_experts=128) + mlp = SimpleNamespace(experts=experts) + model = SimpleNamespace( + model=SimpleNamespace( + start_layer=0, end_layer=1, layers=[SimpleNamespace(mlp=mlp)] + ) + ) + checked = [] + monkeypatch.setattr(route, "unsupported", lambda _: "") + monkeypatch.setattr( + mkm, "routing_unsupported", lambda _: checked.append("routing") or "" + ) + monkeypatch.setattr( + mkm, "unsupported", lambda _: pytest.fail("full-expert engine checked under EP") + ) + + assert route.topk_unsupported(model) == "" + assert checked == ["routing"] + + def test_the_draft_mkm_uses_the_complete_slot_schema(monkeypatch): """Draft widths leave unrelated target slots off instead of omitting them.""" from sgl_nvidia_mega.qwen3_5.draft.moe_mkm import DraftMkm From e514a2d02a659cc3877be42133e3d7c456eed784 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 26 Oct 2026 21:54:36 -0700 Subject: [PATCH 23/62] perf: keep static weights ahead of PDL waits --- sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu | 14 ++++++++++++++ sgl_nvidia_mega/qwen3_5/sm103/hop/src/hop.cu | 4 ++++ sgl_nvidia_mega/qwen3_5/sm103/hop/src/q35_hop.cuh | 9 +++++++++ sgl_nvidia_mega/qwen3_5/sm103/hop/src/rlin.cu | 9 +++++++++ 4 files changed, 36 insertions(+) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu b/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu index 11530ff3..dac44ef9 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/src/mka.cu @@ -243,6 +243,12 @@ __device__ __forceinline__ void mka_prologue_kernel(const ProParams p) { if (is_k) kvs = __ldg(p.k_scale); if (!needs_norm && !is_gate) kvs = __ldg(p.v_scale); + if (!kEarlyTrigger) { + uint32_t weight_dep = __float_as_uint(sx) ^ __float_as_uint(sw); +#pragma unroll + for (int s = 0; s < 16; ++s) weight_dep ^= a[s][0].x ^ a[s][1].x; + q35hop::pin_before_wait(weight_dep, reinterpret_cast(p.w)); + } pdl_wait(); // Keep this release after the wait: q35mka_gate reads x (the hop's output) before its own wait. if (kEarlyTrigger) pdl_trigger(); @@ -459,6 +465,10 @@ __device__ __forceinline__ void mka_epilogue_kernel(const EpiParams p) { for (int b = 0; b < 2; ++b) a[s][b] = ldg_nc_v4(wb + (s * 2 + b) * 32); } const float sx = __ldg(p.sx), sw = __ldg(p.sw); + uint32_t weight_dep = __float_as_uint(sx) ^ __float_as_uint(sw); +#pragma unroll + for (int s = 0; s < kSteps; ++s) weight_dep ^= a[s][0].x ^ a[s][1].x; + q35hop::pin_before_wait(weight_dep, reinterpret_cast(p.w)); pdl_wait(); if (kEarlyTrigger) pdl_trigger(); @@ -741,6 +751,10 @@ __device__ __forceinline__ void mka_epilogue_wide_kernel(const EpiParams p) { for (int b = 0; b < 2; ++b) a[s][b] = ldg_nc_v4(wb + (s * 2 + b) * 32); } const float sx = __ldg(p.sx), sw = __ldg(p.sw); + uint32_t weight_dep = __float_as_uint(sx) ^ __float_as_uint(sw); +#pragma unroll + for (int s = 0; s < kSteps; ++s) weight_dep ^= a[s][0].x ^ a[s][1].x; + q35hop::pin_before_wait(weight_dep, reinterpret_cast(p.w)); pdl_wait(); if (kEarlyTrigger) pdl_trigger(); diff --git a/sgl_nvidia_mega/qwen3_5/sm103/hop/src/hop.cu b/sgl_nvidia_mega/qwen3_5/sm103/hop/src/hop.cu index d638e76a..76c2e3ea 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/hop/src/hop.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/hop/src/hop.cu @@ -188,6 +188,10 @@ __device__ __forceinline__ void oproj_push_kernel(const OprojArgs& o) { for (int b = 0; b < 2; ++b) a[st][b] = ldg_nc_v4(wb + (st * 2 + b) * 32); } const float alpha = __fmul_rn(*o.sx, *o.sw); + uint32_t weight_dep = __float_as_uint(alpha); +#pragma unroll + for (int st = 0; st < S; ++st) weight_dep ^= a[st][0].x ^ a[st][1].x; + pin_before_wait(weight_dep, reinterpret_cast(o.w)); pdl_wait(); const int m = o.m; const bool tok_ok = g < m; diff --git a/sgl_nvidia_mega/qwen3_5/sm103/hop/src/q35_hop.cuh b/sgl_nvidia_mega/qwen3_5/sm103/hop/src/q35_hop.cuh index 9614f84f..ad5c3505 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/hop/src/q35_hop.cuh +++ b/sgl_nvidia_mega/qwen3_5/sm103/hop/src/q35_hop.cuh @@ -56,6 +56,15 @@ struct Site { __device__ __forceinline__ void pdl_wait() { asm volatile("griddepcontrol.wait;" ::: "memory"); } __device__ __forceinline__ void pdl_trigger() { asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); } +// Keep the named static loads above a PDL wait. Without the data dependency, +// ptxas may sink ld.global.nc instructions below the wait. +__device__ __forceinline__ void pin_before_wait(uint32_t dep, uint64_t never) { + asm volatile( + "{ .reg .pred a, b; setp.eq.u64 a, %1, 0; setp.eq.u32 b, %0, 1; " + "and.pred a, a, b; @a trap; }" + :: "r"(dep), "l"(never) + : "memory"); +} __device__ __forceinline__ unsigned long long gtimer() { unsigned long long t; asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); diff --git a/sgl_nvidia_mega/qwen3_5/sm103/hop/src/rlin.cu b/sgl_nvidia_mega/qwen3_5/sm103/hop/src/rlin.cu index c3116600..fb2935f3 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/hop/src/rlin.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/hop/src/rlin.cu @@ -140,6 +140,10 @@ __device__ __forceinline__ void oproj_rows(const OprojArgs& o, const RlArgs& r, for (int b = 0; b < 2; ++b) a[st][b] = ldg_nc_v4(wb + (st * 2 + b) * 32); } const float alpha = __fmul_rn(*o.sx, *o.sw); + uint32_t weight_dep = __float_as_uint(alpha); +#pragma unroll + for (int st = 0; st < S; ++st) weight_dep ^= a[st][0].x ^ a[st][1].x; + pin_before_wait(weight_dep, reinterpret_cast(o.w)); pdl_wait(); const int m = o.m; const bool tok_ok = g < m; @@ -199,6 +203,11 @@ __device__ __forceinline__ void router_rows(const OprojArgs& o, const RlArgs& r, #pragma unroll for (int mt = 0; mt < 2; ++mt) a[st][mt] = ldg_nc_v4(wb + (st * 2 + mt) * 32); } + uint32_t weight_dep = 0; +#pragma unroll + for (int st = 0; st < kRlSteps; ++st) + weight_dep ^= a[st][0].x ^ a[st][1].x; + pin_before_wait(weight_dep, reinterpret_cast(r.rw)); pdl_wait(); const int m = o.m; // stage: token rows < m; x: 7 x 1024 B = 64 chunks of 16 B a row; h0 slice: 7 x 1024 B = 64 chunks a row. All four From e69926b9c2693af6fa7b915b25ae6299778fae79 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 27 Oct 2026 01:57:25 -0700 Subject: [PATCH 24/62] perf: restore four-request one-shot hop --- sgl_nvidia_mega/qwen3_5/boundary/hop.py | 11 ++++++++--- sgl_nvidia_mega/qwen3_5/lowm.py | 3 ++- sgl_nvidia_mega/qwen3_5/plan.py | 13 ++++++++++--- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 7 +++++++ test/unit/test_qwen3_5_spec.py | 14 +++++++++++++- 5 files changed, 40 insertions(+), 8 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/boundary/hop.py b/sgl_nvidia_mega/qwen3_5/boundary/hop.py index 81e3eef5..ada47647 100644 --- a/sgl_nvidia_mega/qwen3_5/boundary/hop.py +++ b/sgl_nvidia_mega/qwen3_5/boundary/hop.py @@ -140,6 +140,7 @@ class Bound(msgspec.Struct, frozen=True): """The boundary package's build result (``st.bound["boundary"]``).""" hop: Optional[object] = None # sm103.hop.Hop for every covered width, or None + mid: bool = False # its one-shot mailbox serves past 15 rows (slot hop_mid) reason: str = "" # why there is no hop @@ -157,8 +158,12 @@ def build(model, ctx) -> Bound: if "hop" not in runtime.SLOTS: return Bound(reason="slot off") - # the widest step the hop covers: its mailbox's rows - rows = plan.widest(ctx.widths, plan.HOP_MAX_ROWS) + # The widest step the one-shot mailbox covers. hop_mid extends it through four + # requests; without the opt-in, the shipped bitwise bound stays unchanged. + max_rows = ( + plan.HOP_MID_MAX_ROWS if "hop_mid" in runtime.SLOTS else plan.HOP_MAX_ROWS + ) + rows = plan.widest(ctx.widths, max_rows) if not rows: return Bound(reason=f"no width of {list(ctx.widths)} within its rows") m = model.model @@ -178,7 +183,7 @@ def build(model, ctx) -> Bound: ) global _ACTIVE _ACTIVE = hop - return Bound(hop=hop) + return Bound(hop=hop, mid=rows > plan.HOP_MAX_ROWS) def active() -> Optional[object]: diff --git a/sgl_nvidia_mega/qwen3_5/lowm.py b/sgl_nvidia_mega/qwen3_5/lowm.py index 15f02f61..a5b613cd 100644 --- a/sgl_nvidia_mega/qwen3_5/lowm.py +++ b/sgl_nvidia_mega/qwen3_5/lowm.py @@ -243,6 +243,7 @@ def build(model, ctx) -> Bound: se_rows = plan.lowm_se_pre_rows(par.tp_size, tep8_m) or se_rows want_push = want_out and "lowm_push" in SLOTS and "hop" in SLOTS want_fp8 = "hop_fp8" in SLOTS and "hop" in SLOTS + fp8_rows = plan.HOP_MID_MAX_ROWS if "hop_mid" in SLOTS else plan.HOP_MAX_ROWS want = want_in or want_out if not (want or want_se or want_fp8): return Bound(reason="stock (SGLANG_NVIDIA_MEGAKERNEL_SLOTS)") @@ -324,7 +325,7 @@ def build(model, ctx) -> Bound: for w in ctx.widths: if w <= se_rows: launches.add((w, k, n, 8, 1, 1, True)) - if want_fp8 and w <= plan.HOP_MAX_ROWS: + if want_fp8 and w <= fp8_rows: launches.add((w, k, n, 8, 0, 1, True)) launches.add((w, kd, nd, 0, 0, 2, False)) for w, k, n8, n16, act, epi8, f32 in sorted(launches): diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index b56ee1fd..21cfa4a2 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -17,7 +17,8 @@ GatedDeltaNet verify chain with its launch choices changed) covers every width, with ``gdn_core`` (numeric: its ReplaySSM verify, gated norm and FP8 output in one launch) up to 224 rows, and ``hop`` (the boundary's fused all-reduce points on the decode hop) the -widths where FlashInfer's low-latency kernels, which it equals bit for bit, run. +widths where FlashInfer's low-latency kernels, which it equals bit for bit, run, or with +``hop_mid`` (numeric: its one-shot protocol past FlashInfer's) up to 28 rows. ``gdn_push`` (numeric: out_proj as the hop's push GEMM, on gdn_verify's FP8 norm) runs up to 8 rows. ``gdn_front`` (numeric: the GDN projections and conv in one launch, then gdn_pdl's verify) covers one request of 7 rows; ``gdn_trt_front`` replaces its projection @@ -69,6 +70,8 @@ MKM_M256_MAX_ROWS = 256 MKM_M512_MAX_ROWS = 512 HOP_MAX_ROWS = 15 +# With hop_mid (numeric): the one-shot hop up to four requests in place of BT routes. +HOP_MID_MAX_ROWS = 28 GDN_PUSH_MAX_ROWS = 8 # the rows slot route_shared's shared stage fits shared memory for ROUTE_SHARED_ROWS = 8 @@ -183,6 +186,7 @@ class Facts(msgspec.Struct, frozen=True): gdn_core: bool = False # with the one-launch verify core gdn_conv: bool = False # with the hoisted conv writing each window row once hop: bool = False # the decode hop, built for the widest step it covers + hop_mid: bool = False # with its one-shot protocol up to HOP_MID_MAX_ROWS gdn_push: bool = False # out_proj packed for the push GEMM on every GDN layer gdn_front: bool = False # in_proj_qkvz packed for the front on every GDN layer gdn_trt_front: bool = False # TensorRT-LLM's front built for every GDN layer @@ -208,7 +212,7 @@ def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> route = facts.route and rows <= ROUTE_MAX_ROWS route_topk = facts.route_topk and not route and rows <= ROUTE_TOPK_MAX_ROWS tep8 = facts.tep8_moe - hop = facts.hop and rows <= HOP_MAX_ROWS + hop = facts.hop and rows <= (HOP_MID_MAX_ROWS if facts.hop_mid else HOP_MAX_ROWS) mka = facts.mka and rows <= (MKA_WIDE_MAX_ROWS if facts.mka_wide else MKA_MAX_ROWS) mkm = ( route @@ -300,7 +304,10 @@ def slots(rows: int, facts: Facts, tokens_per_request: int = GDN_FRONT_ROWS) -> and facts.lowm_push and hop and rows <= max(LOWM_OUTPROJ_MAX_ROWS, LOWM_OPROJ_MAX_ROWS), - hop_fp8=hop and facts.hop_fp8 and rows <= HOP_MAX_ROWS, + # Wherever the one-shot hop runs, including hop_mid's four-request width. + hop_fp8=hop + and facts.hop_fp8 + and rows <= (HOP_MID_MAX_ROWS if facts.hop_mid else HOP_MAX_ROWS), ) diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index c5f205d1..83ec53af 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -40,6 +40,7 @@ "mkm", "mkm_wide", "gdn_core", + "hop_mid", "gdn_push", "gdn_front", "gdn_trt_front", @@ -216,6 +217,8 @@ def slot_reasons() -> list[str]: ) if "route_gemm" in SLOTS and "route_topk" not in SLOTS: reasons.append(f"{SLOTS_ENV}: slot 'route_gemm' runs inside slot 'route_topk'") + if "hop_mid" in SLOTS and "hop" not in SLOTS: + reasons.append(f"{SLOTS_ENV}: slot 'hop_mid' widens slot 'hop'") if "gdn_push" in SLOTS and not {"gdn_verify", "hop"} <= SLOTS: reasons.append( f"{SLOTS_ENV}: slot 'gdn_push' pushes slot 'gdn_verify''s out_proj into " @@ -387,6 +390,7 @@ def build(model, ctx): gdn_core=bound["gdn"].gdn_core, gdn_conv=bound["gdn"].gdn_conv, hop=bound["boundary"].hop is not None, + hop_mid=bound["boundary"].mid, gdn_push=bool(bound["gdn"].packed), gdn_front=bool(bound["gdn"].front), gdn_trt_front=bound["gdn"].gdn_trt_front, @@ -456,6 +460,9 @@ def build(model, ctx): <= slot_plan.MKM_M256_MAX_ROWS ): vacuous.add("mkm_m512") + # hop_mid has nothing to serve when no covered step is past the shipped bound. + if max(pctx.widths) <= slot_plan.HOP_MAX_ROWS: + vacuous.add("hop_mid") # route_topk has nothing to serve where slot route covers every width if "route" in SLOTS and max(pctx.widths) <= slot_plan.ROUTE_MAX_ROWS: vacuous.update({"route_topk", "route_gemm"}) diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 349f9c2e..fcdbab79 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -180,6 +180,7 @@ def test_an_sglang_without_the_eagle_forward_view_is_refused(monkeypatch): "slot 'route_lin' needs 'lowm_se' to replace 'route_shared'", ), ({"route_shared"}, "slot 'route_shared' runs in slot 'route''s launch"), + ({"hop_mid"}, "slot 'hop_mid' widens slot 'hop'"), ({"route", "mkm_wide"}, "slot 'mkm_wide' is slot 'mkm''s 16-token engine"), ( {"route", "mkm", "mkm_wide"}, @@ -331,8 +332,10 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): three = plan.slots(21, every) assert three.lowm_se assert three.gdn_front and three.gdn_trt_front + assert three.hop and three.hop_fp8 and three.mka_push assert not (three.lowm_in or three.lowm_outproj or three.lowm_oproj) - # The 28/56-row low-M paths stay stock; GDN core and MKA serve through 224. + # The 28/56-row low-M projection paths stay stock; hop_mid's one-shot path and + # FP8 copy also serve 28 rows, while GDN core and MKA serve through 224. for rows in (28, 56, 112, 224): expected = [ "gdn_conv", @@ -354,6 +357,8 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): expected.append("mkm_m256") if rows <= plan.ROUTE_GEMM_MAX_ROWS: expected.append("route_gemm") + if rows <= plan.HOP_MID_MAX_ROWS: + expected.extend(("hop", "hop_fp8", "mka_push")) assert _on(plan.slots(rows, every)) == sorted(expected) for rows in (448, 896): expected = [ @@ -369,6 +374,13 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): assert _on(plan.slots(28, plan.Facts(route_topk=True))) == ["route_topk"] assert _on(plan.slots(7, plan.Facts(route_topk=True))) == ["route_topk"] assert not plan.slots(28, plan.Facts(route_gemm=True)).route_gemm + mid = plan.Facts(hop=True, hop_mid=True) + assert _on(plan.slots(28, mid)) == ["hop"] + assert _on(plan.slots(56, mid)) == [] + assert _on(plan.slots(28, plan.Facts(hop=True, hop_mid=True, hop_fp8=True))) == [ + "hop", + "hop_fp8", + ] # TEP8's routed-body plan is internal: S1/S2 use route plus a routing-only # top-k launch, M uses the low-M shared expert, and L uses route-big. It # supersedes the all-expert MKM engines and the numeric tiled router front. From 8498568b2ac998930de4d6353852871fcf076200 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 27 Oct 2026 07:05:16 -0700 Subject: [PATCH 25/62] perf: consume the CP-lite gate epilogue --- sgl_nvidia_mega/qwen3_5/attention/mka.py | 11 ++- sgl_nvidia_mega/qwen3_5/attention/split_kv.py | 1 + test/unit/test_qwen3_5_steps.py | 89 +++++++++++++++++++ 3 files changed, 100 insertions(+), 1 deletion(-) diff --git a/sgl_nvidia_mega/qwen3_5/attention/mka.py b/sgl_nvidia_mega/qwen3_5/attention/mka.py index 08e3c353..e02ccacb 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/mka.py +++ b/sgl_nvidia_mega/qwen3_5/attention/mka.py @@ -67,6 +67,7 @@ "srt/runtime_context.py": ("get_parallel",), "srt/layers/layer_boundary/fusions/norm.py": ("fused_norm_gamma",), "srt/model_executor/forward_batch_info.py": ("out_cache_loc",), + "srt/layers/attention/q35_cp_lite.py": ("CpLite",), "srt/kernel_plugins/nvidia/eagle_binding.py|srt/megakernels/eagle_binding.py": ( "kv", "layer_ids", @@ -304,6 +305,7 @@ def _front_step(layer, x, st): from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) + from sglang.srt.layers.attention.q35_cp_lite import CpLite from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.quantization.fp8_utils import flashinfer_bmm_fp8 from sglang.srt.runtime_context import get_parallel @@ -362,8 +364,15 @@ def _front_step(layer, x, st): k_scale=k_scale, v_scale=v_scale, ) + gate = gate.reshape(m, width) + cp_lite = CpLite._instance + oq = None + if cp_lite is not None: + oq = torch.empty((m, width), dtype=torch.float8_e4m3fn, device=x.device) + cp_lite.offer_gate_epilogue(gate, layer.o_proj.input_scale, oq) o = attn(q.reshape(m, width), None, None, st.fb).reshape(m, width) - oq = gate_quant(o.contiguous(), gate.reshape(m, width), layer.o_proj.input_scale) + if cp_lite is None or not cp_lite.take_gate_epilogue(): + oq = gate_quant(o.contiguous(), gate, layer.o_proj.input_scale) if st.slots.lowm_oproj: # slot lowm_out: o_proj on the low-M launch, pushing into the hop with lowm_push hop = None diff --git a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py index f3946d2c..649b5b26 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py +++ b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py @@ -86,6 +86,7 @@ "8e36076dafabff66", "1445b85387928caf", "d7e7c1fe125739e6", + "066a1e0aa5bd6e5c", ), "srt/layers/attention/trtllm_mha_backend.py:TRTLLMHAAttnBackend._should_use_fused_fp8_path": "e0b14791d523a019", "srt/layers/attention/trtllm_mha_backend.py:TRTLLMHAAttnBackend._fused_fp8_qkv_kv_cache": ( diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index dfb14a44..7e09a281 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -2253,6 +2253,95 @@ def symmetric(group, disabled=False): assert out.shape == (7, 4096) and out.dtype == torch.bfloat16 +@pytest.mark.parametrize("consumed", [False, True]) +def test_the_mka_wide_path_offers_gate_quantization_to_cp_lite(consumed): + from sglang.srt.layers.attention.q35_cp_lite import CpLite + from sglang.srt.layers.quantization import fp8_utils + + from sgl_nvidia_mega.qwen3_5 import lowm + from sgl_nvidia_mega.qwen3_5.attention import mka + from sgl_nvidia_mega.qwen3_5.sm103.attention import gate_quant as gate_module + from sgl_nvidia_mega.qwen3_5.sm103.draft import front + + rows = 21 + log = _log() + attn, st = _skv2_setup(log) + layer = _mka_layer(log, attn) + x = torch.zeros(rows, 4096, dtype=torch.bfloat16) + qkv = torch.zeros(rows, 2560, dtype=torch.bfloat16) + q = torch.zeros(rows, 1024, dtype=torch.float8_e4m3fn) + gate = torch.zeros(rows, 1024, dtype=torch.bfloat16) + attention_out = torch.zeros(rows, 1024, dtype=torch.bfloat16) + fallback = torch.zeros(rows, 1024, dtype=torch.float8_e4m3fn) + result = torch.zeros(rows, 4096, dtype=torch.bfloat16) + qkv_linear = layer.qkv_proj + layer.qkv_proj = _Callable(lambda _: (qkv, None), **vars(qkv_linear)) + attn._fn = lambda *args, **kwargs: attention_out + st.fb.out_cache_loc = torch.zeros(rows, dtype=torch.int64) + st.positions = torch.zeros(3, rows, dtype=torch.int64) + st.slots = plan.Slots(mka=True, lowm_oproj=True) + st.scratch = {} + calls = [] + + class GateEpilogue: + def offer_gate_epilogue(self, offered_gate, scale, out): + calls.append(("offer", offered_gate, scale, out)) + + def take_gate_epilogue(self): + calls.append(("take",)) + return consumed + + def quantize(o, offered_gate, scale): + calls.append(("gate_quant", o, offered_gate, scale)) + return fallback + + def row_parallel(linear, o_proj_input, dtype, what, hop): + calls.append(("o_proj", linear, o_proj_input, dtype, what, hop)) + return result + + with ( + patch.object(CpLite, "_instance", GateEpilogue()), + patch.object( + front, + "fused_qk_gemma_rmsnorm_rope_gate", + return_value=(q, None, gate), + ), + patch.object(gate_module, "gate_quant", quantize), + patch.object(fp8_utils, "flashinfer_bmm_fp8", create=True), + patch.object(lowm, "take_fp8", return_value=None), + patch.object(lowm, "row_parallel", row_parallel), + ): + out = mka.step(layer, x, st) + + assert out is result + assert calls[0][0] == "offer" + offered_gate = calls[0][1] + assert offered_gate.shape == gate.shape + assert offered_gate.data_ptr() == gate.data_ptr() + assert calls[0][2] is layer.o_proj.input_scale + offered = calls[0][3] + assert tuple(offered.shape) == (rows, 1024) + assert offered.dtype == torch.float8_e4m3fn + assert calls[1] == ("take",) + if consumed: + assert [call[0] for call in calls] == ["offer", "take", "o_proj"] + assert calls[2][2] is offered + else: + assert [call[0] for call in calls] == [ + "offer", + "take", + "gate_quant", + "o_proj", + ] + quant_input = calls[2][1] + assert quant_input.shape == attention_out.shape + assert quant_input.data_ptr() == attention_out.data_ptr() + quant_gate = calls[2][2] + assert quant_gate.shape == gate.shape + assert quant_gate.data_ptr() == gate.data_ptr() + assert calls[3][2] is fallback + + @pytest.mark.parametrize("consumes", [False, True]) def test_the_mka_push_slot_pushes_where_the_moe_input_only_consumes(consumes): from sgl_nvidia_mega.qwen3_5.attention import mka From 3fd83d3fbc336985cf4381acf8950105df6bb105 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 27 Oct 2026 08:12:41 -0700 Subject: [PATCH 26/62] test: size replay state pools for wide batches --- test/kernels/test_qwen3_5_gdn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/kernels/test_qwen3_5_gdn.py b/test/kernels/test_qwen3_5_gdn.py index 6fd12b51..bd39fece 100644 --- a/test/kernels/test_qwen3_5_gdn.py +++ b/test/kernels/test_qwen3_5_gdn.py @@ -236,7 +236,7 @@ def test_replayssm_verify_is_sglangs(pdl, gate, bs): g = torch.Generator(device="cuda").manual_seed( 3 + int(pdl) + 2 * int(gate) + 10 * bs ) - ring_len, slots = 16, 9 + ring_len, slots = 16, SLOTS tokens = bs * DRAFT def rand(*shape, dtype=torch.bfloat16): From e7fcac2e68d54565bd28040bdaa68e0c89422419 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 27 Oct 2026 08:28:40 -0700 Subject: [PATCH 27/62] test: check the baked router schedule numerically --- test/kernels/test_qwen3_5_moe.py | 35 ++++++++++++++++++-------------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/test/kernels/test_qwen3_5_moe.py b/test/kernels/test_qwen3_5_moe.py index 46fef75e..6507e939 100644 --- a/test/kernels/test_qwen3_5_moe.py +++ b/test/kernels/test_qwen3_5_moe.py @@ -3,8 +3,9 @@ """The MoE package's kernels on one SM100-family GPU (GB300), bit for bit, eager and replayed from a CUDA graph: - the router front (slot route): router_gemm_quant's logits against the cuBLAS router - GEMM SGLang runs (ReplicatedLinear, bf16), and its NVFP4 input against the trtllm - runner's quantizer (flashinfer.fp4_quantize, linear scale layout), at 1-16 rows; + GEMM SGLang runs (ReplicatedLinear, bf16), within eight BF16 representable steps, + and its NVFP4 input against the trtllm runner's quantizer + (flashinfer.fp4_quantize, linear scale layout), at 1-16 rows; - the fused routed experts (slot mkm), against FlashInfer's trtllm-gen NVFP4 MoE as SGLang serves Qwen3.5 at TP8 and TP4 (intermediate 128 / 256 a rank; routing from the logits, RenormalizeNaive, deferred finalize): each token's 10 expert rows and their @@ -15,10 +16,11 @@ rows (the wide engines, on route_topk's routing), also with every token routed to the same experts; graph replay of each; - the router front with the shared gate/up stage (slot route_shared, numeric): its - logits and NVFP4 input bit for bit the plain launch's, and its gate/up rows against - fp64 on SGLang's static FP8 quantization of the rows, at most 1.05 x the mean relative - L2 and 2 x the max abs error of FlashInfer's FP8 BMM, at 1-8 rows and the TP8 / TP16 - widths (256 / 128 rows); graph replay; + logits against cuBLAS within eight BF16 representable steps, its NVFP4 input bit for + bit the plain launch's, and its gate/up rows against fp64 on SGLang's static FP8 + quantization of the rows, at most 1.05 x the mean relative L2 and 2 x the max abs + error of FlashInfer's FP8 BMM, at 1-8 rows and the TP8 / TP16 widths (256 / 128 rows); + graph replay; - the TEP8 routed path: the routing tables route_topk and route-big write for a rank's window of local experts, and the trtllm-gen MoE body on them against FlashInfer's op as SGLang calls it (on @@ -34,6 +36,7 @@ E, H, TOP_K = mkm.EXPERTS, mkm.HIDDEN, mkm.TOP_K X_SCALE = 448.0 * 6.0 / 12.0 # the static input scale of the routed experts H_SCALE = 30.0 / (448.0 * 6.0) # the static FC2 input scale +ROUTER_BF16_ULPS = 8 # R11's validated four-chain reduction versus cuBLAS's order def _gpu_ready() -> bool: @@ -348,7 +351,7 @@ def test_router_gemm_quant_is_the_stock_router_gemm_and_quantizer(m): x, w, scale = _router_inputs(m, seed=m) got = router.router_gemm_quant(x, w, scale) want = _router_stock(x, w, scale) - assert _bf16_ulps(got[0], want[0]) <= 1 + assert _bf16_ulps(got[0], want[0]) <= ROUTER_BF16_ULPS for name, a, b in zip(("x_fp4", "x_sf"), got[1:], want[1:]): assert torch.equal(a.view(torch.uint8), b.reshape(a.shape).view(torch.uint8)), ( name @@ -392,10 +395,12 @@ def test_router_gemm_quant_shared_holds_its_gate_up_to_fp64(m): *routed, gate_up = router.router_gemm_quant_shared( x, w, scale, sh_w, sh_in, sh_ws ) - for name, a, b in zip( - ("logits", "x_fp4", "x_sf"), routed, router.router_gemm_quant(x, w, scale) - ): - assert torch.equal(a, b), name + want = _router_stock(x, w, scale) + assert _bf16_ulps(routed[0], want[0]) <= ROUTER_BF16_ULPS + for name, a, b in zip(("x_fp4", "x_sf"), routed[1:], want[1:]): + assert torch.equal( + a.view(torch.uint8), b.reshape(a.shape).view(torch.uint8) + ), name xq, _ = static_quant_fp8(x, sh_in, repeat_scale=False) ref = (xq.double() @ sh_w.double().t()) * (sh_in.double() * sh_ws.double()) stock = flashinfer_bmm_fp8(xq, sh_w.t(), sh_in, sh_ws, torch.bfloat16) @@ -886,7 +891,8 @@ def step(): def test_router_front_tiles_are_its_16_row_launches(m): """Slot route_gemm (numeric past 16 rows): the router front over 16-row token tiles in one launch equals the 16-row launch on each tile, bit for bit; its NVFP4 input is the - runner's quantizer's; its logits stay within a bf16 ulp of cuBLAS's GEMM; and + runner's quantizer's; its logits stay within eight BF16 representable steps of + cuBLAS's GEMM; and route_topk's routing-only launch routes as its quantizing launch.""" from flashinfer import fp4_quantize @@ -900,9 +906,8 @@ def test_router_front_tiles_are_its_16_row_launches(m): q, q_sf = fp4_quantize(x, scale, 16, False, False) assert torch.equal(x_fp4, q.view(torch.uint8).reshape(m, -1)) assert torch.equal(x_sf, q_sf.view(torch.uint8).reshape(m, -1)) - stock = torch.nn.functional.linear(x, w).float() - ulp = stock.abs().clamp_min(1e-30) * 2.0**-7 - assert ((logits.float() - stock).abs() <= ulp).all() + stock = torch.nn.functional.linear(x, w) + assert _bf16_ulps(logits, stock) <= ROUTER_BF16_ULPS routed = router.route_topk_quant(logits, x, scale) only = router.route_topk_quant(logits, x, scale, quantize=False) assert only[2] is None and only[3] is None From d1a5c38792e20f05b37d4f479eaff589d79c0214 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 27 Oct 2026 16:47:23 -0700 Subject: [PATCH 28/62] fix: fall back from unavailable draft MoE --- sgl_nvidia_mega/qwen3_5/draft/parts.py | 10 ++----- test/unit/test_qwen3_5_runtime.py | 41 ++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 7 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/draft/parts.py b/sgl_nvidia_mega/qwen3_5/draft/parts.py index 6ed2f353..9effb0b5 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/parts.py +++ b/sgl_nvidia_mega/qwen3_5/draft/parts.py @@ -92,6 +92,8 @@ def build(cls, draft) -> Parts: slots = parts.slots if "draft_moe" in slots: parts.moe = parts._build_moe() + if parts.moe is None: + slots = parts.slots = slots - {"draft_moe"} if "draft_hop" in slots: parts.hop = parts._build_hop() if "draft_xchg" in slots: @@ -127,13 +129,7 @@ def build(cls, draft) -> Parts: def _build_moe(self): from sgl_nvidia_mega.qwen3_5.draft.moe import DraftMoe - moe = DraftMoe.build(self.layer.mlp, self.device.index) - if moe is None: - raise RuntimeError( - "NVIDIA megakernel: slot draft_moe cannot run: " - f"{getattr(DraftMoe, 'last_reason', 'see the log')}" - ) - return moe + return DraftMoe.build(self.layer.mlp, self.device.index) def _build_hop(self): from sgl_nvidia_mega.qwen3_5.draft.hop import GEMV_O_PROJ, DraftHop diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index f24fb3fe..0f800724 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -295,6 +295,47 @@ def test_an_explicit_policy_controls_the_draft_entries(monkeypatch): assert runtime.build(model, target_ctx).summary() == "slots none (stock)" +def test_unavailable_draft_moe_is_dropped_and_runs_stock_mlp(monkeypatch): + from sgl_nvidia_mega.qwen3_5.draft import attention, layer + from sgl_nvidia_mega.qwen3_5.draft.moe import DraftMoe + from sgl_nvidia_mega.qwen3_5.draft.parts import Parts + + calls = [] + + def mlp(hidden): + calls.append("stock mlp") + return hidden + + decoder = SimpleNamespace( + input_layernorm=lambda hidden: hidden, + post_attention_layernorm=lambda hidden, residual: (hidden, residual), + mlp=mlp, + ) + norm = lambda hidden, residual: (hidden, residual) + model = SimpleNamespace(layers=[decoder], start_layer=0, end_layer=1, norm=norm) + draft = SimpleNamespace( + lm=SimpleNamespace(model=model), + binding="binding", + slots=frozenset({"draft_moe"}), + ) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr( + attention, "Attention", lambda _: lambda hidden, *a, **k: hidden + ) + monkeypatch.setattr(DraftMoe, "build", classmethod(lambda cls, mlp, device: None)) + monkeypatch.setattr( + "sglang.srt.distributed.tensor_model_parallel_all_reduce", lambda hidden: hidden + ) + + parts = Parts.build(draft) + assert parts.moe is None and "draft_moe" not in parts.slots + x = torch.ones(1, 4) + assert torch.equal( + layer.layer(parts, x, None, None, None, None, rows_per_request=1), x + ) + assert calls == ["stock mlp"] + + @pytest.mark.parametrize( "batch_size, lower_bound, seq_lens, want", [ From 6324c500f4b0e8aab70625d83084a61b39e33213 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 27 Oct 2026 21:56:45 -0700 Subject: [PATCH 29/62] feat: move Qwen3.5 CP attention into plugin slots --- sgl_nvidia_mega/qwen3_5/attention/__init__.py | 4 +- sgl_nvidia_mega/qwen3_5/attention/mka.py | 14 +- sgl_nvidia_mega/qwen3_5/attention/split_kv.py | 54 +- sgl_nvidia_mega/qwen3_5/draft/__init__.py | 1 + sgl_nvidia_mega/qwen3_5/draft/attention.py | 4 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 48 + .../qwen3_5/sm103/attention/q35_cp_fused.py | 819 ++++ .../qwen3_5/sm103/attention/q35_cp_lite.py | 2364 ++++++++++ .../sm103/attention/q35_cp_partial_attn.py | 3893 +++++++++++++++++ sgl_nvidia_mega/qwen3_5/spec.py | 2 +- test/unit/test_qwen3_5_runtime.py | 114 + test/unit/test_qwen3_5_steps.py | 113 +- 12 files changed, 7389 insertions(+), 41 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_partial_attn.py diff --git a/sgl_nvidia_mega/qwen3_5/attention/__init__.py b/sgl_nvidia_mega/qwen3_5/attention/__init__.py index e4254cb0..3c3a7c68 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/attention/__init__.py @@ -137,9 +137,9 @@ def step(layer, x, st): q, k, v, gate = layer._prepare_qkv_gate( positions=st.positions, hidden_states=x, forward_batch=st.fb ) - if st.slots.skv2: - from sgl_nvidia_mega.qwen3_5.attention import split_kv + from sgl_nvidia_mega.qwen3_5.attention import split_kv + if st.slots.skv2 or split_kv.cp_multi(q.shape[0]): o = split_kv.attend(layer.attn, q, k, v, st) else: o = layer.attn(q, k, v, st.fb) diff --git a/sgl_nvidia_mega/qwen3_5/attention/mka.py b/sgl_nvidia_mega/qwen3_5/attention/mka.py index e02ccacb..4b943e90 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/mka.py +++ b/sgl_nvidia_mega/qwen3_5/attention/mka.py @@ -67,7 +67,6 @@ "srt/runtime_context.py": ("get_parallel",), "srt/layers/layer_boundary/fusions/norm.py": ("fused_norm_gamma",), "srt/model_executor/forward_batch_info.py": ("out_cache_loc",), - "srt/layers/attention/q35_cp_lite.py": ("CpLite",), "srt/kernel_plugins/nvidia/eagle_binding.py|srt/megakernels/eagle_binding.py": ( "kv", "layer_ids", @@ -241,7 +240,7 @@ def step(layer, x, st): nq=layer.num_heads, eps=layer.q_norm.variance_epsilon, ) - if st.slots.skv2: + if st.slots.skv2 or split_kv.cp_multi(m): o = split_kv.core(attn, q, st) else: o = attn(q, None, None, st.fb) @@ -305,7 +304,6 @@ def _front_step(layer, x, st): from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) - from sglang.srt.layers.attention.q35_cp_lite import CpLite from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.quantization.fp8_utils import flashinfer_bmm_fp8 from sglang.srt.runtime_context import get_parallel @@ -313,6 +311,7 @@ def _front_step(layer, x, st): from sgl_nvidia_mega.qwen3_5 import lowm from sgl_nvidia_mega.qwen3_5.attention import split_kv from sgl_nvidia_mega.qwen3_5.sm103.attention.gate_quant import gate_quant + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite from sgl_nvidia_mega.qwen3_5.sm103.draft.front import ( fused_qk_gemma_rmsnorm_rope_gate, ) @@ -365,12 +364,17 @@ def _front_step(layer, x, st): v_scale=v_scale, ) gate = gate.reshape(m, width) - cp_lite = CpLite._instance + cp_lite = CpLite._instance if split_kv.cp_multi(m) else None oq = None if cp_lite is not None: oq = torch.empty((m, width), dtype=torch.float8_e4m3fn, device=x.device) cp_lite.offer_gate_epilogue(gate, layer.o_proj.input_scale, oq) - o = attn(q.reshape(m, width), None, None, st.fb).reshape(m, width) + q = q.reshape(m, width) + o = ( + split_kv.core(attn, q, st) + if cp_lite is not None + else attn(q, None, None, st.fb) + ).reshape(m, width) if cp_lite is None or not cp_lite.take_gate_epilogue(): oq = gate_quant(o.contiguous(), gate, layer.o_proj.input_scale) if st.slots.lowm_oproj: diff --git a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py index 649b5b26..5aaf9df4 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py +++ b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py @@ -22,8 +22,6 @@ from __future__ import annotations -import importlib -import inspect import logging import torch @@ -182,6 +180,27 @@ def core(attn, q, st): step are written: [rows, heads * head_dim] in bf16.""" from sgl_nvidia_mega.qwen3_5.attention import route + if cp_multi(q.shape[0]): + o = cp_core( + attn, + q, + st.view.kv, + st.view.attn, + st.fb, + ROWS, + min_bs=None, + ) + if o is not None: + key = (route.SPLIT_KV_CP, q.shape[0]) + if key not in _ENGAGED: + _ENGAGED.add(key) + logger.info( + "NVIDIA megakernel: Qwen3.5 slot skv2_cp keeps CP-lite for " + "%d requests x %d target-verify rows", + q.shape[0] // ROWS, + ROWS, + ) + return o label = route.label(route.verify_variants(), st.fb) if (label, q.shape[0]) not in _ENGAGED: _ENGAGED.add((label, q.shape[0])) @@ -210,29 +229,22 @@ def core(attn, q, st): def cp_lite(): - """SGLang's process CP-lite (q35_cp_lite.CpLite, built by its TRT-LLM MHA backend - when a CP graph variant is requested) when it can serve one request - (``attend(min_bs=)``), else None.""" - try: - module = importlib.import_module("sglang.srt.layers.attention.q35_cp_lite") - except ImportError: - return None - if "min_bs" not in inspect.signature(module.CpLite.attend).parameters: - return None - return module.CpLite._instance + """The CP state built collectively by the Qwen3.5 plugin, else None.""" + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + return CpLite._instance + + +def cp_multi(rows: int) -> bool: + """Whether a multi-request target verify can try the plugin's CP state.""" + from sgl_nvidia_mega.qwen3_5.runtime import SLOTS + + return "skv2_cp" in SLOTS and rows > ROWS and cp_lite() is not None def cp_reason() -> str: """Why the CP graph variants cannot run ("" when they can).""" - try: - module = importlib.import_module("sglang.srt.layers.attention.q35_cp_lite") - except ImportError: - return "this SGLang has no CP-lite (q35_cp_lite)" - if "min_bs" not in inspect.signature(module.CpLite.attend).parameters: - return "this SGLang's CP-lite serves no single request (attend has no min_bs)" - if module.CpLite._instance is None: - return "SGLang's attention backend did not build CP-lite" - return "" + return "" if cp_lite() is not None else "the plugin did not build CP-lite" def cp_core(attn, q, kv, meta, forward_batch, q_len: int, min_bs: int | None = 1): diff --git a/sgl_nvidia_mega/qwen3_5/draft/__init__.py b/sgl_nvidia_mega/qwen3_5/draft/__init__.py index 712cf21e..46c38fe8 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/draft/__init__.py @@ -214,6 +214,7 @@ def _build(model, hf_config, ctx) -> MtpDraft: from sgl_nvidia_mega.qwen3_5 import runtime runtime.apply_expert_parallel(get_parallel().moe_ep_size) + runtime.build_cp_lite(ctx.binding) return MtpDraft( model, binding=ctx.binding, diff --git a/sgl_nvidia_mega/qwen3_5/draft/attention.py b/sgl_nvidia_mega/qwen3_5/draft/attention.py index 02d2cef6..de3b322a 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/attention.py +++ b/sgl_nvidia_mega/qwen3_5/draft/attention.py @@ -418,8 +418,8 @@ def run(self, attn, q8, kv, meta): class DraftCP: """Slots draft_q1_cp / draft_skv2_cp: in the draft loop's ``Q1_CP`` and the draft extend's ``SPLIT_KV_CP`` graph variants (one request past its measured crossover), - the decode of the request's rows as CP-lite, which then serves one request - (SGLang's q35_cp_lite through ``attention.split_kv.cp_core``'s call). Slot + the decode of the request's rows as the plugin's CP-lite, which then serves one + request through ``attention.split_kv.cp_core``. Slot draft_multi also preserves CP-lite on multi-request graphs where SGLang built it; CP-lite's own shape gate decides whether it runs.""" diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 83ec53af..35812268 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -85,6 +85,7 @@ "draft_q1_cp", "draft_skv2_cp", ) +CP_LITE_SLOTS = frozenset({"skv2_cp", "draft_q1_cp", "draft_skv2_cp"}) # SGLang consumes these names from the shared slot policy itself. They are valid # requests, but are intentionally not built or dispatched by this plugin. RUNTIME_ONLY_SLOTS: tuple[str, ...] = ("gdn_commit", "sharded_gumbel") @@ -118,6 +119,44 @@ def _requested_slots() -> frozenset: SLOTS = _requested_slots() +def build_cp_lite(binding) -> None: + """Collectively build the CP slots' process-local state before validation. + + The allocation and rendezvous order must be identical on every attention-TP + rank, so this runs once from each Qwen3.5 plugin build, before any package checks + ``CpLite._instance`` or CUDA graph variant is created. + """ + from sglang.srt.kernel_plugins.nvidia.binding import megakernel_agree + + requested = tuple(sorted(SLOTS & CP_LITE_SLOTS)) + megakernel_agree( + requested, + what="Qwen3.5 CP slots requested", + ) + if not requested: + return + megakernel_agree( + (binding.page_size, binding.tp_size), + what="Qwen3.5 CP page size and TP size", + ) + if binding.page_size != 64: + raise RuntimeError( + "NVIDIA megakernel: Qwen3.5 CP slots need page size 64 " + f"(got {binding.page_size})" + ) + + import torch + + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + cp = CpLite.maybe_create(binding, torch.device("cuda", torch.cuda.current_device())) + if cp is None: + raise RuntimeError( + "NVIDIA megakernel: Qwen3.5 CP slots need a supported attention-TP " + "Q/KV-head topology" + ) + + def apply_expert_parallel(ep_size: int) -> frozenset: """Drop ``EP_STOCK_SLOTS`` from ``SLOTS`` when the MoE runs expert-parallel (``ep_size`` > 1); returns the dropped names. Every process calls it before reading @@ -320,6 +359,13 @@ def slot_reasons() -> list[str]: "binding", ), "srt/runtime_context.py": ("get_parallel", '"tp_size"', '"moe_ep_size"'), + "srt/kernel_plugins/nvidia/eagle_binding.py|srt/megakernels/eagle_binding.py": ( + "page_size", + "tp_rank", + "tp_size", + "tp_device_group", + ), + "srt/kernel_plugins/nvidia/binding.py": ("megakernel_agree",), } @@ -342,6 +388,8 @@ def build(model, ctx): from sgl_nvidia_mega.qwen3_5.runtime import parity from sgl_nvidia_mega.qwen3_5.runtime.executor import Qwen35Verify + build_cp_lite(ctx.binding) + dropped = apply_expert_parallel(get_parallel().moe_ep_size) if dropped: logger.info( diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py new file mode 100644 index 00000000..d9daa3d1 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py @@ -0,0 +1,819 @@ +# 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. +# Adapted from TensorRT-LLM 4d09932ca7 + 7df45abe5c's plans above 32 requests (Apache-2.0): +# tensorrt_llm/_torch/attention_backend/fmha/q35_cp_fused.py; switches are constants here, the debug ring is dropped, +# and the kernel import points at this package. +"""Host side of the Qwen3.5 fused CP-lite partial decode attention. + +One launch replaces the preparation, FMHA, LSE, and partial-output launches. Each CP +rank attends the pages it owns and writes the all-to-all merge's input layout. + +The module owns the grid-barrier counters and partial workspace. Call +``q35_cp_fused_prepare`` before CUDA graph capture. +""" + +import functools +import logging + +import torch + +logger = logging.getLogger(__name__) + +PAGE = 64 +HEAD_DIM = 256 +MAX_ROWS = 128 +# Batch bound of the plan: one request per lane of the plan warp up to 32; above that the wide-plan kernel variant +# (two requests per lane, q35_cp_partial_attn.Q35CpPartialAttnKernel._plan_wide) up to 64, and the quad-plan variant +# (four per lane, _plan_quad) up to 128. The equal split stays <= 32. +_QUAD_PLAN_MAX_BATCH = 128 +# Above 64 requests the calls take the two-slot kernel (kernel_two / _plan_two, four requests per plan-warp lane): a +# CTA computes up to two requests, so no request is held to whole CTAs. The plans above give every request at least one +# of the SMs' CTAs; past ~SMs / 2 uniform long requests then get exactly one each and the slowest CTA runs a whole +# request (64k contexts, b96: 128 tiles where 81 would balance), which the two-slot plan's per-CTA share removes. +_TWO_SLOT_MAX_BATCH = 128 +_WIDE_PLAN_MAX_BATCH = 64 +_LANE_PLAN_MAX_BATCH = 32 + + +# Verify calls use the wide and two-slot plans through 128 requests; draft calls use +# the wide plan through 64. Batches at or below 32 keep the equal split, so the +# separate balanced-split optimization remains disabled. +MAX_BATCH = 128 +MAX_BATCH_Q1 = 64 +# Pack every two-slot call; the unpacked form remains available as a test hook. +_TWO_SLOT_PACK = True + + +# Mirror of q35_cp_partial_attn.MAX_COMB_ROWS (rows one CTA's combine chunk may touch). +_MAX_COMB_ROWS = 32 +LOG2E = 1.4426950408889634 + +# MTP may write the page immediately before the new tokens earlier in the step. +_FRESH_MARGIN = 1 +# Request the prefix KV below that margin before the dependency wait (overlaps the q all-gather). +_KV_PRE_WAIT = True +# E4M3 P shift (log2 units), as TRTLLM_Q35_ATTN_FUSED_P_SHIFT: P = exp2(x - m + shift) before the FP8 conversion, so +# probabilities down to 2^-(6 + shift) of the row max stay E4M3 normals (0 = the unscaled P that flushes small +# probabilities: 5-13x the error of trtllm-gen with an attention sink). 5: trtllm-gen's error level (within ~5 % of 6) +# at no cost on flat scores and +0.5 us per call with very wide ones (6: +0.9 us). +_P_SHIFT = 5.0 +_DBG_SLOTS = 0 +# SMs left out of the grid (default 0: one CTA per SM). The in-kernel grid barrier needs every CTA resident; a reserve +# keeps the grid co-resident even when other kernels (e.g. a PDL neighbour still polling its peers) hold some SMs. +_SM_RESERVE = 0 +_DBG_RING_CTAS = 256 + +# Declared alignment (bytes) of each pointer argument, used at compile and call time (make_ptr checks the address): +# q and the KV pool are TMA sources and the workspace takes 16-byte stores / bulk copies; tables, seq_lens and the +# scales are read element by element (views such as scales[2:3] are 4-byte aligned); out_o is 16 when its rows allow +# 8-element vector stores, else 4 (alltoall_merge reads it as int32). +_ALIGN = { + "q": 16, + "kv": 16, + "tab": 4, + "seq": 4, + "bmm1": 4, + "bmm2": 4, + "lse": 4, + "ws": 16, + "cnt": 8, + "dbg": 8, +} + + +def cp_partial_geometry(rows: int, n_splits: int) -> dict: + """Mirror of q35_cp_partial_attn.cp_partial_geometry (kept import-free of cutlass; checked at compile time).""" + total = rows * HEAD_DIM + chunk = 8 * ((total + 8 * n_splits - 1) // (8 * n_splits)) + chunk_pad = chunk + (8 if (chunk // 8) % 2 == 0 else 0) + n_chunks = (total + chunk - 1) // chunk + lse_stride = 4 * ((n_splits + 3) // 4) + ws_o_elems = n_chunks * n_splits * chunk_pad + max_cov_rows = (chunk + HEAD_DIM - 2) // HEAD_DIM + 1 + return { + "chunk": chunk, + "ws_req_bytes": ws_o_elems * 2 + rows * lse_stride * 4, + "max_cov_rows": max_cov_rows, + } + + +def cp_partial_balanced_ws_bytes(rows: int, n_ctas: int, slots: int = 1) -> int: + """Mirror of q35_cp_partial_attn.cp_partial_balanced_geometry()["ws_bytes"]: per CTA one O slot (the largest + per-split share n_chunks(c) * chunk_pad(c) of a request's partials over c = 1..n_ctas, FP16) and 16 * rows bytes + of LSE space; ``slots`` slots per CTA (2 for the two-slot plan).""" + total = rows * HEAD_DIM + slot_o = 0 + for c in range(1, n_ctas + 1): + chunk = 8 * ((total + 8 * c - 1) // (8 * c)) + chunk_pad = chunk + (8 if (chunk // 8) % 2 == 0 else 0) + slot_o = max(slot_o, ((total + chunk - 1) // chunk) * chunk_pad * 2) + return slots * n_ctas * (16 * ((slot_o + 15) // 16) + 16 * rows) + + +def balanced_split_plan(seq_lens, cp_rank: int, cp_size: int, n_ctas: int): + """Host mirror of the balanced split's in-kernel plan (q35_cp_partial_attn.Q35CpPartialAttnKernel._plan). + + Returns (T, c, ctas): T = the largest tile count of a CTA's block, c[r] = CTAs of request r, ctas[i] = (request, + index in the request, first tile, tiles) of CTA i (tile k of a request = its owned pages 2k, 2k + 1 on this rank). + """ + t = [ + ((((int(L) + PAGE - 1) // PAGE) - cp_rank + cp_size - 1) // cp_size + 1) // 2 + for L in seq_lens + ] + b = len(t) + assert 1 <= b < n_ctas and b <= _QUAD_PLAN_MAX_BATCH + total, tmax = sum(t), max(t) + lo = max(-(-total // n_ctas), 1) + hi = max(min(tmax, -(-total // (n_ctas - b))), lo) + while lo < hi: + mid = (lo + hi) // 2 + if sum(max(-(-x // mid), 1) for x in t) <= n_ctas: + hi = mid + else: + lo = mid + 1 + c = [max(-(-x // lo), 1) for x in t] + top = max(r for r in range(b) if t[r] == tmax) + c[top] += n_ctas - sum(c) + ctas = [] + for r in range(b): + base, rem = divmod(t[r], c[r]) + for j in range(c[r]): + ctas.append((r, j, j * base + min(j, rem), base + (1 if j < rem else 0))) + return lo, c, ctas + + +def _two_slot_pack_starts(t, T: int): + """Packed two-slot placement for T: request r starts at V_r; at a CTA boundary it takes max(t_r, 1) virtual + tiles, at offset p inside a CTA's range max(t_r, T - p). Returns (starts, end).""" + starts, V = [], 0 + for x in t: + starts.append(V) + p = V % T + V += max(x, 1) if p == 0 else max(x, T - p) + return starts, V + + +def _two_slot_pack_T(t, n_ctas: int, hi: int) -> int: + """T of the packed plan (q35_cp_partial_attn.Q35CpPartialAttnKernel._plan_pack): feasible = end <= n_ctas * T, + not monotone in T; hi (the unpacked plan's T) always is. The smallest feasible T in [lo, lo + 31], lo = ceil(sum + max(t_r, 1) / n_ctas); else the smallest feasible of 32 candidates from lo + 32 to hi in steps of s, then the + smallest feasible of the min(s - 1, 32) values below it (or it).""" + lo = max(-(-sum(max(x, 1) for x in t) // n_ctas), 1) + + def ok(T): + return _two_slot_pack_starts(t, T)[1] <= n_ctas * T + + for T in range(lo, lo + 32): + if ok(T): + return T + lo2 = lo + 32 + s = max(-(-(hi - lo2) // 31), 1) + c1 = next(T for T in (min(lo2 + k * s, hi) for k in range(32)) if ok(T)) + return next((T for T in range(c1 - min(s - 1, 32), c1) if ok(T)), c1) + + +def balanced_split_plan_two( + seq_lens, cp_rank: int, cp_size: int, n_ctas: int, pack: bool = False +): + """Host mirror of the two-slot plan (q35_cp_partial_attn.Q35CpPartialAttnKernel._plan_two). + + T = the smallest tile count with sum_r max(t_r, T) <= n_ctas * T; request r takes the virtual tile range [V_r, V_r + + max(t_r, T)) (requests in order), CTA i the range [i * T, (i + 1) * T). Returns (T, slots): slots[i] = [slot 0, + slot 1] of CTA i, each None or (request, index among the request's CTAs, its CTAs, first workspace slot, first + tile, tiles, CTAs with work); slot 0 is the request holding the CTA's first position, slot 1 the one starting + inside its range. With ``pack``, use the request positions and T from _two_slot_pack_starts / _two_slot_pack_T. + """ + t = [ + ((((int(L) + PAGE - 1) // PAGE) - cp_rank + cp_size - 1) // cp_size + 1) // 2 + for L in seq_lens + ] + b = len(t) + assert 1 <= b < n_ctas and b <= _TWO_SLOT_MAX_BATCH + total, tmax = sum(t), max(t) + lo = max(-(-total // n_ctas), 1) + hi = max(tmax, lo) + while lo < hi: + mid = (lo + hi) // 2 + if sum(max(x, mid) for x in t) <= n_ctas * mid: + hi = mid + else: + lo = mid + 1 + T = lo + if pack: + T = _two_slot_pack_T(t, n_ctas, T) + starts, end = _two_slot_pack_starts(t, T) + else: + starts = [0] + for x in t: + starts.append(starts[-1] + max(x, T)) + end = starts.pop() + slots = [[None, None] for _ in range(n_ctas)] + s0 = 0 + for r, x in enumerate(t): + V = starts[r] + te = max(x, 1) + f, last = V // T, (V + te - 1) // T + c = last - f + 1 + for j in range(c): + i = f + j + lo_t, hi_t = max(V, i * T) - V, min(V + te, (i + 1) * T) - V + slot = ( + r, + j, + c, + s0, + lo_t if x else 0, + hi_t - lo_t if x else 0, + c if x else 0, + ) + k = 0 if V <= i * T else 1 + assert slots[i][k] is None + slots[i][k] = slot + s0 += c + assert end <= n_ctas * T and s0 <= 2 * n_ctas + return T, slots + + +def _num_splits(batch: int, sm_count: int) -> int: + return max((sm_count - _SM_RESERVE) // batch, 1) + + +def _num_ctas(sm_count: int) -> int: + """Grid of the balanced split (any batch).""" + return max(sm_count - _SM_RESERVE, 1) + + +def q35_cp_fused_supported( + q_len: int, + cp_size: int, + hq: int, + batch: int, + sm_count: int, + max_batch: int | None = None, +) -> bool: + """Shapes the kernel handles (D 256, FP8 q / KV and 64-token pages are checked by the caller). max_batch: the + batch bound (default MAX_BATCH, or MAX_BATCH_Q1 at q_len 1; a caller that folds q_len into the heads passes its + own).""" + rows = q_len * cp_size * hq + if max_batch is None: + max_batch = MAX_BATCH if q_len > 1 else MAX_BATCH_Q1 + if not ( + 1 <= rows <= MAX_ROWS + and 1 <= batch <= min(max_batch, MAX_BATCH) + and batch <= sm_count + ): + return False + if batch > _LANE_PLAN_MAX_BATCH: + return batch < _num_ctas( + sm_count + ) # the plan gives every request a CTA; it plans <= 128 requests + return ( + cp_partial_geometry(rows, _num_splits(batch, sm_count))["max_cov_rows"] + <= _MAX_COMB_ROWS + ) + + +class _DeviceState: + """Per-device grid-barrier counters (int64 [MAX_BATCH + 1, MAX_BATCH], row b = batch size b, one epoch per + request; zero-initialized once, never reset) and the partial workspace (sized for every supported call).""" + + def __init__(self, device: torch.device, sm_count: int): + self.sm_count = sm_count + need = 0 + for b in range(1, MAX_BATCH + 1): + if b <= sm_count: + need = max( + need, + b + * cp_partial_geometry(MAX_ROWS, _num_splits(b, sm_count))[ + "ws_req_bytes" + ], + ) + need = max( + need, + cp_partial_balanced_ws_bytes(MAX_ROWS, _num_ctas(sm_count), slots=2), + ) + self.workspace = torch.empty(need // 4 + 64, dtype=torch.float32, device=device) + self.ws_balanced = ( + None # test calls that override the split mode (_balanced_workspace) + ) + self.balanced_logged = ( + False # runtime evidence: one INFO line at the first balanced call + ) + self.wide_logged = ( + False # ... and one at the first call above 32 requests (wide-plan variant) + ) + self.quad_logged = False # ... and one at the first call above 64 requests (two-slot / quad-plan variant) + self.pack_logged = ( + False # ... and one at the first two-slot call on the packed plan + ) + self.counters = torch.zeros( + MAX_BATCH + 1, MAX_BATCH, dtype=torch.int64, device=device + ) + self.dbg_ring = None + if _DBG_SLOTS > 0: + self.dbg_ring = torch.zeros( + _DBG_SLOTS * _DBG_RING_CTAS * 8, dtype=torch.int64, device=device + ) + _start_dbg_dump(self.dbg_ring, device) + + +def _start_dbg_dump(ring: torch.Tensor, device: torch.device) -> None: + raise NotImplementedError("the debug ring is not ported (_DBG_SLOTS = 0)") + + +_STATES: dict[int, _DeviceState] = {} + + +def _sm_count(device: torch.device) -> int: + return torch.cuda.get_device_properties(device).multi_processor_count + + +def q35_cp_fused_prepare(device: torch.device | None = None, sm_count: int = 0) -> None: + """Allocate this device's counters / workspace (call before CUDA graph capture).""" + device = ( + torch.device("cuda", torch.cuda.current_device()) if device is None else device + ) + idx = device.index if device.index is not None else torch.cuda.current_device() + if idx not in _STATES: + _STATES[idx] = _DeviceState( + torch.device("cuda", idx), sm_count or _sm_count(device) + ) + + +def _state(device: torch.device, sm_count: int) -> _DeviceState: + idx = device.index if device.index is not None else torch.cuda.current_device() + st = _STATES.get(idx) + if st is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "q35_cp_fused_partial: call q35_cp_fused_prepare() (or one eager call) before CUDA " + "graph capture: the grid-barrier counters must be zero-initialized persistent memory" + ) + q35_cp_fused_prepare(device, sm_count) + st = _STATES[idx] + return st + + +def _ptr(t: torch.Tensor, dtype, align: int): + from cutlass import cute + + return cute.runtime.make_ptr( + dtype, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=align + ) + + +def _balanced_workspace(state: "_DeviceState", need: int) -> torch.Tensor: + """Workspace for a wide-plan call or a test that overrides the split mode.""" + if state.workspace.numel() * 4 >= need: + return state.workspace + if state.ws_balanced is None or state.ws_balanced.numel() * 4 < need: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "q35_cp_fused_partial(balanced=True): make one eager call before CUDA graph capture" + ) + state.ws_balanced = torch.empty( + need // 4 + 64, dtype=torch.float32, device=state.workspace.device + ) + return state.ws_balanced + + +@functools.cache +def _compiled( + q_len: int, + gh: int, + hql: int, + cp_size: int, + n_splits: int, + bmm1_from_ptr: bool, + bmm2_from_ptr: bool, + o_vec16: bool, + use_pdl: bool, + kv_pre_wait: bool, + fresh_margin: int, + debug_ts: bool, + unsafe_skip_wait: bool, + p_shift: float = 5.0, + debug_ring: int = 0, + balanced: bool = False, + rescale_threshold: float | None = None, + l_from_p8: bool = False, + len_sub: int = 0, + early_trigger: bool = False, + wide_plan: bool = False, + quad_plan: bool = False, + two_slot: bool = False, + two_slot_pack: bool = False, +): + import cuda.bindings.driver as cuda + import cutlass + from cutlass import cute + + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_partial_attn import ( + Q35CpPartialAttnKernel, + ) + + if balanced: + kernel = Q35CpPartialAttnKernel( + q_len=q_len, + gh=gh, + hql=hql, + cp_size=cp_size, + n_splits=n_splits, + bmm1_from_ptr=bmm1_from_ptr, + bmm2_from_ptr=bmm2_from_ptr, + o_vec16=o_vec16, + use_pdl=use_pdl, + kv_pre_wait=kv_pre_wait, + fresh_margin_pages=fresh_margin, + debug_ts=debug_ts, + unsafe_skip_wait=unsafe_skip_wait, + p_shift=p_shift, + debug_ring=debug_ring, + balanced=True, + rescale_threshold=rescale_threshold, + l_from_p8=l_from_p8, + len_sub=len_sub, + early_trigger=early_trigger, + wide_plan=wide_plan, + quad_plan=quad_plan, + two_slot=two_slot, + two_slot_pack=two_slot_pack, + ) + assert kernel.ws_bytes == cp_partial_balanced_ws_bytes( + q_len * gh, n_splits, 2 if two_slot else 1 + ) + else: + kernel = Q35CpPartialAttnKernel( + q_len=q_len, + gh=gh, + hql=hql, + cp_size=cp_size, + n_splits=n_splits, + bmm1_from_ptr=bmm1_from_ptr, + bmm2_from_ptr=bmm2_from_ptr, + o_vec16=o_vec16, + use_pdl=use_pdl, + kv_pre_wait=kv_pre_wait, + fresh_margin_pages=fresh_margin, + debug_ts=debug_ts, + unsafe_skip_wait=unsafe_skip_wait, + p_shift=p_shift, + debug_ring=debug_ring, + rescale_threshold=rescale_threshold, + l_from_p8=l_from_p8, + len_sub=len_sub, + early_trigger=early_trigger, + ) + assert ( + kernel.ws_req_bytes + == cp_partial_geometry(q_len * gh, n_splits)["ws_req_bytes"] + ) + dev = torch.device("cuda") + f32 = torch.empty(1024, dtype=torch.float32, device=dev) + i32 = torch.zeros(64, dtype=torch.int32, device=dev) + u8 = torch.empty(64 * HEAD_DIM, dtype=torch.uint8, device=dev) + bf = torch.empty(1024, dtype=torch.bfloat16, device=dev) + i64 = torch.zeros(16, dtype=torch.int64, device=dev) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + fp8 = cutlass.Float8E4M3FN + a = _ALIGN + I32 = cutlass.Int32 + return cute.compile( + kernel, + _ptr(u8, fp8, a["q"]), + _ptr(u8, fp8, a["kv"]), + _ptr(u8, fp8, a["kv"]), + I32(64), + _ptr(i32, I32, a["tab"]), + I32(2), + I32(1), + _ptr(i32, I32, a["seq"]), + I32(0), + _ptr(f32, cutlass.Float32, a["bmm1"]), + cutlass.Float32(1.0), + _ptr(f32, cutlass.Float32, a["bmm2"]), + cutlass.Float32(1.0), + _ptr(bf, cutlass.BFloat16, 16 if o_vec16 else 4), + I32(HEAD_DIM), + I32(HEAD_DIM), + I32(HEAD_DIM), + _ptr(f32, cutlass.Float32, a["lse"]), + I32(1), + I32(1), + I32(1), + _ptr(f32, cutlass.Float32, a["ws"]), + _ptr(i64, cutlass.Int64, a["cnt"]), + _ptr(i64, cutlass.Int64, a["dbg"]), + I32(1), + stream, + ) + + +def q35_cp_fused_partial( + q_all: torch.Tensor, + kv_pool: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + q_len: int, + cp_rank: int, + cp_size: int, + bmm1_scale: float | torch.Tensor, + bmm2_scale: float | torch.Tensor, + out_o: torch.Tensor, + out_lse: torch.Tensor, + sm_count: int = 0, + enable_pdl: bool = False, + *, + debug: torch.Tensor | None = None, + v_pool: torch.Tensor | None = None, + unsafe_skip_wait: bool = False, + n_splits: int | None = None, + fresh_margin: int | None = None, + counters: torch.Tensor | None = None, + p_shift: float | None = None, + balanced: bool | None = None, + rescale_threshold: float | None = None, + l_from_p8: bool = False, + len_sub: int = 0, + early_trigger: bool = False, + plan: str | None = None, + pack: bool | None = None, +) -> None: + """This CP rank's partial attention of every group head over the pages it owns (q35_cp_fused_partial_ref). + + q_all [cp, rows, hq, D] FP8 (the all-gather output viewed member-major; the native layout, strides (hq * D, + cp * hq * D, D, 1) at a 16-byte aligned address, is read in place, any other view through a copy); kv_pool [blocks, + 1, 64, D] FP8; block_tables [>= b, 2, max_pages] int32 (K / V block ids, unit stride on pages); seq_lens [>= b] + int32 (L incl. the q_len new tokens); bmm1_scale float (natural-log domain) or a [>= 2] FP32 tensor (raw, raw * + log2e); bmm2_scale float or a [>= 1] FP32 tensor. out_o [cp, rows, hq, D] bf16 (unit stride on D), out_lse [cp, + rows, hq] FP32 (base 2). Test hooks: debug (int64 [b * splits, 16] %globaltimer stamps; [CTAs, 16] when balanced), + unsafe_skip_wait (no dependency waits), n_splits (must come with its own ``counters``: every launch on a counter + row must use the same split count; balanced: the grid's CTAs), fresh_margin, balanced (a test override of the + split mode; the wide runtime path selects it automatically and an override needs its own ``counters``), plan + ("lane", "wide", "quad" or "two": the wide split's plan variant instead of the one the batch selects; lane / wide + / quad cover up to 32 / 64 / 128 requests and compute the same one-slot plan, "two" (up to 128) the two-slot plan, + with [CTAs, 24] debug rows), pack (a test override of the packed two-slot plan). + Numerics: p_shift (E4M3 P shift, default _P_SHIFT), rescale_threshold (lazy rescale threshold in log2 units, + default min(8, 8.75 - p_shift); 0 = exact running max), l_from_p8 (row sums from the E4M3-rounded P). + """ + import cuda.bindings.driver as cuda + import cutlass + + cp, rows, hq, D = q_all.shape + assert cp == cp_size and D == HEAD_DIM and q_all.dtype == torch.float8_e4m3fn + assert rows % q_len == 0 + b = rows // q_len + gh = cp * hq + R = q_len * gh + assert 1 <= R <= MAX_ROWS, f"q_len x group heads = {R} > {MAX_ROWS}" + assert 0 <= cp_rank < cp_size + assert ( + kv_pool.dim() == 4 + and kv_pool.shape[1] == 1 + and kv_pool.shape[2] == PAGE + and kv_pool.shape[3] == HEAD_DIM + ) + assert kv_pool.is_contiguous() and kv_pool.dtype in ( + torch.float8_e4m3fn, + torch.uint8, + ) + # v_pool: V in its own pool of kv_pool's shape; block_tables [b, P] then names the same page of both + if v_pool is not None: + assert ( + v_pool.shape == kv_pool.shape + and v_pool.is_contiguous() + and v_pool.dtype == kv_pool.dtype + ) + if block_tables.dim() == 2: + assert v_pool is not None, "a [b, P] page table needs v_pool" + tab_sb, tab_skv = block_tables.stride(0), 0 + else: + assert block_tables.dim() == 3 and block_tables.shape[1] == 2 + tab_sb, tab_skv = block_tables.stride(0), block_tables.stride(1) + assert block_tables.shape[0] >= b + assert block_tables.stride(-1) == 1 and block_tables.dtype == torch.int32 + assert ( + seq_lens.dtype == torch.int32 + and seq_lens.numel() >= b + and seq_lens.stride(0) == 1 + ) + assert ( + tuple(out_o.shape) == (cp, rows, hq, D) + and out_o.dtype == torch.bfloat16 + and out_o.stride(3) == 1 + ) + assert tuple(out_lse.shape) == (cp, rows, hq) and out_lse.dtype == torch.float32 + assert 1 <= b <= MAX_BATCH + device = q_all.device + sms = sm_count or _sm_count(device) + state = _state(device, sms) + bal = b > _LANE_PLAN_MAX_BATCH if balanced is None else bool(balanced) + assert bal or b <= _LANE_PLAN_MAX_BATCH, "the equal split handles <= 32 requests" + # plan variant: lane (<= 32 requests), wide (33..64), two-slot (65..128); `plan` overrides it in tests (also + # "quad": the one-slot plan for 65..128) + if plan is None: + plan = ( + "two" + if b > _WIDE_PLAN_MAX_BATCH + else ("wide" if b > _LANE_PLAN_MAX_BATCH else "lane") + ) + assert bal or plan == "lane", "plan variants are balanced-split plans" + assert ( + b + <= { + "lane": _LANE_PLAN_MAX_BATCH, + "wide": _WIDE_PLAN_MAX_BATCH, + "quad": _QUAD_PLAN_MAX_BATCH, + "two": _TWO_SLOT_MAX_BATCH, + }[plan] + ) + wide = bal and plan == "wide" + quad = bal and plan == "quad" + two = bal and plan == "two" + two_pack = two and (_TWO_SLOT_PACK if pack is None else bool(pack)) + if two and debug is not None: + assert debug.shape[-1] >= 24, ( + "the two-slot kernel's debug rows hold 24 words (both slots' plans)" + ) + if bal: + splits = n_splits if n_splits else _num_ctas(sms) + assert b < splits <= sms, ( + "balanced split: more CTAs than requests, every CTA co-resident (one per SM)" + ) + else: + splits = n_splits if n_splits else _num_splits(b, sms) + assert splits * b <= sms, ( + "the in-kernel grid barrier needs every CTA co-resident (one per SM)" + ) + if n_splits or balanced is not None: + assert counters is not None, ( + "a split-count or split-mode override needs its own counters tensor" + ) + cnt = state.counters if counters is None else counters + assert ( + cnt.dtype == torch.int64 + and cnt.dim() == 2 + and cnt.shape[0] > b + and cnt.shape[1] >= b + ) + assert cnt.is_contiguous() + if bal: + ws = _balanced_workspace( + state, cp_partial_balanced_ws_bytes(R, splits, 2 if two else 1) + ) + if balanced is None and not state.balanced_logged: + state.balanced_logged = True + logger.info( + f"CP-lite first wide-plan call (q_len {q_len}, batch {b}, " + f"cp_rank {cp_rank}, {splits} CTAs)" + ) + if wide and balanced is None and not state.wide_logged: + state.wide_logged = True + logger.info( + f"CP-lite first call above {_LANE_PLAN_MAX_BATCH} requests on the wide-plan kernel " + f"(q_len {q_len}, batch {b}, cp_rank {cp_rank}, {splits} CTAs)" + ) + if (quad or two) and balanced is None and not state.quad_logged: + state.quad_logged = True + logger.info( + f"CP-lite first call above {_WIDE_PLAN_MAX_BATCH} requests on the " + f"{'two-slot' if two else 'quad-plan'} kernel (q_len {q_len}, " + f"batch {b}, cp_rank {cp_rank}, {splits} CTAs)" + ) + if two_pack and balanced is None and not state.pack_logged: + state.pack_logged = True + logger.info( + f"CP-lite first two-slot call on the packed plan (q_len {q_len}, " + f"batch {b}, cp_rank {cp_rank}, {splits} CTAs)" + ) + else: + need = b * cp_partial_geometry(R, splits)["ws_req_bytes"] + ws = state.workspace + assert ws.numel() * 4 >= need, "CP partial workspace too small" + # q: the native all-gather layout is one [b * rows_per_request, 256] matrix read by TMA in place + native = ( + q_all.stride(3) == 1 + and q_all.stride(2) == D + and q_all.stride(0) == hq * D + and q_all.stride(1) == cp * hq * D + and q_all.data_ptr() % 16 == 0 + ) + q_src = q_all if native else q_all.permute(1, 0, 2, 3).contiguous() + b1_ptr = isinstance(bmm1_scale, torch.Tensor) + b2_ptr = isinstance(bmm2_scale, torch.Tensor) + if b1_ptr: + assert ( + bmm1_scale.numel() >= 2 + and bmm1_scale.dtype == torch.float32 + and bmm1_scale.stride(0) == 1 + ) + if b2_ptr: + assert bmm2_scale.numel() >= 1 and bmm2_scale.dtype == torch.float32 + o_vec16 = out_o.data_ptr() % 16 == 0 and all( + out_o.stride(i) % 8 == 0 for i in range(3) + ) + assert out_o.data_ptr() % 4 == 0 and all( + out_o.stride(i) % 2 == 0 for i in range(3) + ), "out_o: 4-byte aligned rows (alltoall_merge reads it as int32)" + ring = state.dbg_ring if debug is None else None + if bal: + fn = _compiled( + q_len, + gh, + hq, + cp, + splits, + b1_ptr, + b2_ptr, + o_vec16, + bool(enable_pdl), + _KV_PRE_WAIT, + _FRESH_MARGIN if fresh_margin is None else fresh_margin, + debug is not None, + unsafe_skip_wait, + _P_SHIFT if p_shift is None else float(p_shift), + _DBG_SLOTS if (ring is not None and not two) else 0, + True, + rescale_threshold, + bool(l_from_p8), + int(len_sub), + bool(early_trigger), + wide, + quad, + two, + two_pack, + ) + else: + fn = _compiled( + q_len, + gh, + hq, + cp, + splits, + b1_ptr, + b2_ptr, + o_vec16, + bool(enable_pdl), + _KV_PRE_WAIT, + _FRESH_MARGIN if fresh_margin is None else fresh_margin, + debug is not None, + unsafe_skip_wait, + _P_SHIFT if p_shift is None else float(p_shift), + _DBG_SLOTS if ring is not None else 0, + False, + rescale_threshold, + bool(l_from_p8), + int(len_sub), + bool(early_trigger), + ) + fp8 = cutlass.Float8E4M3FN + f32 = cutlass.Float32 + I32 = cutlass.Int32 + a = _ALIGN + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + dummy = ws + fn( + _ptr(q_src, fp8, a["q"]), + _ptr(kv_pool, fp8, a["kv"]), + _ptr(kv_pool if v_pool is None else v_pool, fp8, a["kv"]), + I32(kv_pool.shape[0] * PAGE), + _ptr(block_tables, I32, a["tab"]), + I32(tab_sb), + I32(tab_skv), + _ptr(seq_lens, I32, a["seq"]), + I32(cp_rank), + _ptr(bmm1_scale if b1_ptr else dummy, f32, a["bmm1"]), + f32(0.0 if b1_ptr else float(bmm1_scale) * LOG2E), + _ptr(bmm2_scale if b2_ptr else dummy, f32, a["bmm2"]), + f32(0.0 if b2_ptr else float(bmm2_scale)), + _ptr(out_o, cutlass.BFloat16, 16 if o_vec16 else 4), + I32(out_o.stride(0)), + I32(out_o.stride(1)), + I32(out_o.stride(2)), + _ptr(out_lse, f32, a["lse"]), + I32(out_lse.stride(0)), + I32(out_lse.stride(1)), + I32(out_lse.stride(2)), + _ptr(ws, f32, a["ws"]), + _ptr(cnt[b], cutlass.Int64, a["cnt"]), + _ptr( + debug if debug is not None else (ring if ring is not None else cnt[b]), + cutlass.Int64, + a["dbg"], + ), + I32(b), + stream, + ) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py new file mode 100644 index 00000000..4e546958 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py @@ -0,0 +1,2364 @@ +"""Read-split context parallelism ("CP-lite") for TRT-LLM MHA decode-kernel attention. + +When the TP size exceeds the KV head count, each KV head is replicated on the ``cp`` +ranks of its KV group (``cp = tp // kv_heads``) and every rank reads the whole KV of its +head. CP-lite keeps the KV cache as it is (every rank still writes every token) and +splits the read: page ``p`` of a request's context prefix is read by group rank +``p % cp`` only, for every query head of the group. Per call: + +1. ``_cp_gather_q``: all-gather of the group's FP8 queries through peer stores. +2. TRT-LLM GenMHA (FlashInfer) over the owned prefix pages, the ``q_len`` rows folded + with the group heads into one query row (no causal effect on the prefix), with LSE. +3. ``_cp_a2a_merge``: each rank pushes every owner its heads' partials, then merges its + own heads' partials with the causal ``q_len x q_len`` block of this step's rows. + +The exchange buffer is one symmetric-memory allocation on the attention TP group; each +CTA signals its peers with self-resetting CAS flags, so calls stay graph capturable. + +For multi-request calls, the stock and CP paths are captured together. The batch's KV +lengths select the path on the device from fixed, measured thresholds. On the exact path, +the stock kernel runs beside the gather and partial kernels on a second stream. +""" + +from __future__ import annotations + +import logging +import weakref + +import torch +import triton +import triton.language as tl + +logger = logging.getLogger(__name__) + +PAGE, HEAD_DIM = 64, 256 +_MAX_ROWS = 896 # widest call: rows = requests * q_len (TP8 verify at B128: 896) +_KINDS = 2 # flag regions: q gather, partial exchange +# exact path's fused partial: E4M3 P with shift 8 against the exact running max (lazy rescale threshold 0), so the +# row max maps to 2^8, exact in E4M3 +_EXACT_P_SHIFT, _EXACT_RESCALE_THRESHOLD = 8.0, 0.0 +_LL_PWORDS = 520 # LL words of one (row, source) partial: 4 heads x 128 bf16 pairs + 4 LSE, padded to 16 B + + +@triton.jit +def _flat_tid(): + return tl.inline_asm_elementwise( + "mov.u32 $0, %tid.x;", "=r", [], dtype=tl.uint32, is_pure=True, pack=1 + ) + + +@triton.jit +def _sync_threads(): + tl.inline_asm_elementwise( + "bar.sync 0;", "=r", [], dtype=tl.int32, is_pure=False, pack=1 + ) + + +@triton.jit +def _send_release(addrs): + # set the peer's flag 0 -> 1 once its previous signal is consumed + tl.inline_asm_elementwise( + """ + { + .reg .u32 %t; + .reg .pred %p; + send: + atom.global.release.sys.cas.b32 %t, [$1], 0, 1; + setp.eq.u32 %p, %t, 0; + @!%p bra send; + } + """, + "=r, l", + [addrs], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + + +@triton.jit +def _wait_acquire(addrs): + # consume this rank's flag 1 -> 0 + tl.inline_asm_elementwise( + """ + { + .reg .u32 %t; + .reg .pred %p; + wait: + atom.global.sys.acquire.cas.b32 %t, [$1], 1, 0; + setp.eq.u32 %p, %t, 1; + @!%p bra wait; + } + """, + "=r, l", + [addrs], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + + +@triton.jit +def _group_barrier(peers, me, flags_off, slot, CP: tl.constexpr): + """CTA-level release / acquire with the same CTA of every group rank.""" + c = tl.arange(0, CP) + peer = tl.load(peers + c) + own = tl.load(peers + me) + send = (peer + flags_off + (slot * CP + me) * 4).to(tl.pointer_type(tl.uint32)) + wait = (own + flags_off + (slot * CP + c) * 4).to(tl.pointer_type(tl.uint32)) + if _flat_tid() < CP: + _send_release(send) + _wait_acquire(wait) + + +@triton.jit +def _gate_on(seq_lens, nb, thr, NB: tl.constexpr, GMAX: tl.constexpr): + """CP for this call: the sum (GMAX: the max) of the batch's KV lengths reaches thr. Every rank of the group + reads the same lengths, so all of them take the same branch.""" + k = tl.arange(0, NB) + L = tl.load(seq_lens + k, mask=k < nb, other=0) + if GMAX: + return tl.max(L, axis=0) >= thr + else: + return tl.sum(L, axis=0) >= thr + + +@triton.jit +def _gate_lens( + seq_lens, + nb, + thr, + stock_lens, + pre_lens, + QL: tl.constexpr, + NB: tl.constexpr, + GMAX: tl.constexpr, + ZERO: tl.constexpr, +): + """Once per forward: the stock call's lengths (L with the gate off, q_len keys with it on) and the exact + path's prefix lengths (max(L - QL, 0) on, 0 off).""" + k = tl.arange(0, NB) + m = k < nb + L = tl.load(seq_lens + k, mask=m, other=0) + on = _gate_on(seq_lens, nb, thr, NB, GMAX) + on_keys = 0 if ZERO else tl.minimum(L, QL) + tl.store(stock_lens + k, tl.where(on, on_keys, L), mask=m) + tl.store(pre_lens + k, tl.where(on, tl.maximum(L - QL, 0), 0), mask=m) + + +@triton.jit +def _cp_tables( + page_table, + pt_stride, + seq_lens, + tab, + tab_stride, + lens, + ncols, + rank, + nb, + thr, + QL: tl.constexpr, + CP: tl.constexpr, + PG: tl.constexpr, + BLOCK: tl.constexpr, + GATE: tl.constexpr, + NB: tl.constexpr, + GMAX: tl.constexpr, +): + """This rank's compact page table and token count over the prefix [0, L - QL); with GATE and the gate + off, no owned page and no key.""" + b = tl.program_id(0) + cb = tl.program_id(1) + L = tl.load(seq_lens + b) + prefix = tl.maximum(L - QL, 0) + n_pages = (prefix + PG - 1) // PG + n_own = tl.maximum((n_pages - rank + CP - 1) // CP, 0) + on = 1 + if GATE: + on = _gate_on(seq_lens, nb, thr, NB, GMAX).to(tl.int32) + n_own = n_own * on + j = cb * BLOCK + tl.arange(0, BLOCK) + m = j < ncols + pg = tl.load( + page_table + b * pt_stride + j * CP + rank, mask=m & (j < n_own), other=0 + ) + tl.store(tab + b * tab_stride + j, pg, mask=m) + if cb == 0: + last_tok = prefix - (n_pages - 1) * PG + tail = (n_own > 0) & ((n_own - 1) * CP + rank == n_pages - 1) + n = n_own * PG - tl.where(tail, PG - last_tok, 0) + # no owned page: one dummy key, dropped by the merge; gate off: no key (the merge exits) + tl.store(lens + b, tl.maximum(n, on)) + + +@triton.jit +def _cp_chunk_tables( + tab, + tstride, + lens, + base_lens, + vtab, + vlens, + vstart, + vcnt, + nb, + nv, + ragged, + C: tl.constexpr, + KP: tl.constexpr, + NB: tl.constexpr, + PG: tl.constexpr, +): + """Once per forward, program (r, j), chunked fold partial: with a ragged batch (longest owned length + x B >= ragged x the total, and the longest over C pages), chunk j of request r's owned pages into dense row + start_r + j of vtab ([nv, C]) and no key for the one-row-per-request call; otherwise every chunk row empty and that + call keeps its lengths. Row r * KP + j is cleared when it lies past the last chunk.""" + r = tl.program_id(0) + j = tl.program_id(1) + k = tl.arange(0, NB) + L = tl.load(lens + k, mask=k < nb, other=0) + pg = (L + PG - 1) // PG + mp = tl.max(pg, axis=0) + # Chunk rows should not merely repeat a one-row call. + on = ((mp * nb).to(tl.float32) >= ragged * tl.sum(pg, axis=0).to(tl.float32)) & ( + mp > C + ) + n = tl.where(on, (pg + C - 1) // C, 0) + start = tl.cumsum(n, axis=0) - n + total = tl.sum(n, axis=0) + my_n = tl.sum(tl.where(k == r, n, 0), axis=0) + my_start = tl.sum(tl.where(k == r, start, 0), axis=0) + my_pg = tl.sum(tl.where(k == r, pg, 0), axis=0) + my_L = tl.sum(tl.where(k == r, L, 0), axis=0) + z = r * KP + j + if (z < nv) & (z >= total): + tl.store(vlens + z, 0) + if j == 0: + tl.store(vstart + r, my_start) + tl.store(vcnt + r, my_n) + tl.store(base_lens + r, tl.where(on, 0, my_L)) + if j < my_n: + v = my_start + j + cols = tl.arange(0, C) + src = j * C + cols + m = src < my_pg + tl.store( + vtab + v * C + cols, + tl.load(tab + r * tstride + src, mask=m, other=0), + mask=m, + ) + tl.store(vlens + v, tl.minimum(tl.maximum(my_L - j * (C * PG), 0), C * PG)) + + +@triton.jit +def _cp_chunk_q(q, qv, vstart, vcnt, N, BLK: tl.constexpr): + """Request r's query into its chunk rows (program (r, j): row vstart_r + j when j < vcnt_r).""" + r = tl.program_id(0) + j = tl.program_id(1) + if j < tl.load(vcnt + r): + v = tl.load(vstart + r) + j + for c in range(tl.cdiv(N, BLK)): + w = c * BLK + tl.arange(0, BLK) + m = w < N + tl.store( + qv + v.to(tl.int64) * N + w, + tl.load(q + r.to(tl.int64) * N + w, mask=m), + mask=m, + ) + + +@triton.jit +def _cp_chunk_comb( + po_v, pl_v, vstart, vcnt, po, pl, H, HB: tl.constexpr, D: tl.constexpr +): + """One partial per request from its chunk rows (trtllm-gen LSE is base 2); no write for a request without chunk + rows (the one-row-per-request call wrote it).""" + r = tl.program_id(0) + n = tl.load(vcnt + r) + if n > 0: + h = tl.program_id(1) * HB + tl.arange(0, HB) + hm = h < H + d = tl.arange(0, D) + m = tl.full([HB], float("-inf"), tl.float32) + s = tl.zeros([HB], tl.float32) + acc = tl.zeros([HB, D], tl.float32) + v0 = tl.load(vstart + r) + for j in range(n): + v = v0 + j + lj = tl.load(pl_v + v * H + h, mask=hm, other=0.0) + oj = tl.load( + po_v + (v * H + h[:, None]).to(tl.int64) * D + d[None, :], + mask=hm[:, None], + other=0.0, + ) + mn = tl.maximum(m, lj) + a = tl.exp2(m - mn) + w = tl.exp2(lj - mn) + acc = acc * a[:, None] + oj.to(tl.float32) * w[:, None] + s = s * a + w + m = mn + ok = s > 0 + tl.store( + po + (r * H + h[:, None]).to(tl.int64) * D + d[None, :], + (acc / tl.where(ok, s, 1.0)[:, None]).to(tl.bfloat16), + mask=hm[:, None], + ) + tl.store( + pl + r * H + h, + tl.where(ok, m + tl.log2(tl.where(ok, s, 1.0)), float("-inf")), + mask=hm, + ) + + +@triton.jit +def _cp_gather_q( + q, + q_stride, + peers, + me, + q_off, + flags_off, + CP: tl.constexpr, + WORDS: tl.constexpr, +): + """Row i of this rank's queries into row i, slot me of every group rank's buffer.""" + i = tl.program_id(0) + w = tl.arange(0, WORDS) + x = tl.load(q + i * q_stride + w) + for c in tl.static_range(CP): + base = tl.load(peers + c) + q_off + (i * CP + me) * WORDS * 4 + tl.store(base.to(tl.pointer_type(tl.int32)) + w, x) + _sync_threads() + _group_barrier(peers, me, flags_off, i, CP) + + +@triton.jit +def _cp_a2a_merge( + po, + plse, + q, + q_stride, + k_cache, + v_cache, + kv_page_stride, + kv_tok_stride, + table, + t_stride, + seq_lens, + out, + out_stride, + peers, + me, + p_off, + l_off, + flags_off, + slot0, + bmm1, + bmm2, + QL: tl.constexpr, + HQ: tl.constexpr, + CP: tl.constexpr, + D: tl.constexpr, + PG: tl.constexpr, +): + i = tl.program_id(0) + h = tl.arange(0, HQ) + d = tl.arange(0, D) + hd = h[:, None] * D + d[None, :] + # push owner c its heads' partial O (bf16) and LSE (fp32, base 2) + for c in tl.static_range(CP): + x = tl.load(po + (i * CP * HQ + c * HQ) * D + hd) + lse = tl.load(plse + i * CP * HQ + c * HQ + h) + base = tl.load(peers + c) + dst = (base + p_off).to(tl.pointer_type(tl.bfloat16)) + (i * CP + me) * HQ * D + tl.store(dst + hd, x) + ldst = (base + l_off).to(tl.pointer_type(tl.float32)) + (i * CP + me) * HQ + tl.store(ldst + h, lse) + _sync_threads() + _group_barrier(peers, me, flags_off, slot0 + i, CP) + _sync_threads() + own = tl.load(peers + me) + own_p = (own + p_off).to(tl.pointer_type(tl.bfloat16)) + i * CP * HQ * D + own_l = (own + l_off).to(tl.pointer_type(tl.float32)) + i * CP * HQ + b = i // QL + t = i % QL + L = tl.load(seq_lens + b) + # the causal block of this step's rows, for this rank's heads + j = tl.arange(0, 8) + pos = L - QL + j + valid = (j <= t) & (pos >= 0) + pos = tl.maximum(pos, 0) + page = tl.load(table + b * t_stride + pos // PG, mask=valid, other=0) + koff = page.to(tl.int64) * kv_page_stride + (pos % PG) * kv_tok_stride + kk = tl.load( + k_cache + koff[:, None] + d[None, :], mask=valid[:, None], other=0.0 + ).to(tl.float32) + vv = tl.load( + v_cache + koff[:, None] + d[None, :], mask=valid[:, None], other=0.0 + ).to(tl.float32) + qv = tl.load(q + i * q_stride + hd).to(tl.float32) + s = tl.sum(qv[:, None, :] * kk[None, :, :], axis=2) * bmm1 + s = tl.where(valid[None, :], s, float("-inf")) + md = tl.max(s, axis=1) + md0 = tl.where(md > float("-inf"), md, 0.0) + p = tl.exp(s - md0[:, None]) + ld = tl.sum(p, axis=1) + od = ( + tl.sum(p[:, :, None] * vv[None, :, :], axis=1) + / tl.maximum(ld, 1e-30)[:, None] + * bmm2 + ) + lse_d = tl.where(ld > 0, md0 + tl.log(ld), float("-inf")) + # rank c owns a prefix page iff the prefix has more than c pages + n_pages = (tl.maximum(L - QL, 0) + PG - 1) // PG + m = lse_d + for c in tl.static_range(CP): + lc = tl.load(own_l + c * HQ + h, cache_modifier=".cg") * 0.6931471805599453 + lc = tl.where(n_pages > c, lc, float("-inf")) + m = tl.maximum(m, lc) + m0 = tl.where(m > float("-inf"), m, 0.0) + w = tl.where(lse_d > float("-inf"), tl.exp(lse_d - m0), 0.0) + acc = od * w[:, None] + for c in tl.static_range(CP): + lc = tl.load(own_l + c * HQ + h, cache_modifier=".cg") * 0.6931471805599453 + wc = tl.where(n_pages > c, tl.exp(lc - m0), 0.0) + oc = tl.load(own_p + c * HQ * D + hd, cache_modifier=".cg").to(tl.float32) + acc += oc * wc[:, None] + w += wc + o = acc / tl.where(w > 0, w, 1.0)[:, None] + tl.store(out + i * out_stride + hd, o.to(tl.bfloat16)) + + +@triton.jit +def _ll_pack(flag, w): + return flag | (w.to(tl.int64) & 0xFFFFFFFF) + + +@triton.jit +def _ll_spin(addr, e): + # one thread per word: volatile reads until the word's high half is e + return tl.inline_asm_elementwise( + """ + { + .reg .pred %p; + .reg .b32 %lo, %hi; + spin: + ld.volatile.global.b64 $0, [$1]; + mov.b64 {%lo, %hi}, $0; + setp.ne.b32 %p, %hi, $2; + @%p bra spin; + } + """, + "=l,l,r", + args=[addr, e], + dtype=tl.int64, + is_pure=False, + pack=1, + ) + + +@triton.jit +def _ll_wait(src, n: tl.constexpr, e): + """Poll n LL words (data in the low, epoch in the high half) until all carry epoch e.""" + # each thread waits for the words it returns: with a block-wide count, a warp could leave holding a stale copy of + # a word another warp also held (polls narrower than the CTA); the barrier keeps every warp before the epoch store + v = _ll_spin(src + tl.arange(0, n), e) + if n > 32: + tl.debug_barrier() + return v + + +@triton.jit +def _cp_gather_q_ll( + q, + q_stride, + peers, + me, + ll_off, + q_all, + cnt, + seq_lens, + pre_lens, + stock_lens, + nb, + thr, + CP: tl.constexpr, + WORDS: tl.constexpr, + QL: tl.constexpr, + FOLD: tl.constexpr, + GATE: tl.constexpr, + NB: tl.constexpr, + GMAX: tl.constexpr, + GWRITE: tl.constexpr, + TRIG: tl.constexpr = False, + WAIT: tl.constexpr = False, + qv=None, + vstart=None, + vcnt=None, + CHUNK: tl.constexpr = False, +): + """LL variant of _cp_gather_q: every 4-byte word travels with the call's epoch in one 8-byte store. + FOLD: q_all is written source-major per request ([B, CP, QL, heads, D]) for the prefix-only fused partial, + and pre_lens[b] = max(L - QL, 0). + GATE: exits with the gate off; GWRITE: first writes the stock call's lengths (L with the gate off, q_len with it + on) and, FOLD, pre_lens = 0 with the gate off (else _gate_lens wrote them). + TRIG (no GWRITE): releases the dependent grid at entry; the fused partial reads only lengths, tables and prefix KV + before its griddepcontrol.wait, none of them written here. + WAIT: launched as a programmatic dependent of q's writer, which writes none of what the partial reads before its + wait either; every access here follows the wait.""" + if TRIG: + tl.extra.cuda.gdc_launch_dependents() + if WAIT: + _gdc_wait_clobber() + i = tl.program_id(0) + if GATE: + on = _gate_on(seq_lens, nb, thr, NB, GMAX) + if GWRITE and i % QL == 0: + Lb = tl.load(seq_lens + i // QL) + tl.store(stock_lens + i // QL, tl.where(on, tl.minimum(Lb, QL), Lb)) + if FOLD: + tl.store(pre_lens + i // QL, tl.where(on, tl.maximum(Lb - QL, 0), 0)) + if on.to(tl.int32) == 0: + return + e = tl.load(cnt + i) + 1 + flag = e.to(tl.int64) << 32 + w = tl.arange(0, WORDS) + x = _ll_pack(flag, tl.load(q + i * q_stride + w)) + for c in tl.static_range(CP): + dst = (tl.load(peers + c) + ll_off + (i * CP + me) * WORDS * 8).to( + tl.pointer_type(tl.int64) + ) + tl.store(dst + w, x) + own = tl.load(peers + me) + ll_off + for c in tl.static_range(CP): + v = _ll_wait( + (own + (i * CP + c) * WORDS * 8).to(tl.pointer_type(tl.int64)), WORDS, e + ) + if FOLD: + slot = ((i // QL) * CP + c) * QL + i % QL + else: + slot = i * CP + c + val = (v & 0xFFFFFFFF).to(tl.int32) + tl.store(q_all + slot * WORDS + w, val) + if CHUNK: + v0 = tl.load(vstart + i // QL) + for j in range(tl.load(vcnt + i // QL)): + tl.store( + qv + + ((v0 + j) * (QL * CP) + (i % QL) * CP + c).to(tl.int64) * WORDS + + w, + val, + ) + if FOLD: # noqa: SIM102 + if not GATE: # noqa: SIM102 + if i % QL == 0: + tl.store( + pre_lens + i // QL, + tl.maximum(tl.load(seq_lens + i // QL) - QL, 0), + ) + tl.store(cnt + i, e) + + +@triton.jit +def _cp_a2a_merge_ll( + po, + plse, + po_sr, + po_sc, + pl_sr, + pl_sc, + q, + q_stride, + k_cache, + v_cache, + kv_page_stride, + kv_tok_stride, + table, + t_stride, + seq_lens, + out, + out_stride, + peers, + me, + ll_off, + cnt, + bmm1, + bmm2, + nb, + thr, + QL: tl.constexpr, + HQ: tl.constexpr, + CP: tl.constexpr, + D: tl.constexpr, + PG: tl.constexpr, + NWP: tl.constexpr, + GATE: tl.constexpr, + NB: tl.constexpr, + GMAX: tl.constexpr, +): + """LL variant of _cp_a2a_merge; po / plse are read as int32 words, row / member strides po_sr / po_sc (words) + and pl_sr / pl_sc. Online LSE merge of the causal block and the CP partials as they arrive. + GATE: exits with the gate off (out keeps the stock call's rows).""" + i = tl.program_id(0) + if GATE: # noqa: SIM102 + if _gate_on(seq_lens, nb, thr, NB, GMAX).to(tl.int32) == 0: + return + e = tl.load(cnt + i) + 1 + flag = e.to(tl.int64) << 32 + NO: tl.constexpr = HQ * D // 2 + hw = tl.arange(0, NO) + h = tl.arange(0, HQ) + d = tl.arange(0, D) + hd = h[:, None] * D + d[None, :] + for c in tl.static_range(CP): + w = _ll_pack(flag, tl.load(po + i * po_sr + c * po_sc + hw)) + lw = _ll_pack(flag, tl.load(plse + i * pl_sr + c * pl_sc + h)) + dst = (tl.load(peers + c) + ll_off + (i * CP + me) * NWP * 8).to( + tl.pointer_type(tl.int64) + ) + tl.store(dst + hw, w) + tl.store(dst + NO + h, lw) + b = i // QL + t = i % QL + L = tl.load(seq_lens + b) + j = tl.arange(0, 8) + pos = L - QL + j + valid = (j <= t) & (pos >= 0) + pos = tl.maximum(pos, 0) + page = tl.load(table + b * t_stride + pos // PG, mask=valid, other=0) + koff = page.to(tl.int64) * kv_page_stride + (pos % PG) * kv_tok_stride + kk = tl.load( + k_cache + koff[:, None] + d[None, :], mask=valid[:, None], other=0.0 + ).to(tl.float32) + vv = tl.load( + v_cache + koff[:, None] + d[None, :], mask=valid[:, None], other=0.0 + ).to(tl.float32) + qv = tl.load(q + i * q_stride + hd).to(tl.float32) + s = tl.sum(qv[:, None, :] * kk[None, :, :], axis=2) * bmm1 + s = tl.where(valid[None, :], s, float("-inf")) + md = tl.max(s, axis=1) + md0 = tl.where(md > float("-inf"), md, 0.0) + p = tl.exp(s - md0[:, None]) + ld = tl.sum(p, axis=1) + acc = tl.sum(p[:, :, None] * vv[None, :, :], axis=1) * bmm2 + m = tl.where(ld > 0, md0, float("-inf")) + wsum = ld + n_pages = (tl.maximum(L - QL, 0) + PG - 1) // PG + own = tl.load(peers + me) + ll_off + for c in tl.static_range(CP): + src = (own + (i * CP + c) * NWP * 8).to(tl.pointer_type(tl.int64)) + v = _ll_wait(src, NO, e) + lv = _ll_wait(src + NO, HQ, e) + lc = (lv & 0xFFFFFFFF).to(tl.int32).to( + tl.float32, bitcast=True + ) * 0.6931471805599453 + lc = tl.where(n_pages > c, lc, float("-inf")) + wd = (v & 0xFFFFFFFF).to(tl.uint32) + lo = (wd << 16).to(tl.float32, bitcast=True) + hi = (wd & 0xFFFF0000).to(tl.float32, bitcast=True) + oc = tl.reshape(tl.join(lo, hi), (HQ, D)) + mn = tl.maximum(m, lc) + mn0 = tl.where(mn > float("-inf"), mn, 0.0) + a_old = tl.where(m > float("-inf"), tl.exp(m - mn0), 0.0) + a_new = tl.where(lc > float("-inf"), tl.exp(lc - mn0), 0.0) + acc = acc * a_old[:, None] + oc * a_new[:, None] + wsum = wsum * a_old + a_new + m = mn + o = acc / tl.where(wsum > 0, wsum, 1.0)[:, None] + tl.store(out + i * out_stride + hd, o.to(tl.bfloat16)) + tl.store(cnt + i, e) + + +@triton.jit +def _cp_a2a_merge4_ll( + po, + plse, + out, + out_stride, + peers, + me, + ll_off, + cnt, + HQ: tl.constexpr, + CP: tl.constexpr, + D: tl.constexpr, + NWP: tl.constexpr, +): + """_cp_a2a_merge_ll for fused partials, which already hold the causal block (O = 0 / LSE = -inf where a rank + sees no key): push, then a base-2 online LSE merge of the CP partials only.""" + i = tl.program_id(0) + e = tl.load(cnt + i) + 1 + flag = e.to(tl.int64) << 32 + NO: tl.constexpr = HQ * D // 2 + hw = tl.arange(0, NO) + h = tl.arange(0, HQ) + d = tl.arange(0, D) + for c in tl.static_range(CP): + w = _ll_pack(flag, tl.load(po + (i * CP * HQ + c * HQ) * (D // 2) + hw)) + lw = _ll_pack(flag, tl.load(plse + i * CP * HQ + c * HQ + h)) + dst = (tl.load(peers + c) + ll_off + (i * CP + me) * NWP * 8).to( + tl.pointer_type(tl.int64) + ) + tl.store(dst + hw, w) + tl.store(dst + NO + h, lw) + m = tl.full([HQ], float("-inf"), tl.float32) + wsum = tl.zeros([HQ], dtype=tl.float32) + acc = tl.zeros([HQ, D], dtype=tl.float32) + own = tl.load(peers + me) + ll_off + for c in tl.static_range(CP): + src = (own + (i * CP + c) * NWP * 8).to(tl.pointer_type(tl.int64)) + v = _ll_wait(src, NO, e) + lv = _ll_wait(src + NO, HQ, e) + lc = (lv & 0xFFFFFFFF).to(tl.int32).to(tl.float32, bitcast=True) + wd = (v & 0xFFFFFFFF).to(tl.uint32) + lo = (wd << 16).to(tl.float32, bitcast=True) + hi = (wd & 0xFFFF0000).to(tl.float32, bitcast=True) + oc = tl.reshape(tl.join(lo, hi), (HQ, D)) + mn = tl.maximum(m, lc) + mn0 = tl.where(mn > float("-inf"), mn, 0.0) + a_old = tl.where(m > float("-inf"), tl.exp2(m - mn0), 0.0) + a_new = tl.where(lc > float("-inf"), tl.exp2(lc - mn0), 0.0) + acc = acc * a_old[:, None] + oc * a_new[:, None] + wsum = wsum * a_old + a_new + m = mn + o = acc / tl.where(wsum > 0, wsum, 1.0)[:, None] + tl.store(out + i * out_stride + h[:, None] * D + d[None, :], o.to(tl.bfloat16)) + tl.store(cnt + i, e) + + +@triton.jit +def _gdc_wait_clobber(): + # griddepcontrol.wait with a memory clobber: no load may be hoisted above it + tl.inline_asm_elementwise( + "griddepcontrol.wait; // dummy $0", + "=r,~{memory}", + [], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + + +@triton.jit +def _causal_block_h( + q, + q_stride, + i, + h, + L, + k_cache, + v_cache, + kv_page_stride, + kv_tok_stride, + table, + t_stride, + bmm1, + bmm2, + QL: tl.constexpr, + D: tl.constexpr, + PG: tl.constexpr, +): + """Head h of row i: the fp32 causal block over this step's QL keys as _cp_a2a_merge_ll starts its merge, i.e. + unnormalized O (x bmm2) [D], the natural-log max [1] (-inf: no key) and the row sum [1].""" + b = i // QL + t = i % QL + d = tl.arange(0, D) + j = tl.arange(0, 8) + pos = L - QL + j + valid = (j <= t) & (pos >= 0) + pos = tl.maximum(pos, 0) + page = tl.load(table + b * t_stride + pos // PG, mask=valid, other=0) + koff = page.to(tl.int64) * kv_page_stride + (pos % PG) * kv_tok_stride + kk = tl.load( + k_cache + koff[:, None] + d[None, :], mask=valid[:, None], other=0.0 + ).to(tl.float32) + vv = tl.load( + v_cache + koff[:, None] + d[None, :], mask=valid[:, None], other=0.0 + ).to(tl.float32) + qv = tl.load(q + i * q_stride + h * D + d).to(tl.float32) + s = tl.sum(qv[None, :] * kk, axis=1) * bmm1 + s = tl.where(valid, s, float("-inf")) + md = tl.max(s, axis=0) + md0 = tl.where(md > float("-inf"), md, 0.0) + p = tl.exp(s - md0) + ld = tl.sum(p, axis=0) + acc = tl.sum(p[:, None] * vv, axis=0) * bmm2 + m = tl.where(ld > 0, md0, float("-inf")) + return acc, m + tl.zeros([1], tl.float32), ld + tl.zeros([1], tl.float32) + + +@triton.jit +def _cp_merge_h_ll( + po, + plse, + po_sr, + po_sc, + pl_sr, + pl_sc, + q, + q_stride, + k_cache, + v_cache, + kv_page_stride, + kv_tok_stride, + table, + t_stride, + seq_lens, + out, + out_stride, + peers, + me, + ll_off, + cnt, + bmm1, + bmm2, + nb, + thr, + QL: tl.constexpr, + HQ: tl.constexpr, + CP: tl.constexpr, + D: tl.constexpr, + PG: tl.constexpr, + NWP: tl.constexpr, + NB: tl.constexpr, + GMAX: tl.constexpr, + WAIT: tl.constexpr, + EARLY: tl.constexpr = False, + po_v=None, + pl_v=None, + vstart=None, + vcnt=None, + CHUNK: tl.constexpr = False, + gate=None, + gate_stride=0, + o_proj_input_scale=None, + o_proj_input=None, + o_proj_input_stride=0, + GATE_EPILOGUE: tl.constexpr = False, +): + """_cp_a2a_merge_ll (gated) with one program per (row, head): pushes the head's partials, computes its causal + block, then merges the CP partials as they arrive. One epoch per (row, head) in cnt, in an LL region of its own + (the row merge's epochs differ). WAIT: launched as a programmatic dependent. + EARLY (with WAIT): the causal block before the wait, which then guards only the partials; + its predecessor releases it only after its own wait, so the step's q / K / V (written upstream of that wait) are + complete. + CHUNK (fold layout): a request with chunk rows sends the base-2 LSE combine of its rows' partials + (po_v [nv, QL * CP * HQ, D] bf16, pl_v [nv, QL * CP * HQ] fp32) instead of po / plse. + GATE_EPILOGUE also writes the gated, statically scaled FP8 input of o_proj.""" + i = tl.program_id(0) + h = tl.program_id(1) + if WAIT and not EARLY: + _gdc_wait_clobber() + if _gate_on(seq_lens, nb, thr, NB, GMAX).to(tl.int32) == 0: + return + one = tl.arange(0, 1) + e = tl.load(cnt + i * HQ + h) + 1 + flag = e.to(tl.int64) << 32 + DW: tl.constexpr = D // 2 + NO: tl.constexpr = HQ * D // 2 + H: tl.constexpr = QL * CP * HQ + w = tl.arange(0, DW) + L = tl.load(seq_lens + i // QL) + if EARLY: + acc, m, wsum = _causal_block_h( + q, + q_stride, + i, + h, + L, + k_cache, + v_cache, + kv_page_stride, + kv_tok_stride, + table, + t_stride, + bmm1, + bmm2, + QL, + D, + PG, + ) + if WAIT: + _gdc_wait_clobber() + if CHUNK: + n_v = tl.load(vcnt + i // QL) + v0 = tl.load(vstart + i // QL) + for c in tl.static_range(CP): + if CHUNK: + if n_v > 0: + hh = (i % QL) * (CP * HQ) + c * HQ + h + mx = tl.full([1], float("-inf"), tl.float32) + sw = tl.zeros([1], tl.float32) + acc_lo = tl.zeros([DW], tl.float32) + acc_hi = tl.zeros([DW], tl.float32) + for j in range(n_v): + vr = v0 + j + lj = tl.load(pl_v + vr * H + hh + one) + wd = tl.load(po_v + (vr * H + hh).to(tl.int64) * DW + w).to( + tl.uint32, bitcast=True + ) + mn = tl.maximum(mx, lj) + a = tl.exp2(mx - mn) + b2 = tl.exp2(lj - mn) + acc_lo = acc_lo * a + (wd << 16).to(tl.float32, bitcast=True) * b2 + acc_hi = ( + acc_hi * a + (wd & 0xFFFF0000).to(tl.float32, bitcast=True) * b2 + ) + sw = sw * a + b2 + mx = mn + inv = 1.0 / tl.where(sw > 0, sw, 1.0) + lo16 = (acc_lo * inv).to(tl.bfloat16).to(tl.int16, bitcast=True).to( + tl.int32 + ) & 0xFFFF + hi16 = (acc_hi * inv).to(tl.bfloat16).to(tl.int16, bitcast=True).to( + tl.int32 + ) & 0xFFFF + lse2 = tl.where( + sw > 0, + mx + tl.log2(tl.where(sw > 0, sw, 1.0)), + float("-inf"), + ) + ow = _ll_pack(flag, lo16 | (hi16 << 16)) + lw = _ll_pack(flag, lse2.to(tl.int32, bitcast=True)) + else: + ow = _ll_pack(flag, tl.load(po + i * po_sr + c * po_sc + h * DW + w)) + lw = _ll_pack(flag, tl.load(plse + i * pl_sr + c * pl_sc + h + one)) + else: + ow = _ll_pack(flag, tl.load(po + i * po_sr + c * po_sc + h * DW + w)) + lw = _ll_pack(flag, tl.load(plse + i * pl_sr + c * pl_sc + h + one)) + dst = (tl.load(peers + c) + ll_off + (i * CP + me) * NWP * 8).to( + tl.pointer_type(tl.int64) + ) + tl.store(dst + h * DW + w, ow) + tl.store(dst + NO + h + one, lw) + if not EARLY: + acc, m, wsum = _causal_block_h( + q, + q_stride, + i, + h, + L, + k_cache, + v_cache, + kv_page_stride, + kv_tok_stride, + table, + t_stride, + bmm1, + bmm2, + QL, + D, + PG, + ) + n_pages = (tl.maximum(L - QL, 0) + PG - 1) // PG + own = tl.load(peers + me) + ll_off + for c in tl.static_range(CP): + src = (own + (i * CP + c) * NWP * 8).to(tl.pointer_type(tl.int64)) + v = _ll_wait(src + h * DW, DW, e) + lv = _ll_wait(src + NO + h, 1, e) + lc = (lv & 0xFFFFFFFF).to(tl.int32).to( + tl.float32, bitcast=True + ) * 0.6931471805599453 + lc = tl.where(n_pages > c, lc, float("-inf")) + wd = (v & 0xFFFFFFFF).to(tl.uint32) + lo = (wd << 16).to(tl.float32, bitcast=True) + hi = (wd & 0xFFFF0000).to(tl.float32, bitcast=True) + oc = tl.reshape(tl.join(lo, hi), (D,)) + mn = tl.maximum(m, lc) + mn0 = tl.where(mn > float("-inf"), mn, 0.0) + a_old = tl.where(m > float("-inf"), tl.exp(m - mn0), 0.0) + a_new = tl.where(lc > float("-inf"), tl.exp(lc - mn0), 0.0) + acc = acc * a_old + oc * a_new + wsum = wsum * a_old + a_new + m = mn + o = (acc / tl.where(wsum > 0, wsum, 1.0)).to(tl.bfloat16) + d = tl.arange(0, D) + tl.store(out + i * out_stride + h * D + d, o) + if GATE_EPILOGUE: + g = tl.load(gate + i * gate_stride + h * D + d).to(tl.float32) + gated = (o.to(tl.float32) * tl.sigmoid(g)).to(tl.bfloat16).to(tl.float32) + scale_inv = 1.0 / tl.load(o_proj_input_scale).to(tl.float32) + tl.store( + o_proj_input + i * o_proj_input_stride + h * D + d, + tl.clamp(gated * scale_inv, -448.0, 448.0).to(tl.float8e4nv), + ) + tl.store(cnt + i * HQ + h + one, e) + + +class CpLite: + """One process's exchange buffer and scratch, shared by the Qwen3.5 target and + draft slots; built collectively on the plugin binding's TP group.""" + + _instance: CpLite | None = None + + def __init__( + self, group, tp_rank: int, tp_size: int, cp: int, q_heads: int, device + ): + import torch.distributed._symmetric_memory as symm_mem + + self.cp, self.hq, self.device = cp, q_heads, torch.device(device) + self.me = tp_rank % cp + g0 = tp_rank - self.me + self.min_bs = 2 + self.q_bytes = _MAX_ROWS * cp * q_heads * HEAD_DIM # fp8 + self.p_off = self.q_bytes + self.p_bytes = _MAX_ROWS * cp * q_heads * HEAD_DIM * 2 # bf16 + self.l_off = self.p_off + self.p_bytes + self.l_bytes = _MAX_ROWS * cp * q_heads * 4 + self.f_off = self.l_off + self.l_bytes + self.llq_off = self.f_off + _KINDS * _MAX_ROWS * cp * 4 + self.llp_off = self.llq_off + _MAX_ROWS * cp * (q_heads * HEAD_DIM // 4) * 8 + total = self.llp_off + _MAX_ROWS * cp * _LL_PWORDS * 8 + self.cp_ungated = True + self.gather_pdl = self.cp_ungated + self.early_merge = self.cp_ungated + self.merge_h = True + self.fold_mh = True + self.fold_chunk = {33: 128, 41: 64, 65: 128} + self.fold_chunk_q1 = {65: 128} + self.fold_ragged = 1.8 + self.fold_fuse = True + self.llh_off = total + total += _MAX_ROWS * cp * _LL_PWORDS * 8 + self.ll = q_heads * HEAD_DIM // 2 + q_heads <= _LL_PWORDS + self.cnt = torch.zeros(_KINDS, _MAX_ROWS, dtype=torch.int32, device=self.device) + self.sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.fused = self.ll + self.exact = self.fused + if self.fused: + from sgl_nvidia_mega.qwen3_5.sm103.attention import q35_cp_fused + + # grid-barrier counters and workspace: persistent, allocated before any capture + q35_cp_fused.q35_cp_fused_prepare(self.device, self.sms) + if symm_mem.get_backend(self.device) is None: + symm_mem.set_backend("NCCL") + self.buf = symm_mem.empty((total,), dtype=torch.uint8, device=self.device) + self.buf.zero_() + torch.cuda.synchronize(self.device) + handle = symm_mem.rendezvous(self.buf, group) + assert handle.rank == tp_rank, (handle.rank, tp_rank) + self.peers = torch.tensor( + [int(p) for p in handle.buffer_ptrs[g0 : g0 + cp]], + dtype=torch.int64, + device=self.device, + ) + self._handle = handle + self.q_all = ( + self.buf[: self.q_bytes] + .view(torch.float8_e4m3fn) + .view(_MAX_ROWS, cp * q_heads, HEAD_DIM) + ) + rows = _MAX_ROWS + self.po = torch.empty( + rows, cp * q_heads, HEAD_DIM, dtype=torch.bfloat16, device=self.device + ) + self.plse = torch.empty( + rows, cp * q_heads, dtype=torch.float32, device=self.device + ) + # multi-CTA counters of the trtllm-gen calls: rows x heads (the chunk rows: up to 8 B x 112 at q_len 7) + self.counter = torch.zeros(1 << 20, dtype=torch.int32, device=self.device) + # exact path: the fused kernel's [cp, rows, heads, D] partials and the prefix lengths the gather writes + self.po_x = torch.empty( + cp, rows, q_heads, HEAD_DIM, dtype=torch.bfloat16, device=self.device + ) + self.pl_x = torch.empty( + cp, rows, q_heads, dtype=torch.float32, device=self.device + ) + self.pre_lens = torch.zeros(rows, dtype=torch.int32, device=self.device) + self.min_kv = {} + self.gate_max = False + self.stock_lens = torch.zeros(rows, dtype=torch.int32, device=self.device) + self.gate_overlap = 1 + self.cp_main = True + self.cp_pdl = self.cp_main + self._gate_epilogue = None + self._gate_epilogue_done = False + self.cnt_h = torch.zeros( + _MAX_ROWS * q_heads, dtype=torch.int32, device=self.device + ) + self._side = torch.cuda.Stream(device=self.device) + self._ws = None + if self.min_kv and self.gate_overlap: + self._own_workspace() + # (lengths, b, q_len, threshold) -> [(stock lengths, prefix lengths), forward, last layer id that used them] + self._glens: dict = {} + self._logged: set = set() + # (page table, lengths, q_len) -> [tables, owned lengths, forward, last layer id that used them] + self._tabs: dict = {} + torch.distributed.barrier(group) + logger.info( + "TRT-L1 CP-lite: tp rank %d, KV group of %d (group rank %d), min batch %d, LL %s, fused %s", + tp_rank, + cp, + self.me, + self.min_bs, + self.ll, + self.fused, + ) + logger.info( + "TRT-L1 CP-lite exact (prefix-only fused partial + fp32 causal block): %s", + self.exact, + ) + logger.info( + "TRT-L18 CP-lite current-stream CP: %s, PDL: %s, per-head merge: %s, ungated exact: %s, gather PDL: %s, early merge: %s", + self.cp_main, + self.cp_pdl, + self.merge_h, + self.cp_ungated, + self.gather_pdl, + self.early_merge, + ) + logger.info( + "TRT-L1 CP-lite KV gate: minimum mean KV per request %s", self.min_kv + ) + + def offer_gate_epilogue( + self, gate: torch.Tensor, scale: torch.Tensor, out: torch.Tensor + ) -> None: + """Offer the next exact CP-lite call the attention gate and FP8 o_proj input.""" + self._gate_epilogue = (gate, scale, out) + self._gate_epilogue_done = False + + def take_gate_epilogue(self) -> bool: + """Return whether the last CP-lite call wrote the offered o_proj input.""" + done = self._gate_epilogue_done + self._gate_epilogue = None + self._gate_epilogue_done = False + return done + + def _gate_thr(self, b: int) -> int | None: + """The device threshold of a b-request call (on the length sum, or with GATE_MAX on the max), or None.""" + keys = [k for k in self.min_kv if k <= b] + if not keys: + return None + t = self.min_kv[max(keys)] + return t if self.gate_max else t * b + + @classmethod + def maybe_create(cls, binding, device) -> CpLite | None: + """Build the Qwen3.5 CP state collectively on the attention TP group. + + The plugin spec pins the global Q/KV head counts and head dimension. Keeping + those model-specific facts here avoids a runtime-side prebuild hook. Every + rank calls this from plugin build before graph variants are validated. + """ + from sgl_nvidia_mega.qwen3_5 import TEXT_CONFIG + + if cls._instance is not None: + return cls._instance + tp, rank = binding.tp_size, binding.tp_rank + kv = TEXT_CONFIG["num_key_value_heads"] + if tp <= kv or tp % kv or TEXT_CONFIG["head_dim"] != HEAD_DIM: + return None + cp = tp // kv + q_heads = TEXT_CONFIG["num_attention_heads"] // tp + if cp not in (2, 4, 8) or q_heads * cp > 16: + return None + group = binding.tp_device_group + cls._instance = cls(group, rank, tp, cp, q_heads, device) + return cls._instance + + def _tables(self, page_table, seq_lens, q_len, forward_batch, layer_id, gate): + """Once per forward: reused only by later layers of the same forward, so every capture and + every single-layer draft step computes its own.""" + b = seq_lens.shape[0] + ncols = triton.cdiv(page_table.shape[1], self.cp) + key = ( + page_table.data_ptr(), + page_table.shape, + seq_lens.data_ptr(), + b, + q_len, + gate["GATE"], + gate["thr"], + ) + ent = self._tabs.get(key) + if ent is not None and ent[2]() is forward_batch and layer_id > ent[3]: + ent[3] = layer_id + return ent[0], ent[1], ncols * PAGE + if len(self._tabs) > 64: + self._tabs.clear() + tab = torch.empty(b, ncols, dtype=torch.int32, device=self.device) + lens = torch.empty(b, dtype=torch.int32, device=self.device) + self._tabs[key] = [tab, lens, weakref.ref(forward_batch), layer_id] + _cp_tables[(b, triton.cdiv(ncols, 256))]( + page_table, + page_table.stride(0), + seq_lens, + tab, + tab.stride(0), + lens, + ncols, + self.me, + **gate, + QL=q_len, + CP=self.cp, + PG=PAGE, + BLOCK=256, + ) + return tab, lens, ncols * PAGE + + def attend( + self, + q: torch.Tensor, + kv_cache, + page_table: torch.Tensor, + seq_lens: torch.Tensor, + q_len: int, + bmm1_scale, + bmm2_scale, + workspace: torch.Tensor, + layer, + forward_batch, + stock=None, + min_bs: int | None = None, + ) -> torch.Tensor | None: + """CP-lite attention of this rank's heads: [rows, heads, head_dim] bf16, or None + where the call is not served (decided by shapes only, so every rank agrees). + stock(seq_lens, out): the caller's own kernel call into out; with it and a KV gate for this batch size, + both are captured and the device picks one per call. + min_bs: override the smallest batch this call serves; a caller that picks CP by + its own host-side KV length can pass 1.""" + rows = q.shape[0] + b = seq_lens.shape[0] + if ( + q.dtype != torch.float8_e4m3fn + or q.dim() != 3 + or q.shape[1] != self.hq + or q.shape[2] != HEAD_DIM + or layer.tp_k_head_num != 1 + or layer.sliding_window_size not in (None, -1) + or rows != b * q_len + or b < (self.min_bs if min_bs is None else min_bs) + or rows > _MAX_ROWS + or not isinstance(bmm1_scale, float) + or not isinstance(bmm2_scale, float) + ): + return None + if (b, q_len) not in self._logged: + self._logged.add((b, q_len)) + logger.info( + "TRT-L1 CP-lite engaged: %d requests x %d rows (layer %d)", + b, + q_len, + layer.layer_id, + ) + k_cache, v_cache = kv_cache + q = q if q.is_contiguous() else q.contiguous() + words = self.hq * HEAD_DIM // 4 + exact = self.exact and self._exact_supported(q_len, b) + # Every path but the non-exact fused one can be gated; exact batches above + # the fused kernel's limit take the TRT-LLM partial. + thr = ( + self._gate_thr(b) + if stock is not None and self.ll and (self.exact or not self.fused) + else None + ) + gate = { + "nb": b, + "thr": 0 if thr is None else thr, + "GATE": thr is not None, + "NB": triton.next_power_of_2(b), + "GMAX": self.gate_max, + } + out = None + if thr is not None and self.gate_overlap: + return self._gated_overlap( + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + forward_batch, + layer.layer_id, + stock, + gate, + exact, + ) + if exact and thr is None and self.ll and self.cp_ungated: + return self._ungated_main( + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + gate, + ) + if self.ll: + _cp_gather_q_ll[(rows,)]( + q.view(torch.int32), + words, + self.peers, + self.me, + self.llq_off, + self.q_all.view(torch.int32), + self.cnt[0], + seq_lens, + self.pre_lens, + self.stock_lens, + **gate, + CP=self.cp, + WORDS=words, + QL=q_len, + FOLD=exact, + GWRITE=True, + num_warps=4, + ) + if thr is not None: + out = torch.empty( + rows, self.hq, HEAD_DIM, dtype=torch.bfloat16, device=q.device + ) + stock(self.stock_lens[:b], out) + else: + _cp_gather_q[(rows,)]( + q.view(torch.int32), + words, + self.peers, + self.me, + 0, + self.f_off, + CP=self.cp, + WORDS=words, + num_warps=4, + ) + if exact: + return self._exact( + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + out, + gate, + ) + if self.fused and not self.exact: + out = self._fused( + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + ) + if out is not None: + return out + self._fold_partial( + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + forward_batch, + layer.layer_id, + gate, + workspace, + ) + if out is None: + out = torch.empty( + rows, self.hq, HEAD_DIM, dtype=torch.bfloat16, device=q.device + ) + if self.ll and self.fold_mh: + fold_gate = gate if gate["GATE"] else {**gate, "thr": 0} + self._fold_merge_h( + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + rows, + bmm1_scale, + bmm2_scale, + out, + fold_gate, + ) + elif self.ll: + self._fold_merge( + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + rows, + bmm1_scale, + bmm2_scale, + out, + gate, + ) + else: + _cp_a2a_merge[(rows,)]( + self.po, + self.plse, + q, + q.stride(0), + k_cache, + v_cache, + k_cache.stride(0), + k_cache.stride(2), + page_table, + page_table.stride(0), + seq_lens, + out, + out.stride(0), + self.peers, + self.me, + self.p_off, + self.l_off, + self.f_off, + _MAX_ROWS, + bmm1_scale, + bmm2_scale, + QL=q_len, + HQ=self.hq, + CP=self.cp, + D=HEAD_DIM, + PG=PAGE, + num_warps=8, + ) + return out + + def _exact_supported(self, q_len, b): + """The fused partial (q_len folded into the heads) takes <= MAX_BATCH verify / MAX_BATCH_Q1 draft requests; + wider batches take the trtllm-gen partial.""" + from sgl_nvidia_mega.qwen3_5.sm103.attention import q35_cp_fused as F + + return F.q35_cp_fused_supported( + 1, + self.cp, + q_len * self.hq, + b, + self.sms, + max_batch=F.MAX_BATCH if q_len > 1 else F.MAX_BATCH_Q1, + ) + + def _exact( + self, + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + out, + gate, + ): + """The fused partial on the context prefix (q_len folded into the heads: one query row of q_len x heads per + request, no causal effect), then the trtllm-gen path's merge with the causal block in fp32.""" + self._exact_partial( + k_cache, + v_cache, + page_table, + self.pre_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + ) + return self._exact_merge( + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + rows, + bmm1_scale, + bmm2_scale, + out, + gate, + ) + + def _gate_lens_fwd( + self, seq_lens, b, q_len, thr, forward_batch, layer_id, zero=False + ): + """Once per forward, make the stock and exact-prefix lengths. + + With ``zero``, the stock call gets no key while CP is selected. + """ + key = (seq_lens.data_ptr(), b, q_len, thr, zero) + ent = self._glens.get(key) + if ent is not None and ent[1]() is forward_batch and layer_id > ent[2]: + ent[2] = layer_id + return ent[0] + if len(self._glens) > 64: + self._glens.clear() + lens = torch.empty(2, b, dtype=torch.int32, device=self.device) + self._glens[key] = [(lens[0], lens[1]), weakref.ref(forward_batch), layer_id] + _gate_lens[(1,)]( + seq_lens, + b, + thr, + lens[0], + lens[1], + QL=q_len, + NB=triton.next_power_of_2(b), + GMAX=self.gate_max, + ZERO=zero, + ) + return lens[0], lens[1] + + def _gated_overlap( + self, + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + forward_batch, + layer_id, + stock, + gate, + exact, + ): + """Gated call: stock(lengths, out) on this stream beside gather + partial on the side stream (the + TRT-LLM partial with its own workspace), then the merge after the streams join.""" + if self.cp_main and exact and self.gate_overlap == 1: + return self._gated_main( + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + forward_batch, + layer_id, + stock, + gate, + ) + fold_mh = self.fold_mh and not exact + stock_lens, pre_lens = self._gate_lens_fwd( + seq_lens, + b, + q_len, + gate["thr"], + forward_batch, + layer_id, + zero=fold_mh, + ) + out = torch.empty( + rows, self.hq, HEAD_DIM, dtype=torch.bfloat16, device=self.device + ) + main = torch.cuda.current_stream() + self._side.wait_stream(main) + stock(stock_lens, out) + words = self.hq * HEAD_DIM // 4 + fused = None + with torch.cuda.stream(self._side): + if fold_mh and self.fold_fuse: + # Fill chunk-row queries in the gather and combine them in the merge. + tabs = self._tables( + page_table, + seq_lens, + q_len, + forward_batch, + layer_id, + gate, + ) + chunk = self._chunk_tables( + tabs[0], tabs[1], b, forward_batch, layer_id, q_len + ) + if chunk is not None: + nv, h = chunk[1].shape[0], q_len * self.cp * self.hq + fused = ( + chunk, + torch.empty(nv, h, HEAD_DIM, dtype=q.dtype, device=self.device), + torch.empty( + nv, + h, + HEAD_DIM, + dtype=torch.bfloat16, + device=self.device, + ), + torch.empty(nv, h, dtype=torch.float32, device=self.device), + tabs, + ) + _cp_gather_q_ll[(rows,)]( + q.view(torch.int32), + words, + self.peers, + self.me, + self.llq_off, + self.q_all.view(torch.int32), + self.cnt[0], + seq_lens, + pre_lens, + self.stock_lens, + **gate, + CP=self.cp, + WORDS=words, + QL=q_len, + FOLD=exact, + GWRITE=False, + num_warps=4, + **( + {} + if fused is None + else { + "qv": fused[1].view(torch.int32), + "vstart": fused[0][3], + "vcnt": fused[0][4], + "CHUNK": True, + } + ), + ) + if exact: + self._exact_partial( + k_cache, + v_cache, + page_table, + pre_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + ) + else: + self._fold_partial( + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + forward_batch, + layer_id, + gate, + self._own_workspace(), + fused, + ) + main.wait_stream(self._side) + merge = ( + self._exact_merge + if exact + else (self._fold_merge_h if fold_mh else self._fold_merge) + ) + merge_args = ( + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + rows, + bmm1_scale, + bmm2_scale, + out, + gate, + ) + if fused is None: + return merge(*merge_args) + return merge(*merge_args, fused=fused) + + def _gated_main( + self, + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + forward_batch, + layer_id, + stock, + gate, + ): + """Gated exact call (TRT-L18): gather, fused partial and merge (all exit with the gate off) on this stream. The + stock kernel would hold SMs the partial's grid barrier waits for, so it runs on the side stream after the + partial; it gets no key with the gate on (then writes nothing), so it and the merge write out in disjoint cases.""" + stock_lens, pre_lens = self._gate_lens_fwd( + seq_lens, b, q_len, gate["thr"], forward_batch, layer_id, zero=True + ) + out = torch.empty( + rows, self.hq, HEAD_DIM, dtype=torch.bfloat16, device=self.device + ) + words = self.hq * HEAD_DIM // 4 + _cp_gather_q_ll[(rows,)]( + q.view(torch.int32), + words, + self.peers, + self.me, + self.llq_off, + self.q_all.view(torch.int32), + self.cnt[0], + seq_lens, + pre_lens, + self.stock_lens, + **gate, + CP=self.cp, + WORDS=words, + QL=q_len, + FOLD=True, + GWRITE=False, + TRIG=self.cp_pdl, + num_warps=4, + ) + self._exact_partial( + k_cache, + v_cache, + page_table, + pre_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + pdl=self.cp_pdl, + ) + main = torch.cuda.current_stream() + self._side.wait_stream(main) + with torch.cuda.stream(self._side): + stock(stock_lens, out) + if self.merge_h: + self._merge_h( + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + rows, + bmm1_scale, + bmm2_scale, + out, + gate, + ) + else: + self._exact_merge( + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + rows, + bmm1_scale, + bmm2_scale, + out, + gate, + ) + main.wait_stream(self._side) + return out + + def _ungated_main( + self, + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + gate, + ): + """Run an ungated exact call with PDL and the per-head merge.""" + words = self.hq * HEAD_DIM // 4 + _cp_gather_q_ll[(rows,)]( + q.view(torch.int32), + words, + self.peers, + self.me, + self.llq_off, + self.q_all.view(torch.int32), + self.cnt[0], + seq_lens, + self.pre_lens, + self.stock_lens, + **gate, + CP=self.cp, + WORDS=words, + QL=q_len, + FOLD=True, + GWRITE=False, + TRIG=True, + WAIT=self.gather_pdl, + num_warps=4, + launch_pdl=self.gather_pdl, + ) + self._exact_partial( + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + pdl=True, + len_sub=q_len, + early_trigger=self.early_merge, + ) + out = torch.empty( + rows, self.hq, HEAD_DIM, dtype=torch.bfloat16, device=self.device + ) + epilogue = self._gate_epilogue + if epilogue is not None and not ( + tuple(epilogue[0].shape) == (rows, self.hq * HEAD_DIM) + and epilogue[0].stride(1) == 1 + and tuple(epilogue[2].shape) == (rows, self.hq * HEAD_DIM) + and epilogue[2].stride(1) == 1 + ): + epilogue = None + self._merge_h( + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + rows, + bmm1_scale, + bmm2_scale, + out, + {**gate, "thr": 0}, + pdl=True, + early=self.early_merge, + gate_epilogue=epilogue, + ) + if epilogue is not None: + self._gate_epilogue_done = True + if ("gate_epilogue", rows) not in self._logged: + self._logged.add(("gate_epilogue", rows)) + logger.info( + "TRT-L1 CP-lite merge writes the gated FP8 o_proj input at %d rows", + rows, + ) + return out + + def _merge_h( + self, + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + rows, + bmm1_scale, + bmm2_scale, + out, + gate, + pdl=None, + early=False, + gate_epilogue=None, + ): + pdl = self.cp_pdl if pdl is None else pdl + hq = self.hq + epilogue_args = ( + {} + if gate_epilogue is None + else { + "gate": gate_epilogue[0], + "gate_stride": gate_epilogue[0].stride(0), + "o_proj_input_scale": gate_epilogue[1], + "o_proj_input": gate_epilogue[2], + "o_proj_input_stride": gate_epilogue[2].stride(0), + "GATE_EPILOGUE": True, + } + ) + _cp_merge_h_ll[(rows, hq)]( + self.po_x.view(torch.int32), + self.pl_x.view(torch.int32), + hq * HEAD_DIM // 2, + self.po_x.stride(0) // 2, + hq, + self.pl_x.stride(0), + q, + q.stride(0), + k_cache, + v_cache, + k_cache.stride(0), + k_cache.stride(2), + page_table, + page_table.stride(0), + seq_lens, + out, + out.stride(0), + self.peers, + self.me, + self.llh_off, + self.cnt_h, + bmm1_scale, + bmm2_scale, + gate["nb"], + gate["thr"], + QL=q_len, + HQ=hq, + CP=self.cp, + D=HEAD_DIM, + PG=PAGE, + NWP=_LL_PWORDS, + NB=gate["NB"], + GMAX=gate["GMAX"], + WAIT=pdl, + EARLY=early and pdl, + num_warps=4, + launch_pdl=pdl, + **epilogue_args, + ) + + def _own_workspace(self): + """The overlapped trtllm-gen partial's workspace (the caller's belongs to the stock call).""" + if self._ws is None: + # the chunked fold partial's workspace grows with its rows (~0.23 MiB a row at q_len 7: 225 MiB at B128) + mb = 384 if self.fold_chunk else 128 + self._ws = torch.empty(mb << 20, dtype=torch.uint8, device=self.device) + return self._ws + + def _fold_partial( + self, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + forward_batch, + layer_id, + gate, + workspace, + fused=None, + ): + """trtllm-gen over this rank's owned prefix pages for every group head, the q_len rows folded into the + heads (q_len_per_req 1), into po / plse with LSE.""" + import flashinfer + + if fused is None: + tab, lens, max_kv = self._tables( + page_table, seq_lens, q_len, forward_batch, layer_id, gate + ) + else: + tab, lens, max_kv = fused[4] + h = q_len * self.cp * self.hq + query = self.q_all[:rows].view(b, h, HEAD_DIM) + out = self.po[:rows].view(b, h, HEAD_DIM) + lse = self.plse[:rows].view(b, h) + chunk = ( + self._chunk_tables(tab, lens, b, forward_batch, layer_id, q_len) + if fused is None + else fused[0] + ) + flashinfer.decode.trtllm_batch_decode_with_kv_cache( + query=query, + kv_cache=(k_cache, v_cache), + workspace_buffer=workspace, + block_tables=tab, + seq_lens=lens if chunk is None else chunk[0], + max_seq_len=max_kv, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + out=out, + lse=lse, + q_len_per_req=1, + multi_ctas_kv_counter_buffer=self.counter, + ) + if chunk is None: + return + # ragged batch: the call above got no key; the chunk rows (empty otherwise) and the combine + _, vtab, vlens, vstart, vcnt = chunk + nv, C = vtab.shape + n = h * HEAD_DIM // 4 + if fused is None: + qv = torch.empty(nv, h, HEAD_DIM, dtype=query.dtype, device=self.device) + po_v = torch.empty( + nv, h, HEAD_DIM, dtype=torch.bfloat16, device=self.device + ) + pl_v = torch.empty(nv, h, dtype=torch.float32, device=self.device) + _cp_chunk_q[(b, nv // b)]( + query.view(torch.int32), + qv.view(torch.int32), + vstart, + vcnt, + n, + BLK=1024, + ) + else: # the gather wrote the chunk rows' queries; the merge combines their partials + qv, po_v, pl_v = fused[1:4] + flashinfer.decode.trtllm_batch_decode_with_kv_cache( + query=qv, + kv_cache=(k_cache, v_cache), + workspace_buffer=self._own_workspace(), + block_tables=vtab, + seq_lens=vlens, + max_seq_len=C * PAGE, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + out=po_v, + lse=pl_v, + q_len_per_req=1, + multi_ctas_kv_counter_buffer=self.counter, + ) + if fused is None: + _cp_chunk_comb[(b, triton.cdiv(h, 16))]( + po_v, pl_v, vstart, vcnt, out, lse, h, HB=16, D=HEAD_DIM + ) + + def _chunk_tables(self, tab, lens, b, forward_batch, layer_id, q_len=7): + """Build the one-row and chunk-row tables once per forward. + + Draft calls use their own chunk policy. Ragged batches clear the one-row + lengths and use dense ``[B * K, C]`` chunk rows instead. + """ + spec = self.fold_chunk_q1 if q_len == 1 else self.fold_chunk + keys = [k for k in spec if k <= b] + C = spec[max(keys)] if keys else 0 + ncols = tab.shape[1] + if not C or ncols <= C or ncols % C: + return None + key = ("chunk", tab.data_ptr(), lens.data_ptr(), b) + ent = self._tabs.get(key) + if ent is not None and ent[1]() is forward_batch and layer_id > ent[2]: + ent[2] = layer_id + return ent[0] + K = ncols // C + nv = b * K + base_lens = torch.empty(b, dtype=torch.int32, device=self.device) + vtab = torch.empty(nv, C, dtype=torch.int32, device=self.device) + vlens = torch.empty(nv, dtype=torch.int32, device=self.device) + vstart = torch.empty(b, dtype=torch.int32, device=self.device) + vcnt = torch.empty(b, dtype=torch.int32, device=self.device) + _cp_chunk_tables[(b, K)]( + tab, + tab.stride(0), + lens, + base_lens, + vtab, + vlens, + vstart, + vcnt, + b, + nv, + float(self.fold_ragged), + C=C, + KP=K, + NB=triton.next_power_of_2(b), + PG=PAGE, + ) + out = (base_lens, vtab, vlens, vstart, vcnt) + self._tabs[key] = [out, weakref.ref(forward_batch), layer_id] + if (b, "chunk") not in self._logged: + self._logged.add((b, "chunk")) + logger.info( + "CP-lite fold partial: %d requests, chunk rows of %d pages (%d rows) when the longest " + "owned length x B >= %.2f x the total", + b, + C, + nv, + self.fold_ragged, + ) + return out + + def _fold_merge( + self, + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + rows, + bmm1_scale, + bmm2_scale, + out, + gate, + ): + gh_w = self.cp * self.hq * HEAD_DIM // 2 + _cp_a2a_merge_ll[(rows,)]( + self.po.view(torch.int32), + self.plse.view(torch.int32), + gh_w, + self.hq * HEAD_DIM // 2, + self.cp * self.hq, + self.hq, + q, + q.stride(0), + k_cache, + v_cache, + k_cache.stride(0), + k_cache.stride(2), + page_table, + page_table.stride(0), + seq_lens, + out, + out.stride(0), + self.peers, + self.me, + self.llp_off, + self.cnt[1], + bmm1_scale, + bmm2_scale, + **gate, + QL=q_len, + HQ=self.hq, + CP=self.cp, + D=HEAD_DIM, + PG=PAGE, + NWP=_LL_PWORDS, + num_warps=8, + ) + return out + + def _fold_merge_h( + self, + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + rows, + bmm1_scale, + bmm2_scale, + out, + gate, + fused=None, + ): + """Merge the fold layout with one program per row and query head.""" + hq = self.hq + _cp_merge_h_ll[(rows, hq)]( + self.po.view(torch.int32), + self.plse.view(torch.int32), + self.cp * hq * HEAD_DIM // 2, + hq * HEAD_DIM // 2, + self.cp * hq, + hq, + q, + q.stride(0), + k_cache, + v_cache, + k_cache.stride(0), + k_cache.stride(2), + page_table, + page_table.stride(0), + seq_lens, + out, + out.stride(0), + self.peers, + self.me, + self.llh_off, + self.cnt_h, + bmm1_scale, + bmm2_scale, + gate["nb"], + gate["thr"], + QL=q_len, + HQ=hq, + CP=self.cp, + D=HEAD_DIM, + PG=PAGE, + NWP=_LL_PWORDS, + NB=gate["NB"], + GMAX=gate["GMAX"], + WAIT=False, + num_warps=4, + **( + {} + if fused is None + else { + "po_v": fused[2].view(torch.int32), + "pl_v": fused[3], + "vstart": fused[0][3], + "vcnt": fused[0][4], + "CHUNK": True, + } + ), + ) + return out + + def _exact_partial( + self, + k_cache, + v_cache, + page_table, + pre_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + pdl=False, + len_sub=0, + early_trigger=False, + ): + from sgl_nvidia_mega.qwen3_5.sm103.attention import q35_cp_fused as F + + cp, hq = self.cp, self.hq + F.q35_cp_fused_partial( + self.q_all[:rows].view(b, cp, q_len * hq, HEAD_DIM).permute(1, 0, 2, 3), + k_cache, + page_table, + pre_lens, + 1, + self.me, + cp, + bmm1_scale, + bmm2_scale, + self.po_x[:, :rows].view(cp, b, q_len * hq, HEAD_DIM), + self.pl_x[:, :rows].view(cp, b, q_len * hq), + sm_count=self.sms, + enable_pdl=pdl, + v_pool=v_cache, + p_shift=_EXACT_P_SHIFT, + rescale_threshold=_EXACT_RESCALE_THRESHOLD, + len_sub=len_sub, + early_trigger=early_trigger, + ) + + def _exact_merge( + self, + q, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + rows, + bmm1_scale, + bmm2_scale, + out, + gate, + ): + hq = self.hq + if out is None: + out = torch.empty( + rows, hq, HEAD_DIM, dtype=torch.bfloat16, device=self.device + ) + _cp_a2a_merge_ll[(rows,)]( + self.po_x.view(torch.int32), + self.pl_x.view(torch.int32), + hq * HEAD_DIM // 2, + self.po_x.stride(0) // 2, + hq, + self.pl_x.stride(0), + q, + q.stride(0), + k_cache, + v_cache, + k_cache.stride(0), + k_cache.stride(2), + page_table, + page_table.stride(0), + seq_lens, + out, + out.stride(0), + self.peers, + self.me, + self.llp_off, + self.cnt[1], + bmm1_scale, + bmm2_scale, + **gate, + QL=q_len, + HQ=hq, + CP=self.cp, + D=HEAD_DIM, + PG=PAGE, + NWP=_LL_PWORDS, + num_warps=8, + ) + return out + + def _fused( + self, + k_cache, + v_cache, + page_table, + seq_lens, + q_len, + b, + rows, + bmm1_scale, + bmm2_scale, + ): + """The fused partial (TRT-LLM's kernel) on the gathered queries, then the 4-way merge; None where the kernel + does not take the shape (decided by shapes only).""" + from sgl_nvidia_mega.qwen3_5.sm103.attention import q35_cp_fused as F + + if not F.q35_cp_fused_supported(q_len, self.cp, self.hq, b, self.sms): + return None + cp, hq = self.cp, self.hq + F.q35_cp_fused_partial( + self.q_all[:rows].view(rows, cp, hq, HEAD_DIM).permute(1, 0, 2, 3), + k_cache, + page_table, + seq_lens, + q_len, + self.me, + cp, + bmm1_scale, + bmm2_scale, + self.po[:rows].view(rows, cp, hq, HEAD_DIM).permute(1, 0, 2, 3), + self.plse[:rows].view(rows, cp, hq).permute(1, 0, 2), + sm_count=self.sms, + enable_pdl=False, + v_pool=v_cache, + ) + out = torch.empty(rows, hq, HEAD_DIM, dtype=torch.bfloat16, device=self.device) + _cp_a2a_merge4_ll[(rows,)]( + self.po.view(torch.int32), + self.plse.view(torch.int32), + out, + out.stride(0), + self.peers, + self.me, + self.llp_off, + self.cnt[1], + HQ=hq, + CP=cp, + D=HEAD_DIM, + NWP=_LL_PWORDS, + num_warps=8, + ) + return out diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_partial_attn.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_partial_attn.py new file mode 100644 index 00000000..fd5ab58c --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_partial_attn.py @@ -0,0 +1,3893 @@ +# 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. +# CuTe constexpr branches must remain nested for staged compilation. +# ruff: noqa: SIM102 +# Adapted from TensorRT-LLM 4d09932ca7 + 7df45abe5c's wide / quad / two-slot plans (Apache-2.0): +# tensorrt_llm/_torch/cute_dsl_kernels/blackwell/q35_cp_partial_attn.py; plus a V pool, len_sub, numerics options. +"""Qwen3.5 context-parallel (CP-lite) partial decode attention, CuTe DSL (sm100 / GB300): one launch per call. + +Replaces the prep + trtllm-gen FMHA + LSE (+ own-head draft) launches of CP-lite's merge-free path. One CP rank of a KV +group (cp ranks share one KV head; every rank holds the whole KV cache) reads only the pages it owns -- page p belongs +to rank p % cp -- for every q row of the group: rows = q_len tokens x gh group heads (row n = t * gh + g, <= 128), and +writes one partial (O normalized bf16, base-2 LSE fp32) per (token, group head) into the layout alltoall_merge takes: +out_o [cp, b * q_len, hq, D], out_lse [cp, b * q_len, hq] with group head g = d * hq + h going to member d. The +contract is q35_cp.q35_cp_fused_partial_ref: row (r, t) sees key positions j < L_r - q_len + t + 1 (prefix + causal +draft block; the owner of each new token's page adds that part of the draft block); a rank that sees no key of a row +writes O = 0 and LSE = -inf. + +Unlike the single-rank kernel (q35_decode_attn.py: keys on the MMA M side, <= 32 rows on N), the q rows are the MMA M +dimension here (FA style, softmax thread = row = TMEM lane), which keeps row max / sums thread-local at 112-128 rows and +makes the partial epilogue row-contiguous: + +* grid (n_splits, batch), one CTA per SM, all co-resident. Request r's owned pages (owned index i -> global page + cp_rank + i * cp) form tiles of two pages (128 keys); CTA (s, r) takes tiles [s * ts, (s + 1) * ts). +* warp 0, TMA: the split's block-table entries -> SMEM, then per tile 4 boxes of K and 4 of V (64 keys x 128 B, SW128, + EVICT_FIRST) into a 2-stage ring. Pages >= max((L - q_len) // 64 - fresh_margin_pages, 0) are requested only after + griddepcontrol.wait (the new tokens, and the positions MTP passes wrote earlier in the step); earlier pages may be + requested before it (kv_pre_wait). Warp 3 loads q (128 x 256 FP8 box, rows >= q_len * gh zero-filled) after the wait. +* warp 1, MMA (tcgen05 kind::f8f6f4, FP32 accumulators in TMEM, 512 columns): GEMM1 S[128 rows, 128 keys] = Q . K^T into + one of two S buffers; GEMM2 O[128, 256] += P . V as two N = 128 MMAs (one per d half), issued one tile behind GEMM1. +* warps 4-7, softmax (thread = row): x = S * bmm1_log2, mask, row max; lazy rescale (the reference max of a row only + moves when a score exceeds it by more than 8 in log2 units, so P = exp2(x - m) <= 256 fits E4M3; a warp rescales its O + rows in TMEM when any of its rows moved); P -> SMEM as E4M3 (K-major SW128), row sums in FP32. Warps whose rows are + all padding (rows <= 96) skip the softmax; their P rows are zero. +* epilogue: O / l of every row -> FP16 -> the CTA's chunk-major workspace slots (staged through SMEM, coalesced), + LSE per (row, split); grid barrier per request (64-bit epoch counter, see below); CTA s then bulk-loads chunk s of the + request's row-major [rows][256] partials of all valid splits plus the LSEs of its rows, merges them, applies bmm2 and + writes bf16 O / fp32 LSE. Dependents are released after a CTA barrier that follows the last global store. + +PDL: the only global accesses before griddepcontrol.wait are reads of seq_lens, the block tables and prefix KV pages +(TMA); q, the scales, the epoch counter, the workspace and the outputs are touched after it. + +Grid-barrier assumptions (counter row ``counters[b]``, see q35_cp_fused.py), as in q35_decode_attn.py: +(a) launches on one counter row never overlap (every kernel between two launches is a plain launch or waits before it + exits); (b) no other kernel holds SMs while it waits for work ordered after this launch; (c) every launch on a row + uses the same n_splits. + +Wide split (``balanced=True`` internally): the grid is (n_ctas, 1) and the CTAs are shared among the requests by their +owned KV, so a ragged batch costs about what its total KV costs instead of what its longest request costs: + +* plan (in every CTA, warp 2, from seq_lens before the dependency wait; graph-safe, no extra launch): request r owns + t_r tiles on this rank; T = the smallest tile count with sum_r max(1, ceil(t_r / T)) <= n_ctas (binary search + between ceil(total / n_ctas) and min(max t_r, ceil(total / (n_ctas - b)))); request r gets c_r = max(1, ceil(t_r / + T)) consecutive CTAs, the longest request also the spare ones; CTA j of a request takes the j-th of c_r near-equal + blocks (floor / ceil) of its tiles, so the first min(c_r, t_r) CTAs have work. One request per lane up to 32 + requests; ``wide_plan`` (33..64 requests) gives each lane requests r and r + 32, ``quad_plan`` (65..128) requests + r, r + 32, r + 64 and r + 96, and both compute the same plan. Every request needs a CTA: batch < n_ctas. +* partials: a request with c_r CTAs keeps the equal split's layout and combine geometry with c_r splits + ([n_chunks][c_r][chunk_pad] FP16, LSEs [rows][lse_stride]) in the workspace slots of its CTAs. +* grid barrier per request: the request's first CTA adds n_ctas - c_r + 1 and every other one 1, so every launch adds + n_ctas to each counter of its row whatever the plan (assumption (c) holds with n_splits := n_ctas). +* combine: CTA j of the request finalizes chunk j (one bulk copy of the chunk of every valid partial, as above); the + chunk geometry follows c_r at run time, the merge loop is dispatched to a compile-time lanes-per-granule variant. + +Two-slot plan (``two_slot``, kernel_two, 65..128 requests): the plans above give every request at least one CTA, so past +~n_ctas / 2 uniform requests a CTA runs a whole request (64k contexts, b96: 128 tiles where 81 would balance). Here T = +the smallest tile count with sum_r max(t_r, T) <= n_ctas * T; request r takes the virtual tile range [V_r, V_r + +max(t_r, T)) (requests in order) and CTA i the range [i * T, (i + 1) * T), so a CTA computes at most T tiles, of at most +two requests (slot 0, then slot 1, through one K / V ring and one O accumulator). Each request keeps the layout above +with c_r = the CTAs meeting its tiles; its workspace slots start at the exclusive prefix sum of c_r (sum <= 2 n_ctas, +twice the slots); the grid barrier and the combine run once per slot. + +Packed two-slot plan (``two_slot_pack``): above, every request takes at least T virtual +tiles, the empty requests of a padded CUDA graph too, so ragged batches near n_ctas requests raise T (AgentX B113-128: +about 1.3x the total / n_ctas). Packed, the requests are placed in order at a running position V, p = V mod T: at a CTA +boundary (p = 0) request r takes max(t_r, 1) tiles, inside a CTA's range (p > 0) max(t_r, T - p), so a request that +would end inside the range it started in pads to the next boundary. At most one request starts inside each range: a +CTA still meets at most two (slots, c_r, workspace, barrier and combine as above; kernel_two unchanged); an empty +request is written (O = 0, LSE = -inf) by the CTA holding V_r. T must satisfy end <= n_ctas * T, which is not monotone +in T: T = the smallest feasible value in [lo, lo + 31] (lo = ceil(sum_r max(t_r, 1) / n_ctas)), one candidate per plan +lane; else the first feasible of 32 candidates lo + 32 + k s (the last one hi = the plan above's T, always feasible: +every packed share is <= max(t_r, T)), then the smallest feasible of the min(s - 1, 32) values below it, or it (host +mirror: q35_cp_fused._two_slot_pack_T). Each lane scans the requests for its candidate (max(t_r, 1) from SMEM, the +quotient by a scaled reciprocal off the dependency chain, one branch-free step per request, the start of every request +stored in the lane's SMEM row), so the winning lane's row gives the starts without another scan; hi is only searched +when [lo, lo + 31] has no feasible T. Tile counts must stay below 2^15. The packed and unpacked kernels are separate +builds. +""" + +from __future__ import annotations + +import cutlass +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Float32, Int32, Int64, Uint32, const_expr, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 +from cutlass.cutlass_dsl import T, dsl_user_op + +try: + from cutlass.utils import SmemAllocator +except ImportError: # pragma: no cover - older CuTe DSL + from cutlass.memory import SmemAllocator + +FP8 = cutlass.Float8E4M3FN +PAGE = 64 # tokens per KV page +HEAD_DIM = 256 +TILE_KEYS = 128 # keys per tile (two pages) = MMA N of GEMM1 and MMA K of GEMM2 +M_ROWS = 128 # MMA M: q rows (q_len x group heads), zero-padded +SW_BYTES = 128 # bytes per swizzle row (one TMA box column) +KV_STAGES = 2 +EVICT_FIRST = 0x12F0000000000000 # CUTLASS TMA::CacheHintSm90::EVICT_FIRST +RESCALE_THRESHOLD = 8.0 # log2 units: P <= 2^8 = 256 < 448 (E4M3 max) without a P shift +P_TOP = 8.75 # log2 of the largest P (2^8.75 = 430 < 448): a P shift s lowers the lazy-rescale threshold to 8.75 - s +NUM_THREADS = 256 # warp 0 TMA, warp 1 MMA, warp 2 scalars, warp 3 q, warps 4-7 softmax (thread = row) +SOFTMAX_WARP0 = 4 +TAB_MAX = 512 # block-table entries per split staged in SMEM (K and V each) +STAGE_ROW = ( + HEAD_DIM * 2 + 16 +) # bytes per row of the FP16 partial staging (16-byte pad: conflict-free v4 stores) +KV_REGION = ( + 4 * KV_STAGES * 16384 +) # bytes of the K + V rings; reused for the staging and the combine +MAX_COMB_ROWS = 32 # rows touched by one CTA's combine chunk (bound of its LSE buffer) +DBG_RING_CTAS = 256 # CTA stride of the in-model debug ring (>= SMs) +TMEM_COLS = 512 # 2 S buffers (2 x 128) + O (2 x 128) +COMB_LSE_OFF = ( + 96 * 1024 +) # balanced split: byte offset of the combine's LSE buffer in the K / V ring region +PLAN_WORDS = 8 # balanced split: per-CTA plan in SMEM (request, index, CTAs, first tile, tiles, first CTA, L, valid) +# packed two-slot plan: scratch at the start of the K ring (free until the first TMA load). max(t_r, 1) of the requests +# (+ the word the scan reads ahead), then two regions of one row per plan lane with each request's start; 129-word rows +# put the 32 lanes' stores for one request in different banks. +PACK_TT_BYTES = 1024 +PACK_ROW_BYTES = 129 * 4 +PACK_REGION_BYTES = 32 * PACK_ROW_BYTES +assert PACK_TT_BYTES + 2 * PACK_REGION_BYTES <= KV_STAGES * TILE_KEYS * HEAD_DIM + + +def cp_partial_geometry(rows: int, n_splits: int) -> dict: + """Combine geometry for ``rows`` q rows and ``n_splits`` CTAs per request (host and kernel share it). + + The request's partial of one split is row-major [rows][256] FP16; CTA s finalizes elements [s * chunk, (s + 1) * + chunk). Workspace per request: O partials chunk-major [n_chunks][n_splits][chunk_pad] FP16 (one contiguous block per + finalizing CTA), then LSEs [rows][lse_stride] FP32.""" + total = rows * HEAD_DIM + chunk = 8 * ((total + 8 * n_splits - 1) // (8 * n_splits)) + # an odd number of 16-byte granules per split row: lanes reading one granule of consecutive splits hit different + # SMEM bank groups + chunk_pad = chunk + (8 if (chunk // 8) % 2 == 0 else 0) + n_chunks = (total + chunk - 1) // chunk + lse_stride = 4 * ((n_splits + 3) // 4) + ws_o_elems = n_chunks * n_splits * chunk_pad + n_gran = chunk // 8 + gps = 32 + while gps > 1 and gps * n_gran > NUM_THREADS: + gps //= 2 + max_cov_rows = (chunk + HEAD_DIM - 2) // HEAD_DIM + 1 + comb_o_bytes = 16 * ((n_splits * chunk_pad * 2 + 15) // 16) + return { + "chunk": chunk, + "chunk_pad": chunk_pad, + "n_chunks": n_chunks, + "lse_stride": lse_stride, + "ws_o_elems": ws_o_elems, + "ws_req_bytes": ws_o_elems * 2 + rows * lse_stride * 4, + "n_gran": n_gran, + "gps": gps, + "max_cov_rows": max_cov_rows, + "comb_o_bytes": comb_o_bytes, + "comb_bytes": comb_o_bytes + MAX_COMB_ROWS * lse_stride * 4, + } + + +def cp_partial_balanced_geometry(rows: int, n_ctas: int, slots: int = 1) -> dict: + """Workspace and combine bounds of the balanced split for ``rows`` q rows and ``n_ctas`` CTAs (host and kernel). + + A request with c CTAs uses the equal split's combine geometry with c splits (cp_partial_geometry): its partials + [n_chunks(c)][c][chunk_pad(c)] FP16 occupy the O slots of its c consecutive CTAs (slot_o_bytes >= n_chunks(c) * + chunk_pad(c) * 2 for every c), its LSEs [rows][lse_stride(c)] FP32 the LSE slots (16 * rows bytes per CTA >= rows * + lse_stride(c) * 4 / c). The combine stages c x chunk_pad FP16 (<= COMB_LSE_OFF bytes) and nrows x lse_stride + FP32. ``slots`` = 2 (the two-slot plan, where a CTA works for up to two requests) gives every CTA two slots: the + sum of c is <= 2 x n_ctas there.""" + total = rows * HEAD_DIM + slot_o, max_o, max_l = 0, 0, 0 + for c in range(1, n_ctas + 1): + chunk = 8 * ((total + 8 * c - 1) // (8 * c)) + chunk_pad = chunk + (8 if (chunk // 8) % 2 == 0 else 0) + n_chunks = (total + chunk - 1) // chunk + lse_stride = 4 * ((c + 3) // 4) + nrows = min( + (chunk + HEAD_DIM - 2) // HEAD_DIM + 1, rows + ) # rows one chunk can touch + slot_o = max(slot_o, n_chunks * chunk_pad * 2) + max_o = max(max_o, c * chunk_pad * 2) + max_l = max(max_l, nrows * lse_stride * 4) + slot_o = 16 * ((slot_o + 15) // 16) + n_slots = slots * n_ctas + return { + "slot_o_bytes": slot_o, + "lse_area_off": n_slots * slot_o, + "lse_cta_bytes": 16 * rows, + "ws_bytes": n_slots * (slot_o + 16 * rows), + "max_comb_o_bytes": max_o, + "max_comb_lse_bytes": max_l, + } + + +# -------------------------------------------------------------------------------------------------------------------- +# inline PTX helpers +# -------------------------------------------------------------------------------------------------------------------- +def _asm(ret, args, text, cons, side_effects, loc, ip): + return llvm.inline_asm( + ret, + args, + text, + cons, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None) -> Float32: + return Float32( + _asm( + T.f32(), + [Float32(x).ir_value(loc=loc, ip=ip)], + "ex2.approx.ftz.f32 $0, $1;", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _lg2(x, *, loc=None, ip=None) -> Float32: + return Float32( + _asm( + T.f32(), + [Float32(x).ir_value(loc=loc, ip=ip)], + "lg2.approx.ftz.f32 $0, $1;", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma(a, b, c, *, loc=None, ip=None) -> Float32: + return Float32( + _asm( + T.f32(), + [Float32(v).ir_value(loc=loc, ip=ip) for v in (a, b, c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _pack_e4m3x4(a, b, c, d, *, loc=None, ip=None) -> Uint32: + """Four FP32 -> four E4M3 (cvt.rn.satfinite), ``a`` in the low byte.""" + out = _asm( + T.i32(), + [Float32(v).ir_value(loc=loc, ip=ip) for v 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", + False, + loc, + ip, + ) + return Uint32(out) + + +@dsl_user_op +def _e4m3x4_pair_sums(w, *, loc=None, ip=None): + """Four E4M3 (as packed by _pack_e4m3x4) -> FP32 (byte 0 + byte 1, byte 2 + byte 3).""" + out = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [Uint32(w).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 lo, hi, a0, a1, a2, a3;\n\t.reg .b32 h0, h1;\n\t.reg .f32 f0, f1, f2, f3;\n\t" + "mov.b32 {lo, hi}, $2;\n\tcvt.rn.f16x2.e4m3x2 h0, lo;\n\tcvt.rn.f16x2.e4m3x2 h1, hi;\n\t" + "mov.b32 {a0, a1}, h0;\n\tmov.b32 {a2, a3}, h1;\n\t" + "cvt.f32.f16 f0, a0;\n\tcvt.f32.f16 f1, a1;\n\tcvt.f32.f16 f2, a2;\n\tcvt.f32.f16 f3, a3;\n\t" + "add.f32 $0, f0, f1;\n\tadd.f32 $1, f2, f3;\n\t}\n", + "=f,=f,r", + False, + loc, + ip, + ) + return tuple( + Float32(llvm.extractvalue(T.f32(), out, [i], loc=loc, ip=ip)) for i in range(2) + ) + + +@dsl_user_op +def _pack_f16x2(lo, hi, *, loc=None, ip=None) -> Uint32: + """Two FP32 -> FP16x2 (round to nearest), ``lo`` in the low half.""" + return Uint32( + _asm( + T.i32(), + [ + Float32(hi).ir_value(loc=loc, ip=ip), + Float32(lo).ir_value(loc=loc, ip=ip), + ], + "cvt.rn.f16x2.f32 $0, $1, $2;", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _pack_bf16x2(lo, hi, *, loc=None, ip=None) -> Uint32: + """Two FP32 -> BF16x2 (round to nearest), ``lo`` in the low half.""" + return Uint32( + _asm( + T.i32(), + [ + Float32(hi).ir_value(loc=loc, ip=ip), + Float32(lo).ir_value(loc=loc, ip=ip), + ], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _unpack_f16x2(w, *, loc=None, ip=None): + """FP16x2 -> (low, high) FP32.""" + out = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [Uint32(w).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 l, h;\n\tmov.b32 {l, h}, $2;\n\tcvt.f32.f16 $0, l;\n\tcvt.f32.f16 $1, h;\n\t}\n", + "=f,=f,r", + False, + loc, + ip, + ) + return tuple( + Float32(llvm.extractvalue(T.f32(), out, [i], loc=loc, ip=ip)) for i in range(2) + ) + + +@dsl_user_op +def _keep_lt(x, j, cnt, *, loc=None, ip=None) -> Float32: + """x if j < cnt else -inf (one setp + selp; straight-line masking without data-dependent branches).""" + return Float32( + _asm( + T.f32(), + [ + Float32(x).ir_value(loc=loc, ip=ip), + Int32(j).ir_value(loc=loc, ip=ip), + Int32(cnt).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .pred p;\n\tsetp.lt.s32 p, $2, $3;\n\tselp.f32 $0, $1, 0fFF800000, p;\n\t}\n", + "=f,f,r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _vote_any(pred, *, loc=None, ip=None) -> Int32: + """Warp-wide OR of a predicate (all 32 lanes must execute it).""" + return Int32( + _asm( + T.i32(), + [Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred pi, po;\n\tsetp.ne.s32 pi, $1, 0;\n\t" + "vote.sync.any.pred po, pi, 0xffffffff;\n\tselp.s32 $0, 1, 0, po;\n\t}\n", + "=r,r", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_add_s32(x, *, loc=None, ip=None) -> Int32: + """Warp-wide sum (all 32 lanes must execute it).""" + return Int32( + _asm( + T.i32(), + [Int32(x).ir_value(loc=loc, ip=ip)], + "redux.sync.add.s32 $0, $1, 0xffffffff;", + "=r,r", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_s32(x, *, loc=None, ip=None) -> Int32: + """Warp-wide max (all 32 lanes must execute it).""" + return Int32( + _asm( + T.i32(), + [Int32(x).ir_value(loc=loc, ip=ip)], + "redux.sync.max.s32 $0, $1, 0xffffffff;", + "=r,r", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ballot_top(pred, *, loc=None, ip=None) -> Int32: + """Highest lane whose predicate is set (-1 if none; all 32 lanes must execute it).""" + return Int32( + _asm( + T.i32(), + [Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred pi;\n\t.reg .b32 bal;\n\tsetp.ne.s32 pi, $1, 0;\n\t" + "vote.sync.ballot.b32 bal, pi, 0xffffffff;\n\tbfind.u32 $0, bal;\n\t}\n", + "=r,r", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _shfl_up_s32(x, off, *, loc=None, ip=None) -> Int32: + """Value of lane - off (own value for lanes < off; all 32 lanes must execute it).""" + return Int32( + _asm( + T.i32(), + [Int32(x).ir_value(loc=loc, ip=ip), Int32(off).ir_value(loc=loc, ip=ip)], + "shfl.sync.up.b32 $0, $1, $2, 0, 0xffffffff;", + "=r,r,r", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bfind_u32(x, *, loc=None, ip=None) -> Int32: + """Index of the most significant set bit (-1 for 0).""" + return Int32( + _asm( + T.i32(), + [Int32(x).ir_value(loc=loc, ip=ip)], + "bfind.u32 $0, $1;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_floor_s32(a, b, *, loc=None, ip=None) -> Int32: + """floor(a / b) for 0 <= a < 2^24, 1 <= b < 2^24: FP32 reciprocal estimate (within one of the quotient there) + corrected by one step either way; a short dependency chain instead of the integer division sequence.""" + return Int32( + _asm( + T.i32(), + [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 fa, fb, fq;\n\t.reg .s32 q, p, r;\n\t.reg .pred hi, lo;\n\t" + "cvt.rn.f32.s32 fa, $1;\n\tcvt.rn.f32.s32 fb, $2;\n\trcp.approx.ftz.f32 fb, fb;\n\t" + "mul.rn.f32 fq, fa, fb;\n\tcvt.rzi.s32.f32 q, fq;\n\tmul.lo.s32 p, q, $2;\n\t" + "setp.gt.s32 hi, p, $1;\n\t@hi sub.s32 q, q, 1;\n\t@hi sub.s32 p, p, $2;\n\t" + "sub.s32 r, $1, p;\n\tsetp.ge.s32 lo, r, $2;\n\t@lo add.s32 q, q, 1;\n\tmov.b32 $0, q;\n\t}\n", + "=r,r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _rcp_up_f32(b, *, loc=None, ip=None) -> Float32: + """1 / b scaled up by 2^-16 (b >= 1): _div_rcp_s32(a, _rcp_up_f32(b)) = floor(a / b) for 0 <= a < 2^15 (the scale + exceeds the rcp.approx / mul rounding, and a * 2^-16 stays below the 1 / b margin to the next integer).""" + return Float32( + _asm( + T.f32(), + [Int32(b).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 f;\n\tcvt.rn.f32.s32 f, $1;\n\trcp.approx.ftz.f32 f, f;\n\t" + "mul.rn.f32 $0, f, 0f3F800080;\n\t}\n", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_rcp_s32(a, rb, *, loc=None, ip=None) -> Int32: + """floor(a * rb) for a >= 0: the quotient when rb = _rcp_up_f32(b) (three instructions, no correction step).""" + return Int32( + _asm( + T.i32(), + [Int32(a).ir_value(loc=loc, ip=ip), Float32(rb).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 f;\n\tcvt.rn.f32.s32 f, $1;\n\tmul.rn.f32 f, f, $2;\n\tcvt.rzi.s32.f32 $0, f;\n\t}\n", + "=r,r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _pack_step(p, k, q, m, b, *, loc=None, ip=None): + """One request of the packed placement (position k * b + p, 0 <= p < b; the request takes max(t, 1) = q * b + m + tiles): p + m wraps into the next range, or with p > 0 and q = 0 a short request pads to it. Returns (p', k') + without branches (the plan lanes scan different b).""" + out = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [Int32(v).ir_value(loc=loc, ip=ip) for v in (p, k, q, m, b)], + "{\n\t.reg .pred c, a;\n\t.reg .s32 u, w, x, y, z;\n\t" + "add.s32 u, $2, $5;\n\tsetp.ge.s32 c, u, $6;\n\tsetp.gt.s32 a, $2, 0;\n\tsetp.eq.and.s32 a, $4, 0, a;\n\t" + "sub.s32 w, u, $6;\n\tselp.s32 x, 0, u, a;\n\tor.pred a, a, c;\n\tselp.s32 y, 1, 0, a;\n\t" + "add.s32 z, $3, $4;\n\tadd.s32 $1, z, y;\n\tselp.s32 $0, w, x, c;\n\t}\n", + "=r,=r,r,r,r,r,r", + False, + loc, + ip, + ) + return tuple( + Int32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(2) + ) + + +@dsl_user_op +def _st_shared_v4(addr, a, b, c, d, *, loc=None, ip=None) -> None: + _asm( + None, + [Int32(addr).ir_value(loc=loc, ip=ip)] + + [Uint32(v).ir_value(loc=loc, ip=ip) for v in (a, b, c, d)], + "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_shared_v4_u32(addr, *, loc=None, ip=None): + out = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [Int32(addr).ir_value(loc=loc, ip=ip)], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + True, + loc, + ip, + ) + return tuple( + Uint32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(4) + ) + + +@dsl_user_op +def _st_shared_f32(addr, v, *, loc=None, ip=None) -> None: + _asm( + None, + [Int32(addr).ir_value(loc=loc, ip=ip), Float32(v).ir_value(loc=loc, ip=ip)], + "st.shared.f32 [$0], $1;", + "r,f", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_shared_f32(addr, *, loc=None, ip=None) -> Float32: + return Float32( + _asm( + T.f32(), + [Int32(addr).ir_value(loc=loc, ip=ip)], + "ld.shared.f32 $0, [$1];", + "=f,r", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_shared_s32(addr, v, *, loc=None, ip=None) -> None: + _asm( + None, + [Int32(addr).ir_value(loc=loc, ip=ip), Int32(v).ir_value(loc=loc, ip=ip)], + "st.shared.b32 [$0], $1;", + "r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_shared_s32(addr, *, loc=None, ip=None) -> Int32: + return Int32( + _asm( + T.i32(), + [Int32(addr).ir_value(loc=loc, ip=ip)], + "ld.shared.b32 $0, [$1];", + "=r,r", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_global_v4_u32(addr, a, b, c, d, *, loc=None, ip=None) -> None: + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip)] + + [Uint32(v).ir_value(loc=loc, ip=ip) for v in (a, b, c, d)], + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _st_global_u32(addr, a, *, loc=None, ip=None) -> None: + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(a).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _st_global_f32(addr, v, *, loc=None, ip=None) -> None: + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), Float32(v).ir_value(loc=loc, ip=ip)], + "st.global.f32 [$0], $1;", + "l,f", + True, + loc, + ip, + ) + + +@dsl_user_op +def _griddep_wait(*, loc=None, ip=None) -> None: + _asm(None, [], "griddepcontrol.wait;", "~{memory}", True, loc, ip) + + +@dsl_user_op +def _griddep_launch(*, loc=None, ip=None) -> None: + _asm(None, [], "griddepcontrol.launch_dependents;", "~{memory}", True, loc, ip) + + +@dsl_user_op +def _red_release_add_u64(addr, val, *, loc=None, ip=None) -> None: + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), Int64(val).ir_value(loc=loc, ip=ip)], + "red.release.gpu.global.add.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_acquire_gpu_u64(addr, *, loc=None, ip=None) -> Int64: + return Int64( + _asm( + T.i64(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.acquire.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ld_relaxed_gpu_u64(addr, *, loc=None, ip=None) -> Int64: + return Int64( + _asm( + T.i64(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bulk_g2s(dst_smem, src_gmem, nbytes, mbar_smem, *, loc=None, ip=None) -> None: + """cp.async.bulk global -> shared (16-byte aligned, size a multiple of 16), completion on an mbarrier.""" + _asm( + None, + [ + Int32(dst_smem).ir_value(loc=loc, ip=ip), + Int64(src_gmem).ir_value(loc=loc, ip=ip), + Int32(nbytes).ir_value(loc=loc, ip=ip), + Int32(mbar_smem).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1], $2, [$3];", + "r,l,r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _nanosleep(ns, *, loc=None, ip=None) -> None: + _asm( + None, + [Int32(ns).ir_value(loc=loc, ip=ip)], + "nanosleep.u32 $0;", + "r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None) -> Int32: + return Int32(_asm(T.i32(), [], "mov.u32 $0, %smid;", "=r", True, loc, ip)) + + +@dsl_user_op +def _globaltimer(*, loc=None, ip=None) -> Int64: + return Int64(_asm(T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", True, loc, ip)) + + +@dsl_user_op +def _clock64(*, loc=None, ip=None) -> Int64: + return Int64(_asm(T.i64(), [], "mov.u64 $0, %clock64;", "=l", True, loc, ip)) + + +# -------------------------------------------------------------------------------------------------------------------- +# kernel +# -------------------------------------------------------------------------------------------------------------------- +class Q35CpPartialAttnKernel: + """See the module docstring. Compile-time parameters: + + q_len, gh, hql, cp_size: tokens per request, group heads (= cp_size * hql), q heads per rank, CP group size; + rows = q_len * gh <= 128. + n_splits: CTAs per request (n_splits * batch <= SMs); with ``balanced`` the CTAs of the whole grid + (batch < n_splits <= SMs, shared among the requests by the in-kernel plan). + balanced: wide split (module docstring): grid (n_splits, 1); also exposed as a test hook for batches <= 32. + wide_plan: balanced split for 33..64 requests: the plan warp handles requests lane and lane + 32 (_plan_wide, + same plan as _plan would compute with 64 lanes); batches <= 32 keep the one-request-per-lane + _plan (the host compiles this variant only for batch > 32). + quad_plan: balanced split for 65..128 requests: lanes plan requests lane, lane + 32, lane + 64 and lane + 96 + (_plan_quad, same plan again); the host compiles it only for batch > 64, which keeps the 33..64 + calls on wide_plan. + two_slot: the two-slot plan (module docstring; kernel_two, _plan_two, up to 128 requests): the host's + default above 64 requests. + two_slot_pack: the packed two-slot plan (module docstring; _plan_pack). + bmm1_from_ptr: read the log2-domain QK scale from ``bmm1_ptr[1]`` (trtllm-gen [raw, raw * log2e] layout) after the + dependency wait, else use ``bmm1_log2``. Same for ``bmm2`` (``bmm2_ptr[0]``). + o_vec16: out_o rows are 16-byte aligned (8-element vector stores), else 4-byte stores. + use_pdl: launched as a programmatic dependent (waits, releases its dependents at the end). + early_trigger: with use_pdl, also release the dependents right after the dependency wait (warp 3), so a dependent + that waits before touching this kernel's outputs launches under it. + kv_pre_wait: request the prefix pages below the fresh margin before griddepcontrol.wait. + fresh_margin_pages: prefix pages just below the new tokens' first page that are read only after the wait. + debug_ts / unsafe_skip_wait: test hooks (%globaltimer stamps per CTA; no dependency waits = positive control of the + delayed-producer tests). + """ + + def __init__( + self, + *, + q_len: int, + gh: int, + hql: int, + cp_size: int, + n_splits: int, + bmm1_from_ptr: bool, + bmm2_from_ptr: bool, + o_vec16: bool, + use_pdl: bool = True, + kv_pre_wait: bool = True, + fresh_margin_pages: int = 1, + evict_first: bool = True, + debug_ts: bool = False, + unsafe_skip_wait: bool = False, + p_shift: float = 5.0, + debug_ring: int = 0, + balanced: bool = False, + rescale_threshold: float | None = None, + l_from_p8: bool = False, + len_sub: int = 0, + early_trigger: bool = False, + wide_plan: bool = False, + quad_plan: bool = False, + two_slot: bool = False, + two_slot_pack: bool = False, + ): + assert gh == cp_size * hql, "group heads = cp_size x q heads per rank" + self.early_trigger = bool(early_trigger) and use_pdl + # L = max(seq_lens - len_sub, 0): a folded caller's prefix lengths from the batch's own lengths, so nothing + # written by the q all-gather is read before the dependency wait + self.len_sub = int(len_sub) + self.balanced = bool(balanced) + self.wide_plan = bool(wide_plan) + self.quad_plan = bool(quad_plan) + self.two_slot = bool(two_slot) + self.two_slot_pack = bool(two_slot_pack) + assert self.two_slot or not self.two_slot_pack, ( + "the packed plan is a two-slot plan" + ) + assert self.balanced or not ( + self.wide_plan or self.quad_plan or self.two_slot + ), "balanced-split plans" + assert int(self.wide_plan) + int(self.quad_plan) + int(self.two_slot) <= 1, ( + "one plan variant" + ) + assert not (self.two_slot and debug_ring), ( + "the two-slot kernel has no debug ring" + ) + self.q_len = q_len + self.gh = gh + self.hql = hql + self.cp = cp_size + self.rows = q_len * gh + assert 1 <= self.rows <= M_ROWS, ( + f"q_len * group heads = {self.rows} must be in [1, {M_ROWS}]" + ) + self.act_warps = (self.rows + 31) // 32 # softmax warps that hold real rows + self.n_splits = n_splits + self.bmm1_from_ptr = bmm1_from_ptr + self.bmm2_from_ptr = bmm2_from_ptr + self.o_vec16 = o_vec16 + self.use_pdl = use_pdl + self.kv_pre_wait = kv_pre_wait and use_pdl + self.fresh_margin_pages = fresh_margin_pages + self.evict_first = evict_first + self.debug_ts = debug_ts + self.unsafe_skip_wait = unsafe_skip_wait + # in-model diagnostics: every launch writes one record per CTA into a ring of `debug_ring` launches + # [slot][cta][8] int64 = (entry, wait done, barrier arrive, barrier release, end (all %globaltimer), smid, + # launch index * 64 + batch, setup done), slot = launch index % debug_ring (launch index from the epoch; a + # batch of 64 or more reads as batch - 64 of a later launch index) + self.debug_ring = int(debug_ring) + assert not (self.debug_ring and debug_ts), ( + "debug_ring and debug_ts share the debug pointer" + ) + # P = exp2(x - m_ref + p_shift) before the E4M3 conversion (see q35_decode_attn.py): the reference max maps to + # 2^p_shift; l carries the same factor (the LSE subtracts it); the lazy rescale threshold is P_TOP - p_shift. + self.p_shift = float(p_shift) + if rescale_threshold is None: + rescale_threshold = min(RESCALE_THRESHOLD, P_TOP - self.p_shift) + assert 0.0 <= rescale_threshold <= P_TOP - self.p_shift + self.rescale_threshold = float(rescale_threshold) + # l_from_p8: the row sum adds the E4M3-rounded P that GEMM2 multiplies (numerator and denominator agree) + self.l_from_p8 = bool(l_from_p8) + assert 32 * self.act_warps * STAGE_ROW <= KV_REGION + if self.balanced: + geo_b = cp_partial_balanced_geometry( + self.rows, n_splits, 2 if self.two_slot else 1 + ) + self.ws_bytes = geo_b["ws_bytes"] + self.chunk = self.chunk_pad = self.lse_stride = ( + 0 # per request at run time (kernel: bchunk, ...) + ) + self.slot_o_bytes = geo_b["slot_o_bytes"] + self.lse_area_off = geo_b["lse_area_off"] + self.lse_cta_bytes = geo_b["lse_cta_bytes"] + assert geo_b["max_comb_o_bytes"] <= COMB_LSE_OFF, ( + "balanced combine: O partials exceed their buffer" + ) + assert geo_b["max_comb_lse_bytes"] <= KV_REGION - COMB_LSE_OFF, ( + "balanced combine: LSEs exceed their buffer" + ) + return + geo = cp_partial_geometry(self.rows, n_splits) + self.chunk = geo["chunk"] + self.chunk_pad = geo["chunk_pad"] + self.n_chunks = geo["n_chunks"] + self.lse_stride = geo["lse_stride"] + self.ws_o_elems = geo["ws_o_elems"] + self.ws_req_bytes = geo["ws_req_bytes"] + self.gps = geo["gps"] + self.comb_o_bytes = geo["comb_o_bytes"] + assert geo["max_cov_rows"] <= MAX_COMB_ROWS, ( + "too few splits for the combine (rows per chunk)" + ) + # granule passes of the combine: one when every granule of a chunk has its own gps lanes (<= 8-10 requests at + # 112 rows), more at larger batches (fewer splits, longer chunks; gps is then 1) + self.gran_iters = (geo["n_gran"] + NUM_THREADS // self.gps - 1) // ( + NUM_THREADS // self.gps + ) + assert geo["comb_bytes"] <= KV_REGION, "combine buffers exceed the K / V ring" + + # ---------------------------------------------------------------- host-side (JIT) entry + @cute.jit + def __call__( + self, + q_ptr: cute.Pointer, + kv_ptr: cute.Pointer, + v_ptr: cute.Pointer, + n_kv_rows: Int32, + tab_ptr: cute.Pointer, + tab_stride_b: Int32, + tab_stride_kv: Int32, + seqlen_ptr: cute.Pointer, + cp_rank: Int32, + bmm1_ptr: cute.Pointer, + bmm1_log2: Float32, + bmm2_ptr: cute.Pointer, + bmm2: Float32, + o_ptr: cute.Pointer, + o_sd: Int32, + o_st: Int32, + o_sh: Int32, + lse_ptr: cute.Pointer, + l_sd: Int32, + l_st: Int32, + l_sh: Int32, + ws_ptr: cute.Pointer, + cnt_ptr: cute.Pointer, + dbg_ptr: cute.Pointer, + batch: Int32, + stream, + ): + R = self.rows + NS = self.n_splits + # q of request r: rows [r * R, (r + 1) * R) of the all-gather buffer viewed as [b * R, 256] (row t * gh + g) + mQ = cute.make_tensor( + q_ptr, + cute.make_layout((R, HEAD_DIM, batch), stride=(HEAD_DIM, 1, R * HEAD_DIM)), + ) + mKV = cute.make_tensor( + kv_ptr, cute.make_layout((n_kv_rows, HEAD_DIM), stride=(HEAD_DIM, 1)) + ) + # V may live in its own pool (SGLang: per-layer K and V buffers); tab_stride_kv 0 shares one page table + mKVt = cute.make_tensor( + v_ptr, cute.make_layout((HEAD_DIM, n_kv_rows), stride=(1, HEAD_DIM)) + ) + mTab = cute.make_tensor( + tab_ptr, + cute.make_layout( + (batch, 2, 1 << 24), # page extent unused (no bounds) + stride=(tab_stride_b, tab_stride_kv, 1), + ), + ) + mSeq = cute.make_tensor(seqlen_ptr, cute.make_layout((batch,))) + mBmm1 = cute.make_tensor(bmm1_ptr, cute.make_layout((2,))) + mBmm2 = cute.make_tensor(bmm2_ptr, cute.make_layout((1,))) + # outputs / workspace are addressed explicitly (vector stores, bulk copies) + mO = cute.make_tensor(o_ptr, cute.make_layout((1,))) + mLse = cute.make_tensor(lse_ptr, cute.make_layout((1,))) + mWs = cute.make_tensor(ws_ptr, cute.make_layout((1,))) + mCnt = cute.make_tensor( + cnt_ptr, cute.make_layout((8,)) + ) # grid-barrier epoch per request + if const_expr(self.debug_ring > 0): + mDbg = cute.make_tensor( + dbg_ptr, + cute.make_layout( + (self.debug_ring, DBG_RING_CTAS, 8), + stride=(DBG_RING_CTAS * 8, 8, 1), + ), + ) + elif const_expr(self.two_slot): + mDbg = cute.make_tensor( + dbg_ptr, cute.make_layout((NS, 24), stride=(24, 1)) + ) # + slot 1's plan words + elif const_expr(self.balanced): + mDbg = cute.make_tensor(dbg_ptr, cute.make_layout((NS, 16), stride=(16, 1))) + else: + mDbg = cute.make_tensor( + dbg_ptr, cute.make_layout((batch * NS, 16), stride=(16, 1)) + ) + + mma_qk = sm100_utils.make_trivial_tiled_mma( + FP8, + FP8, + OperandMajorMode.K, + OperandMajorMode.K, + Float32, + tcgen05.CtaGroup.ONE, + (M_ROWS, TILE_KEYS), + ) + mma_pv = sm100_utils.make_trivial_tiled_mma( + FP8, + FP8, + OperandMajorMode.K, + OperandMajorMode.MN, + Float32, + tcgen05.CtaGroup.ONE, + (M_ROWS, 128), + ) + qk_tiler = (M_ROWS, TILE_KEYS, HEAD_DIM) + pv_tiler = (M_ROWS, 128, TILE_KEYS) # one d half per MMA + sQ_l = sm100_utils.make_smem_layout_a(mma_qk, qk_tiler, FP8, 1) + sK_l = sm100_utils.make_smem_layout_b(mma_qk, qk_tiler, FP8, KV_STAGES) + sP_l = sm100_utils.make_smem_layout_a(mma_pv, pv_tiler, FP8, 2) + # V: (d half h of stage s) = MMA stage 2 s + h; physically the [stage][d half][key atoms] layout of the TMA view + sV_l = sm100_utils.make_smem_layout_b(mma_pv, pv_tiler, FP8, 2 * KV_STAGES) + # the same buffers seen by the TMA (one box = one page x 128 B) and by the softmax P stores ((row, key, buffer)) + sK_tma_l = cute.tiled_divide( + sm100_utils.make_smem_layout( + OperandMajorMode.K, (TILE_KEYS, HEAD_DIM), FP8, KV_STAGES + ), + (PAGE, SW_BYTES), + ) + sV_tma_l = cute.tiled_divide( + sm100_utils.make_smem_layout( + OperandMajorMode.MN, (HEAD_DIM, TILE_KEYS), FP8, KV_STAGES + ), + (SW_BYTES, PAGE), + ) + sP_mk_l = sm100_utils.make_smem_layout( + OperandMajorMode.K, (M_ROWS, TILE_KEYS), FP8, 2 + ) + box_k = sm100_utils.make_smem_layout( + OperandMajorMode.K, (PAGE, SW_BYTES), FP8, 1 + ) + box_v = sm100_utils.make_smem_layout( + OperandMajorMode.MN, (SW_BYTES, PAGE), FP8, 1 + ) + + op = cpasync.CopyBulkTensorTileG2SOp() + tma_k, gK_tma = cpasync.make_tiled_tma_atom(op, mKV, box_k, (PAGE, SW_BYTES)) + tma_v, gV_tma = cpasync.make_tiled_tma_atom(op, mKVt, box_v, (SW_BYTES, PAGE)) + cluster_vmnk = cute.tiled_divide( + cute.make_layout((1, 1, 1)), (mma_qk.thr_id.shape,) + ) + sQ_one = cute.slice_(sQ_l, (None, None, None, 0)) + tma_q, gQ_tma = cute.nvgpu.make_tiled_tma_atom_A( + sm100_utils.cluster_shape_to_tma_atom_A((1, 1), mma_qk.thr_id), + mQ, + sQ_one, + qk_tiler, + mma_qk, + cluster_vmnk.shape, + ) + q_bytes = cute.size_in_bytes(FP8, sQ_one) + + if const_expr(self.balanced): + grid = [NS, 1, 1] + else: + grid = [NS, batch, 1] + if const_expr(self.two_slot): + self.kernel_two( + mma_qk, + mma_pv, + tma_k, + gK_tma, + tma_v, + gV_tma, + tma_q, + gQ_tma, + mTab, + mSeq, + cp_rank, + mBmm1, + bmm1_log2, + mBmm2, + bmm2, + mO, + o_sd, + o_st, + o_sh, + mLse, + l_sd, + l_st, + l_sh, + mWs, + mCnt, + mDbg, + sQ_l, + sK_l, + sV_l, + sP_l, + sK_tma_l, + sV_tma_l, + sP_mk_l, + q_bytes, + ).launch( + grid=grid, + block=[NUM_THREADS, 1, 1], + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=1, + ) + else: + self.kernel( + mma_qk, + mma_pv, + tma_k, + gK_tma, + tma_v, + gV_tma, + tma_q, + gQ_tma, + mTab, + mSeq, + cp_rank, + mBmm1, + bmm1_log2, + mBmm2, + bmm2, + mO, + o_sd, + o_st, + o_sh, + mLse, + l_sd, + l_st, + l_sh, + mWs, + mCnt, + mDbg, + sQ_l, + sK_l, + sV_l, + sP_l, + sK_tma_l, + sV_tma_l, + sP_mk_l, + q_bytes, + ).launch( + grid=grid, + block=[NUM_THREADS, 1, 1], + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=1, + ) + + # ---------------------------------------------------------------- device helpers + @cute.jit + def _produce( + self, + i, + kb0, + kb1, + vb0, + vb1, + tma_k: cute.CopyAtom, + tKg: cute.Tensor, + tKs: cute.Tensor, + tma_v: cute.CopyAtom, + tVg: cute.Tensor, + tVs: cute.Tensor, + full_k, + empty_k, + full_v, + empty_v, + ): + """TMA of tile i of the ring: pages (slot 0, slot 1) of K (blocks kb0, kb1) and V (vb0, vb1).""" + S: cutlass.Constexpr = KV_STAGES + s = i % S + ph = (i // S) & 1 + kbs = (kb0, kb1) + vbs = (vb0, vb1) + cute.arch.mbarrier_wait(empty_k + s, ph ^ 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(full_k + s, 2 * PAGE * HEAD_DIM) + for p in cutlass.range_constexpr(2): + for kb in cutlass.range_constexpr(2): + if const_expr(self.evict_first): + cute.copy( + tma_k, + tKg[(None, kbs[p], kb)], + tKs[(None, p, kb, s)], + tma_bar_ptr=full_k + s, + cache_policy=Int64(EVICT_FIRST), + ) + else: + cute.copy( + tma_k, + tKg[(None, kbs[p], kb)], + tKs[(None, p, kb, s)], + tma_bar_ptr=full_k + s, + ) + cute.arch.mbarrier_wait(empty_v + s, ph ^ 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(full_v + s, 2 * PAGE * HEAD_DIM) + for p in cutlass.range_constexpr(2): + for h in cutlass.range_constexpr(2): + if const_expr(self.evict_first): + cute.copy( + tma_v, + tVg[(None, h, vbs[p])], + tVs[(None, h, p, s)], + tma_bar_ptr=full_v + s, + cache_policy=Int64(EVICT_FIRST), + ) + else: + cute.copy( + tma_v, + tVg[(None, h, vbs[p])], + tVs[(None, h, p, s)], + tma_bar_ptr=full_v + s, + ) + + @cute.jit + def _gemm2( + self, + j, + mma_pv: cute.TiledMma, + tOtO: cute.Tensor, + tCrP: cute.Tensor, + tCrV: cute.Tensor, + full_v, + p_full, + empty_v, + p_empty, + o_done, + j_loc=None, + ): + """O += P(tile j) . V(tile j) (two d halves); releases the V stage and the P buffer, signals o_done. Returns + mma_pv: its ACCUMULATE field is updated here, and the caller must rebind it. j counts the CTA's tiles (ring and + barrier phases); j_loc (two-slot kernel) is the tile's index in its request: O restarts at j_loc 0.""" + S: cutlass.Constexpr = KV_STAGES + sj = j % S + pb = j % 2 + jl = j if j_loc is None else j_loc + cute.arch.mbarrier_wait(full_v + sj, (j // S) & 1) + cute.arch.mbarrier_wait(p_full + pb, (j // 2) & 1) + for h in cutlass.range_constexpr(2): + for kk in cutlass.range_constexpr(TILE_KEYS // 32): + mma_pv.set(tcgen05.Field.ACCUMULATE, (jl > 0) or (kk != 0)) + cute.gemm( + mma_pv, + tOtO[(None, 0, h)], + tCrP[(None, 0, kk, pb)], + tCrV[(None, 0, kk, 2 * sj + h)], + tOtO[(None, 0, h)], + ) + with cute.arch.elect_one(): + tcgen05.commit(empty_v + sj) + tcgen05.commit(p_empty + pb) + tcgen05.commit(o_done) + return mma_pv + + @cute.jit + def _softmax_tile( + self, + i, + xs, + ob64, + stt, + tStS: cute.Tensor, + tOtO: cute.Tensor, + t2r, + thr_ld, + t2r64, + thr_ld64, + r2t64, + thr_st64, + bars_sm, + sP_base, + sP_mk_l: cute.ComposedLayout, + st, + cnt_a, + cnt_b, + scale, + mDbg: cute.Tensor, + dbg_row, + apply_mask: cutlass.Constexpr, + i_loc=None, + ): + """Softmax of tile i (CTA-local) for row st (TMEM lane st): S -> x, mask (apply_mask: keys of slot 0 / 1 visible + iff key < cnt_a / cnt_b), lazy rescale of (m, l, O row), P -> SMEM as E4M3. stt = [m_ref, l_acc]. + The O rescale of a lazy-rescale tile runs after P(i) is in SMEM (by then GEMM2(i - 1) has normally completed, + and the S registers are free for 64-column TMEM round trips), before P(i) is released to GEMM2(i). + i_loc (two-slot kernel): the tile's index in its request (O holds nothing of the request before i_loc 1).""" + s_full, s_empty, p_full, p_empty, o_done = bars_sm + il = i if i_loc is None else i_loc + sb = i % 2 + cute.arch.mbarrier_wait(s_full + sb, (i // 2) & 1) + if const_expr(self.debug_ts): + if st == 0: + if i == 0: + mDbg[dbg_row, 2] = _globaltimer() + accS = tStS[(None, None), 0, 0, sb] + for c in cutlass.range_constexpr(4): + cute.copy( + t2r, + thr_ld.partition_S(cute.local_tile(accS, (M_ROWS, 32), (0, c))), + xs[c], + ) + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(s_empty + sb) + # ---- mask (tiles with the new tokens, the request tail or a duplicate second slot; straight-line) ---- + if const_expr(apply_mask): + for c in cutlass.range_constexpr(4): + for v in cutlass.range_constexpr(32): + if const_expr(32 * c + v < PAGE): + xs[c][v] = _keep_lt(xs[c][v], 32 * c + v, cnt_a) + else: + xs[c][v] = _keep_lt(xs[c][v], 32 * c + v - PAGE, cnt_b) + # ---- row max of the raw scores (scale > 0) ---- + m0 = xs[0][0] + m1 = xs[1][0] + m2 = xs[2][0] + m3 = xs[3][0] + for v in cutlass.range_constexpr(1, 32): + m0 = cute.arch.fmax(m0, xs[0][v]) + m1 = cute.arch.fmax(m1, xs[1][v]) + m2 = cute.arch.fmax(m2, xs[2][v]) + m3 = cute.arch.fmax(m3, xs[3][v]) + mx = cute.arch.fmax(cute.arch.fmax(m0, m1), cute.arch.fmax(m2, m3)) * scale + m_ref = stt[0] + need = Int32(0) + if mx > m_ref + self.rescale_threshold: + need = Int32(1) + resc = _vote_any(need) + alpha = Float32(1.0) + if resc != 0: + # slow path (warp-uniform): move the reference max of the rows that exceeded it and rescale l now; the O + # rows are rescaled below, once P(i) is written + m_new = m_ref + if need != 0: + m_new = mx + alpha = Float32(0.0) + if m_ref != -Float32.inf: + alpha = _ex2(m_ref - m_new) + stt[1] = stt[1] * alpha + stt[0] = m_new + # ---- P = exp2(s * scale - m), row sum (FP32), E4M3 P row -> SMEM (K-major SW128) ---- + m_use = stt[0] + if m_use == -Float32.inf: + m_use = Float32(0.0) # no visible key yet: every x is -inf, P = 0 + neg_m = self.p_shift - m_use + la = Float32(0.0) + lb = Float32(0.0) + lc = Float32(0.0) + ld = Float32(0.0) + pw = cute.make_rmem_tensor((32,), Uint32) + for c in cutlass.range_constexpr(4): + for q4 in cutlass.range_constexpr(8): + p0 = _ex2(_fma(xs[c][4 * q4], scale, neg_m)) + p1 = _ex2(_fma(xs[c][4 * q4 + 1], scale, neg_m)) + p2 = _ex2(_fma(xs[c][4 * q4 + 2], scale, neg_m)) + p3 = _ex2(_fma(xs[c][4 * q4 + 3], scale, neg_m)) + pw[8 * c + q4] = _pack_e4m3x4(p0, p1, p2, p3) + if const_expr(self.l_from_p8): + s01, s23 = _e4m3x4_pair_sums(pw[8 * c + q4]) + la = la + s01 + lb = lb + s23 + else: + la = la + p0 + lb = lb + p1 + lc = lc + p2 + ld = ld + p3 + stt[1] = stt[1] + ((la + lb) + (lc + ld)) + pb = i % 2 + cute.arch.mbarrier_wait(p_empty + pb, ((i // 2) & 1) ^ 1) + for c16 in cutlass.range_constexpr(TILE_KEYS // 16): + off = cute.crd2idx((st, 16 * c16, pb), sP_mk_l) + _st_shared_v4( + sP_base + off, + pw[4 * c16], + pw[4 * c16 + 1], + pw[4 * c16 + 2], + pw[4 * c16 + 3], + ) + if resc != 0: + if il > 0: + # O holds GEMM2 of tiles < i once GEMM2(i - 1) completed; GEMM2(i) waits for P(i) (released below). + cute.arch.mbarrier_wait(o_done, (i - 1) & 1) + for h in cutlass.range_constexpr(2): + accO = tOtO[(None, None), 0, h] + for c in cutlass.range_constexpr(2): + tO = cute.local_tile(accO, (M_ROWS, 64), (0, c)) + cute.copy(t2r64, thr_ld64.partition_S(tO), ob64) + cute.arch.fence_view_async_tmem_load() + for v in cutlass.range_constexpr(64): + ob64[v] = ob64[v] * alpha + cute.copy(r2t64, ob64, thr_st64.partition_D(tO)) + cute.arch.fence_view_async_tmem_store() + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.mbarrier_arrive(p_full + pb) + + def _out_coords(self, row, req): + """(member d, token, local head h) of q row ``row`` (= t * gh + g, g = d * hq + h) of request ``req`` + (trace-time helper, inlined).""" + g = row % self.gh + return g // self.hql, req * self.q_len + row // self.gh, g % self.hql + + @cute.jit + def _plan( + self, + mSeq: cute.Tensor, + cp_rank: Int32, + cta: Int32, + lane: Int32, + sPlan: cute.Tensor, + ): + """Balanced split plan, one warp (lane r = request r; batch <= 32): this CTA's (request, index in the request, + CTAs of the request, first tile, tiles, request's first CTA, L, CTAs with work) -> sPlan. Every CTA of the grid + derives the same plan from seq_lens (see the module docstring).""" + CP: cutlass.Constexpr = self.cp + NC: cutlass.Constexpr = self.n_splits + batch = mSeq.shape[0] + L = Int32(0) + if lane < batch: + L = mSeq[lane] + if const_expr(self.len_sub > 0): + L = cutlass.max(L - self.len_sub, 0) + # owned tiles of request `lane` on this rank (0 for lanes >= batch: L = 0) + t = ((((L + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + 1) // 2 + total = _redux_add_s32(t) + tmax = _redux_max_s32(t) + # smallest T (tiles per CTA) with sum max(1, ceil(t / T)) <= NC; ceil(total / (NC - batch)) and max t fit + # (tile counts stay far below 2^24: _div_floor_s32's range) + lo = cutlass.max((total + NC - 1) // NC, 1) + hi = cutlass.max( + cutlass.min(tmax, _div_floor_s32(total + NC - batch - 1, NC - batch)), lo + ) + while lo < hi: + mid = (lo + hi) // 2 + need = Int32(0) + if lane < batch: + need = cutlass.max(_div_floor_s32(t + mid - 1, mid), 1) + if _redux_add_s32(need) <= NC: + hi = mid + else: + lo = mid + 1 + c = Int32(0) + if lane < batch: + c = cutlass.max(_div_floor_s32(t + lo - 1, lo), 1) + spare = NC - _redux_add_s32(c) + top = Int32(0) + if lane < batch: + if t == tmax: + top = Int32(1) + if lane == _ballot_top(top): + c = c + spare # the spare CTAs go to (one of) the longest requests + incl = c + for k in cutlass.range_constexpr(5): + up = _shfl_up_s32(incl, 1 << k) + if lane >= (1 << k): + incl = incl + up + first = incl - c + if lane < batch: + if first <= cta: + if cta < incl: + j = cta - first + base = _div_floor_s32(t, c) + rem = t - base * c + n_t = base + if j < rem: + n_t = base + 1 + sPlan[0] = lane + sPlan[1] = j + sPlan[2] = c + sPlan[3] = j * base + cutlass.min(j, rem) + sPlan[4] = n_t + sPlan[5] = first + sPlan[6] = L + sPlan[7] = cutlass.min(c, t) + + @cute.jit + def _plan_wide( + self, + mSeq: cute.Tensor, + cp_rank: Int32, + cta: Int32, + lane: Int32, + sPlan: cute.Tensor, + ): + """_plan for up to 64 requests (wide_plan): lane r plans request r (lower half) and request r + 32 (upper + half). Same plan as _plan with 64 lanes and as the host mirror q35_cp_fused.balanced_split_plan: T by the same + binary search, the spare CTAs to the highest-numbered longest request, CTAs in request order.""" + CP: cutlass.Constexpr = self.cp + NC: cutlass.Constexpr = self.n_splits + batch = mSeq.shape[0] + lane_b = lane + 32 + La = Int32(0) + Lb = Int32(0) + if lane < batch: + La = mSeq[lane] + if const_expr(self.len_sub > 0): + La = cutlass.max(La - self.len_sub, 0) + if lane_b < batch: + Lb = mSeq[lane_b] + if const_expr(self.len_sub > 0): + Lb = cutlass.max(Lb - self.len_sub, 0) + # owned tiles on this rank of requests `lane` / `lane + 32` (0 past the batch: L = 0) + ta = ((((La + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + 1) // 2 + tb = ((((Lb + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + 1) // 2 + total = _redux_add_s32(ta + tb) + tmax = _redux_max_s32(cutlass.max(ta, tb)) + lo = cutlass.max((total + NC - 1) // NC, 1) + hi = cutlass.max( + cutlass.min(tmax, _div_floor_s32(total + NC - batch - 1, NC - batch)), lo + ) + while lo < hi: + mid = (lo + hi) // 2 + need = Int32(0) + if lane < batch: + need = cutlass.max(_div_floor_s32(ta + mid - 1, mid), 1) + if lane_b < batch: + need = need + cutlass.max(_div_floor_s32(tb + mid - 1, mid), 1) + if _redux_add_s32(need) <= NC: + hi = mid + else: + lo = mid + 1 + ca = Int32(0) + cb = Int32(0) + if lane < batch: + ca = cutlass.max(_div_floor_s32(ta + lo - 1, lo), 1) + if lane_b < batch: + cb = cutlass.max(_div_floor_s32(tb + lo - 1, lo), 1) + spare = NC - _redux_add_s32(ca + cb) + top_a = Int32(0) + top_b = Int32(0) + if lane < batch: + if ta == tmax: + top_a = Int32(1) + if lane_b < batch: + if tb == tmax: + top_b = Int32(1) + # the spare CTAs go to the highest-numbered longest request (an upper-half one when there is one) + hb = _ballot_top(top_b) + ha = _ballot_top(top_a) + if hb >= 0: + if lane == hb: + cb = cb + spare + else: + if lane == ha: + ca = ca + spare + # first CTA of each request: requests 0..31 (lower half) take CTAs first, then 32..63 + incl_a = ca + incl_b = cb + for k in cutlass.range_constexpr(5): + up_a = _shfl_up_s32(incl_a, 1 << k) + up_b = _shfl_up_s32(incl_b, 1 << k) + if lane >= (1 << k): + incl_a = incl_a + up_a + incl_b = incl_b + up_b + incl_b = incl_b + _redux_add_s32(ca) + first_a = incl_a - ca + first_b = incl_b - cb + if lane < batch: + if first_a <= cta: + if cta < incl_a: + self._plan_store(sPlan, lane, cta - first_a, ta, ca, first_a, La) + if lane_b < batch: + if first_b <= cta: + if cta < incl_b: + self._plan_store(sPlan, lane_b, cta - first_b, tb, cb, first_b, Lb) + + @cute.jit + def _plan_quad( + self, + mSeq: cute.Tensor, + cp_rank: Int32, + cta: Int32, + lane: Int32, + sPlan: cute.Tensor, + ): + """_plan for up to 128 requests (quad_plan): lane r plans requests r, r + 32, r + 64 and r + 96 (quarters a, b, + c, d). Same plan as _plan with 128 lanes and as the host mirror q35_cp_fused.balanced_split_plan: T by the same + binary search, the spare CTAs to the highest-numbered longest request, CTAs in request order.""" + CP: cutlass.Constexpr = self.cp + NC: cutlass.Constexpr = self.n_splits + batch = mSeq.shape[0] + lane_b = lane + 32 + lane_c = lane + 64 + lane_d = lane + 96 + La = Int32(0) + Lb = Int32(0) + Lc = Int32(0) + Ld = Int32(0) + if lane < batch: + La = mSeq[lane] + if const_expr(self.len_sub > 0): + La = cutlass.max(La - self.len_sub, 0) + if lane_b < batch: + Lb = mSeq[lane_b] + if const_expr(self.len_sub > 0): + Lb = cutlass.max(Lb - self.len_sub, 0) + if lane_c < batch: + Lc = mSeq[lane_c] + if const_expr(self.len_sub > 0): + Lc = cutlass.max(Lc - self.len_sub, 0) + if lane_d < batch: + Ld = mSeq[lane_d] + if const_expr(self.len_sub > 0): + Ld = cutlass.max(Ld - self.len_sub, 0) + # owned tiles on this rank of the lane's four requests (0 past the batch: L = 0) + ta = ((((La + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + 1) // 2 + tb = ((((Lb + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + 1) // 2 + tc = ((((Lc + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + 1) // 2 + td = ((((Ld + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + 1) // 2 + total = _redux_add_s32(ta + tb + tc + td) + tmax = _redux_max_s32(cutlass.max(cutlass.max(ta, tb), cutlass.max(tc, td))) + lo = cutlass.max((total + NC - 1) // NC, 1) + hi = cutlass.max( + cutlass.min(tmax, _div_floor_s32(total + NC - batch - 1, NC - batch)), lo + ) + while lo < hi: + mid = (lo + hi) // 2 + need = Int32(0) + if lane < batch: + need = cutlass.max(_div_floor_s32(ta + mid - 1, mid), 1) + if lane_b < batch: + need = need + cutlass.max(_div_floor_s32(tb + mid - 1, mid), 1) + if lane_c < batch: + need = need + cutlass.max(_div_floor_s32(tc + mid - 1, mid), 1) + if lane_d < batch: + need = need + cutlass.max(_div_floor_s32(td + mid - 1, mid), 1) + if _redux_add_s32(need) <= NC: + hi = mid + else: + lo = mid + 1 + ca = Int32(0) + cb = Int32(0) + cc = Int32(0) + cd = Int32(0) + if lane < batch: + ca = cutlass.max(_div_floor_s32(ta + lo - 1, lo), 1) + if lane_b < batch: + cb = cutlass.max(_div_floor_s32(tb + lo - 1, lo), 1) + if lane_c < batch: + cc = cutlass.max(_div_floor_s32(tc + lo - 1, lo), 1) + if lane_d < batch: + cd = cutlass.max(_div_floor_s32(td + lo - 1, lo), 1) + spare = NC - _redux_add_s32(ca + cb + cc + cd) + top_a = Int32(0) + top_b = Int32(0) + top_c = Int32(0) + top_d = Int32(0) + if lane < batch: + if ta == tmax: + top_a = Int32(1) + if lane_b < batch: + if tb == tmax: + top_b = Int32(1) + if lane_c < batch: + if tc == tmax: + top_c = Int32(1) + if lane_d < batch: + if td == tmax: + top_d = Int32(1) + # the spare CTAs go to the highest-numbered longest request (the highest quarter that has one); every lane + # takes part in the four ballots, the branches below are warp-uniform + hd = _ballot_top(top_d) + hc = _ballot_top(top_c) + hb = _ballot_top(top_b) + ha = _ballot_top(top_a) + if hd >= 0: + if lane == hd: + cd = cd + spare + else: + if hc >= 0: + if lane == hc: + cc = cc + spare + else: + if hb >= 0: + if lane == hb: + cb = cb + spare + else: + if lane == ha: + ca = ca + spare + # first CTA of each request: requests 0..31 take CTAs first, then 32..63, 64..95 and 96..127 + incl_a = ca + incl_b = cb + incl_c = cc + incl_d = cd + for k in cutlass.range_constexpr(5): + up_a = _shfl_up_s32(incl_a, 1 << k) + up_b = _shfl_up_s32(incl_b, 1 << k) + up_c = _shfl_up_s32(incl_c, 1 << k) + up_d = _shfl_up_s32(incl_d, 1 << k) + if lane >= (1 << k): + incl_a = incl_a + up_a + incl_b = incl_b + up_b + incl_c = incl_c + up_c + incl_d = incl_d + up_d + sum_a = _redux_add_s32(ca) + sum_b = _redux_add_s32(cb) + sum_c = _redux_add_s32(cc) + incl_b = incl_b + sum_a + incl_c = incl_c + sum_a + sum_b + incl_d = incl_d + sum_a + sum_b + sum_c + first_a = incl_a - ca + first_b = incl_b - cb + first_c = incl_c - cc + first_d = incl_d - cd + if lane < batch: + if first_a <= cta: + if cta < incl_a: + self._plan_store(sPlan, lane, cta - first_a, ta, ca, first_a, La) + if lane_b < batch: + if first_b <= cta: + if cta < incl_b: + self._plan_store(sPlan, lane_b, cta - first_b, tb, cb, first_b, Lb) + if lane_c < batch: + if first_c <= cta: + if cta < incl_c: + self._plan_store(sPlan, lane_c, cta - first_c, tc, cc, first_c, Lc) + if lane_d < batch: + if first_d <= cta: + if cta < incl_d: + self._plan_store(sPlan, lane_d, cta - first_d, td, cd, first_d, Ld) + + @cute.jit + def _plan_two( + self, + mSeq: cute.Tensor, + cp_rank: Int32, + cta: Int32, + lane: Int32, + sPlan: cute.Tensor, + scr: Int32, + ): + """Two-slot plan (two_slot, up to 128 requests; lanes plan requests r, r + 32, r + 64, r + 96). T = the smallest + tile count with sum_r max(t_r, T) <= NC * T. Request r occupies the virtual tile range [V_r, V_r + max(t_r, T)) + (requests in order) and CTA i the range [i * T, (i + 1) * T): a CTA meets at most two requests (every range is + at least T long), slot 0 = the request whose range holds the CTA's first position, slot 1 = the request that + starts inside it. The CTAs of request r are those meeting its real tiles [V_r, V_r + t_r) (the one holding V_r + when t_r = 0; it writes O = 0 / LSE = -inf). sPlan[8 k + w], slot k: (request or -1 for none, index among the + request's CTAs, its CTAs, first tile, tiles, first workspace slot, L, CTAs with work). two_slot_pack: T and + V_r of the packed plan (_plan_pack, scratch SMEM at scr). Returns T. Host mirror: + q35_cp_fused.balanced_split_plan_two.""" + CP: cutlass.Constexpr = self.cp + batch = mSeq.shape[0] + if lane < 2 * PLAN_WORDS: + sPlan[lane] = Int32(0) + if lane % PLAN_WORDS == 0: + sPlan[lane] = Int32( + -1 + ) # no request in the slot unless one is stored below + cute.arch.sync_warp() + lane_b = lane + 32 + lane_c = lane + 64 + lane_d = lane + 96 + La = Int32(0) + Lb = Int32(0) + Lc = Int32(0) + Ld = Int32(0) + if lane < batch: + La = mSeq[lane] + if const_expr(self.len_sub > 0): + La = cutlass.max(La - self.len_sub, 0) + if lane_b < batch: + Lb = mSeq[lane_b] + if const_expr(self.len_sub > 0): + Lb = cutlass.max(Lb - self.len_sub, 0) + if lane_c < batch: + Lc = mSeq[lane_c] + if const_expr(self.len_sub > 0): + Lc = cutlass.max(Lc - self.len_sub, 0) + if lane_d < batch: + Ld = mSeq[lane_d] + if const_expr(self.len_sub > 0): + Ld = cutlass.max(Ld - self.len_sub, 0) + ta = ((((La + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + 1) // 2 + tb = ((((Lb + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + 1) // 2 + tc = ((((Lc + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + 1) // 2 + td = ((((Ld + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + 1) // 2 + total = _redux_add_s32(ta + tb + tc + td) + tmax = _redux_max_s32(cutlass.max(cutlass.max(ta, tb), cutlass.max(tc, td))) + if const_expr(self.two_slot_pack): + T, Va, Vb, Vc, Vd = self._plan_pack( + batch, lane, ta, tb, tc, td, total, tmax, scr + ) + else: + T = self._plan_two_T(batch, lane, ta, tb, tc, td, total, tmax) + # virtual lengths, then their exclusive prefix sums V in request order (quarters a, b, c, d) + va = Int32(0) + vb = Int32(0) + vc = Int32(0) + vd = Int32(0) + if lane < batch: + va = cutlass.max(ta, T) + if lane_b < batch: + vb = cutlass.max(tb, T) + if lane_c < batch: + vc = cutlass.max(tc, T) + if lane_d < batch: + vd = cutlass.max(td, T) + iva = va + ivb = vb + ivc = vc + ivd = vd + for k in cutlass.range_constexpr(5): + ua = _shfl_up_s32(iva, 1 << k) + ub = _shfl_up_s32(ivb, 1 << k) + uc = _shfl_up_s32(ivc, 1 << k) + ud = _shfl_up_s32(ivd, 1 << k) + if lane >= (1 << k): + iva = iva + ua + ivb = ivb + ub + ivc = ivc + uc + ivd = ivd + ud + sva = _redux_add_s32(va) + svb = _redux_add_s32(vb) + svc = _redux_add_s32(vc) + Va = iva - va + Vb = ivb + sva - vb + Vc = ivc + sva + svb - vc + Vd = ivd + sva + svb + svc - vd + # first / last CTA meeting the real tiles (one position when t = 0) and the request's CTA count + fa = _div_floor_s32(Va, T) + fb = _div_floor_s32(Vb, T) + fc = _div_floor_s32(Vc, T) + fd = _div_floor_s32(Vd, T) + ca = Int32(0) + cb = Int32(0) + cc = Int32(0) + cd = Int32(0) + if lane < batch: + ca = _div_floor_s32(Va + cutlass.max(ta, 1) - 1, T) - fa + 1 + if lane_b < batch: + cb = _div_floor_s32(Vb + cutlass.max(tb, 1) - 1, T) - fb + 1 + if lane_c < batch: + cc = _div_floor_s32(Vc + cutlass.max(tc, 1) - 1, T) - fc + 1 + if lane_d < batch: + cd = _div_floor_s32(Vd + cutlass.max(td, 1) - 1, T) - fd + 1 + # first workspace slot of each request: exclusive prefix sums of the CTA counts + ica = ca + icb = cb + icc = cc + icd = cd + for k in cutlass.range_constexpr(5): + ua = _shfl_up_s32(ica, 1 << k) + ub = _shfl_up_s32(icb, 1 << k) + uc = _shfl_up_s32(icc, 1 << k) + ud = _shfl_up_s32(icd, 1 << k) + if lane >= (1 << k): + ica = ica + ua + icb = icb + ub + icc = icc + uc + icd = icd + ud + sca = _redux_add_s32(ca) + scb = _redux_add_s32(cb) + scc = _redux_add_s32(cc) + Sa = ica - ca + Sb = icb + sca - cb + Sc = icc + sca + scb - cc + Sd = icd + sca + scb + scc - cd + if lane < batch: + self._plan_two_store(sPlan, cta, T, lane, Va, ta, fa, ca, Sa, La) + if lane_b < batch: + self._plan_two_store(sPlan, cta, T, lane_b, Vb, tb, fb, cb, Sb, Lb) + if lane_c < batch: + self._plan_two_store(sPlan, cta, T, lane_c, Vc, tc, fc, cc, Sc, Lc) + if lane_d < batch: + self._plan_two_store(sPlan, cta, T, lane_d, Vd, td, fd, cd, Sd, Ld) + return T + + @cute.jit + def _plan_two_T( + self, + batch: Int32, + lane: Int32, + ta: Int32, + tb: Int32, + tc: Int32, + td: Int32, + total: Int32, + tmax: Int32, + ): + """The two-slot plan's T (_plan_two): binary search over the warp's four requests per lane.""" + NC: cutlass.Constexpr = self.n_splits + lane_b = lane + 32 + lane_c = lane + 64 + lane_d = lane + 96 + # smallest T with sum max(t_r, T) <= NC * T (the left side grows by at most batch < NC per step of T) + lo = cutlass.max((total + NC - 1) // NC, 1) + hi = cutlass.max(tmax, lo) + while lo < hi: + mid = (lo + hi) // 2 + need = Int32(0) + if lane < batch: + need = cutlass.max(ta, mid) + if lane_b < batch: + need = need + cutlass.max(tb, mid) + if lane_c < batch: + need = need + cutlass.max(tc, mid) + if lane_d < batch: + need = need + cutlass.max(td, mid) + if _redux_add_s32(need) <= NC * mid: + hi = mid + else: + lo = mid + 1 + return lo + + @cute.jit + def _plan_pack( + self, + batch: Int32, + lane: Int32, + ta: Int32, + tb: Int32, + tc: Int32, + td: Int32, + total: Int32, + tmax: Int32, + scr: Int32, + ): + """Packed two-slot plan (two_slot_pack, module docstring): (T, V_r of requests lane, lane + 32, lane + 64, lane + + 96; 0 past the batch). Round 0 scans T = lo + k, round 1 (only when round 0 found none) lo + 32 + k s up to hi, + round 2 the m values below round 1's pick; lane l scans k = 31 - l, so the highest feasible lane holds the + smallest feasible T. The rounds share one scan site: the plan's code precedes the hot loops and costs time even + unexecuted (three inlined scans in the unpacked build: +8-9 us per call at B68-128, 120K). SMEM scratch at scr + (PACK_* layout). Host mirror: q35_cp_fused._two_slot_pack_T / _two_slot_pack_starts.""" + NC: cutlass.Constexpr = self.n_splits + lane_b = lane + 32 + lane_c = lane + 64 + lane_d = lane + 96 + tta = Int32(0) + ttb = Int32(0) + ttc = Int32(0) + ttd = Int32(0) + if lane < batch: + tta = cutlass.max(ta, 1) + if lane_b < batch: + ttb = cutlass.max(tb, 1) + if lane_c < batch: + ttc = cutlass.max(tc, 1) + if lane_d < batch: + ttd = cutlass.max(td, 1) + _st_shared_s32(scr + 4 * lane, tta) + _st_shared_s32(scr + 4 * lane_b, ttb) + _st_shared_s32(scr + 4 * lane_c, ttc) + _st_shared_s32(scr + 4 * lane_d, ttd) + cute.arch.sync_warp() + lo = cutlass.max((_redux_add_s32(tta + ttb + ttc + ttd) + NC - 1) // NC, 1) + k = 31 - lane + rows_a = scr + PACK_TT_BYTES + rows_b = rows_a + PACK_REGION_BYTES + lo2 = lo + 32 + T = lo + row = rows_a + c1 = Int32(0) + s = Int32(1) + m = Int32(0) + hi = Int32(0) + rnd = Int32(0) + while rnd < 3: + if rnd == 1: + # nothing feasible in [lo, lo + 31] (then hi > lo + 31): strided candidates up to hi (k = 0 is hi) + hi = self._plan_two_T(batch, lane, ta, tb, tc, td, total, tmax) + s = cutlass.max((hi - lo2 + 30) // 31, 1) + cand = lo + k + base = rows_a + act = Int32(1) + if rnd == 1: + cand = cutlass.min(lo2 + k * s, hi) + base = rows_b + if rnd == 2: + # the m values below round 1's pick (lanes k >= m rescan T - 1, not counted) + cand = c1 - m + cutlass.min(k, m - 1) + if k >= m: + act = Int32(0) + end = self._pack_scan(batch, cand, scr, base + lane * PACK_ROW_BYTES) + ok = Int32(0) + if end <= NC * cand: + ok = act + win = _ballot_top(ok) + if rnd == 0: + rnd = Int32(1) + if win >= 0: + T = lo + 31 - win + row = rows_a + win * PACK_ROW_BYTES + rnd = Int32(3) + else: + if rnd == 1: + c1 = cutlass.min(lo2 + (31 - win) * s, hi) + T = c1 + row = rows_b + win * PACK_ROW_BYTES + m = cutlass.min(s - 1, 32) + rnd = Int32(3) + if m > 0: + rnd = Int32(2) + else: + if win >= 0: + T = c1 - m + 31 - win + row = rows_a + win * PACK_ROW_BYTES + rnd = Int32(3) + cute.arch.sync_warp() + Va = Int32(0) + Vb = Int32(0) + Vc = Int32(0) + Vd = Int32(0) + if lane < batch: + Va = _ld_shared_s32(row + 4 * lane) + if lane_b < batch: + Vb = _ld_shared_s32(row + 4 * lane_b) + if lane_c < batch: + Vc = _ld_shared_s32(row + 4 * lane_c) + if lane_d < batch: + Vd = _ld_shared_s32(row + 4 * lane_d) + # the K ring's TMA loads (async proxy) follow the CTA barrier after the plan + cute.arch.fence_proxy("async.shared", space="cta") + return T, Va, Vb, Vc, Vd + + @cute.jit + def _pack_scan(self, batch: Int32, T: Int32, tt: Int32, row: Int32): + """Packed placement for T over the batch's max(t_r, 1) at SMEM tt: stores each request's start at row + 4 r and + returns the end position (k T + p). One request per iteration, its max(t, 1) loaded one iteration ahead (the tt + area holds the word past the batch); the quotient is off the _pack_step dependency chain.""" + rb = _rcp_up_f32(T) + k = Int32(0) + p = Int32(0) + x = _ld_shared_s32(tt) + for r in cutlass.range(batch): + xr = x + x = _ld_shared_s32(tt + 4 * r + 4) + q = _div_rcp_s32(xr, rb) + _st_shared_s32(row + 4 * r, k * T + p) + p, k = _pack_step(p, k, q, xr - q * T, T) + return k * T + p + + @cute.jit + def _plan_two_store( + self, + sPlan: cute.Tensor, + cta: Int32, + T: Int32, + r: Int32, + V: Int32, + t: Int32, + f: Int32, + c: Int32, + s0: Int32, + L: Int32, + ): + """This CTA's slot for request r (if the CTA meets its real tiles): see _plan_two.""" + if f <= cta: + if cta < f + c: + lo_t = cutlass.max(V, cta * T) - V + hi_t = cutlass.min(V + cutlass.max(t, 1), (cta + 1) * T) - V + n_t = hi_t - lo_t + n_valid = c + if t == 0: + lo_t = Int32(0) + n_t = Int32(0) + n_valid = Int32(0) + base = Int32(8) + if V <= cta * T: + base = Int32(0) + sPlan[base] = r + sPlan[base + 1] = cta - f + sPlan[base + 2] = c + sPlan[base + 3] = lo_t + sPlan[base + 4] = n_t + sPlan[base + 5] = s0 + sPlan[base + 6] = L + sPlan[base + 7] = n_valid + + @cute.jit + def _plan_store( + self, + sPlan: cute.Tensor, + r: Int32, + j: Int32, + t: Int32, + c: Int32, + first: Int32, + L: Int32, + ): + """This CTA's plan words (as in _plan): CTA j of request r's c CTAs takes the j-th of c near-equal blocks of + its t tiles.""" + base = _div_floor_s32(t, c) + rem = t - base * c + n_t = base + if j < rem: + n_t = base + 1 + sPlan[0] = r + sPlan[1] = j + sPlan[2] = c + sPlan[3] = j * base + cutlass.min(j, rem) + sPlan[4] = n_t + sPlan[5] = first + sPlan[6] = L + sPlan[7] = cutlass.min(c, t) + + @cute.jit + def _combine_balanced( + self, + tidx: Int32, + warp: Int32, + lane: Int32, + req: Int32, + split: Int32, + c_req: Int32, + n_valid: Int32, + chunk: Int32, + chunk_pad: Int32, + lse_stride: Int32, + ws_req, + ws_lse, + kv_base, + comb_bar, + o_base, + o_sd: Int32, + o_st: Int32, + o_sh: Int32, + l_base, + l_sd: Int32, + l_st: Int32, + l_sh: Int32, + sScal: cute.Tensor, + mDbg: cute.Tensor, + dbg_row: Int32, + comb_phase=0, + ): + """Balanced split: CTA ``split`` of the request's ``c_req`` CTAs finalizes elements [o_lo, o_hi) of the request: + the equal split's combine with the request's chunk geometry (run time) and its lanes per granule dispatched to + a compile-time variant (_merge_balanced). comb_phase: comb_bar's phase (the two-slot kernel's second combine + after a first one that loaded partials uses 1).""" + TOT: cutlass.Constexpr = self.rows * HEAD_DIM + o_lo = split * chunk + if o_lo < TOT: + o_hi = cutlass.min(o_lo + chunk, TOT) + row0 = o_lo // HEAD_DIM + nrows = (o_hi - 1) // HEAD_DIM - row0 + 1 + n_gran = (o_hi - o_lo) // 8 + s_comb = kv_base # [n_valid][chunk_pad] FP16 + s_lse = kv_base + COMB_LSE_OFF # [nrows][lse_stride] FP32 + if n_valid > 0: + if warp == 0: + cute.arch.fence_proxy( + "async.global" + ) # acquire (generic proxy) -> bulk reads (async proxy) + if lane == 0: + nb_o = n_valid * (chunk_pad * 2) + nb_l = nrows * (lse_stride * 4) + cute.arch.mbarrier_arrive_and_expect_tx(comb_bar, nb_o + nb_l) + _bulk_g2s( + s_comb, + ws_req + (split * c_req * chunk_pad) * 2, + nb_o, + comb_bar.toint(), + ) + _bulk_g2s( + s_lse, + ws_lse + (row0 * lse_stride) * 4, + nb_l, + comb_bar.toint(), + ) + # spin (no suspend hint): every lane of a warp must be awake for the shuffles below + landed = cute.arch.mbarrier_try_wait(comb_bar, comb_phase) + while not landed: + landed = cute.arch.mbarrier_try_wait(comb_bar, comb_phase) + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 6] = _globaltimer() + # lanes per granule: the largest power of two <= 32 with gps * (chunk // 8) <= NUM_THREADS + lg_gps = cutlass.min( + _bfind_u32(cutlass.max(_div_floor_s32(NUM_THREADS, chunk // 8), 1)), 5 + ) + args = ( + tidx, + req, + o_lo, + row0, + n_gran, + n_valid, + chunk_pad, + lse_stride, + s_comb, + s_lse, + o_base, + o_sd, + o_st, + o_sh, + l_base, + l_sd, + l_st, + l_sh, + sScal, + ) + if lg_gps == 0: + self._merge_balanced(1, *args) + else: + if lg_gps == 1: + self._merge_balanced(2, *args) + else: + if lg_gps == 2: + self._merge_balanced(4, *args) + else: + if lg_gps == 3: + self._merge_balanced(8, *args) + else: + if lg_gps == 4: + self._merge_balanced(16, *args) + else: + self._merge_balanced(32, *args) + + @cute.jit + def _merge_balanced( + self, + lanes: cutlass.Constexpr, + tidx: Int32, + req: Int32, + o_lo: Int32, + row0: Int32, + n_gran: Int32, + n_valid: Int32, + chunk_pad: Int32, + lse_stride: Int32, + s_comb, + s_lse, + o_base, + o_sd: Int32, + o_st: Int32, + o_sh: Int32, + l_base, + l_sd: Int32, + l_st: Int32, + l_sh: Int32, + sScal: cute.Tensor, + ): + """Granule merge of the balanced combine for ``lanes`` (compile time) lanes per granule: the equal split's + loop body.""" + gps: cutlass.Constexpr = ( + lanes # annotated: stays a Python constant inside the dynamic loops below + ) + n_bfly: cutlass.Constexpr = lanes.bit_length() - 1 + per_pass: cutlass.Constexpr = NUM_THREADS // lanes + jg = tidx % gps + for gi in cutlass.range((n_gran + per_pass - 1) // per_pass): + g = gi * per_pass + tidx // gps + gc = cutlass.min(g, n_gran - 1) + e0 = o_lo + gc * 8 + orow = e0 // HEAD_DIM + dd, tok, hh = self._out_coords(orow, req) + addr = o_base + (dd * o_sd + tok * o_st + hh * o_sh + e0 % HEAD_DIM) * 2 + if n_valid > 0: + lrow = s_lse + (orow - row0) * (lse_stride * 4) + mrow = -Float32.inf + for j in cutlass.range(jg, n_valid, gps): + mrow = cute.arch.fmax(mrow, _ld_shared_f32(lrow + j * 4)) + for lg in cutlass.range_constexpr(n_bfly): + mrow = cute.arch.fmax( + mrow, cute.arch.shuffle_sync_bfly(mrow, offset=gps >> (lg + 1)) + ) + m_use = mrow + if mrow == -Float32.inf: + m_use = Float32( + 0.0 + ) # no visible key in any split: every weight is exp2(-inf) = 0 + den = Float32(0.0) + a0 = Float32(0.0) + a1 = Float32(0.0) + a2 = Float32(0.0) + a3 = Float32(0.0) + a4 = Float32(0.0) + a5 = Float32(0.0) + a6 = Float32(0.0) + a7 = Float32(0.0) + for j in cutlass.range(jg, n_valid, gps): + wj = _ex2(_ld_shared_f32(lrow + j * 4) - m_use) + den = den + wj + u4 = _ld_shared_v4_u32(s_comb + (j * chunk_pad + gc * 8) * 2) + f0, f1 = _unpack_f16x2(u4[0]) + f2, f3 = _unpack_f16x2(u4[1]) + f4, f5 = _unpack_f16x2(u4[2]) + f6, f7 = _unpack_f16x2(u4[3]) + a0 = _fma(wj, f0, a0) + a1 = _fma(wj, f1, a1) + a2 = _fma(wj, f2, a2) + a3 = _fma(wj, f3, a3) + a4 = _fma(wj, f4, a4) + a5 = _fma(wj, f5, a5) + a6 = _fma(wj, f6, a6) + a7 = _fma(wj, f7, a7) + for lg in cutlass.range_constexpr(n_bfly): + sh_off = gps >> (lg + 1) + den = den + cute.arch.shuffle_sync_bfly(den, offset=sh_off) + a0 = a0 + cute.arch.shuffle_sync_bfly(a0, offset=sh_off) + a1 = a1 + cute.arch.shuffle_sync_bfly(a1, offset=sh_off) + a2 = a2 + cute.arch.shuffle_sync_bfly(a2, offset=sh_off) + a3 = a3 + cute.arch.shuffle_sync_bfly(a3, offset=sh_off) + a4 = a4 + cute.arch.shuffle_sync_bfly(a4, offset=sh_off) + a5 = a5 + cute.arch.shuffle_sync_bfly(a5, offset=sh_off) + a6 = a6 + cute.arch.shuffle_sync_bfly(a6, offset=sh_off) + a7 = a7 + cute.arch.shuffle_sync_bfly(a7, offset=sh_off) + if jg == 0: + if g < n_gran: + scl = Float32( + 0.0 + ) # a row without any visible key gives 0, not NaN + if den > Float32(0.0): + scl = sScal[0] / den + w0 = _pack_bf16x2(a0 * scl, a1 * scl) + w1 = _pack_bf16x2(a2 * scl, a3 * scl) + w2 = _pack_bf16x2(a4 * scl, a5 * scl) + w3 = _pack_bf16x2(a6 * scl, a7 * scl) + if const_expr(self.o_vec16): + _st_global_v4_u32(addr, w0, w1, w2, w3) + else: + _st_global_u32(addr, w0) + _st_global_u32(addr + 4, w1) + _st_global_u32(addr + 8, w2) + _st_global_u32(addr + 12, w3) + if e0 % HEAD_DIM == 0: + lse_out = -Float32.inf + if den > Float32(0.0): + lse_out = mrow + _lg2(den) + _st_global_f32( + l_base + (dd * l_sd + tok * l_st + hh * l_sh) * 4, + lse_out, + ) + else: + # this rank sees no key of the request: O = 0, LSE = -inf + if jg == 0: + if g < n_gran: + if const_expr(self.o_vec16): + _st_global_v4_u32( + addr, Uint32(0), Uint32(0), Uint32(0), Uint32(0) + ) + else: + _st_global_u32(addr, Uint32(0)) + _st_global_u32(addr + 4, Uint32(0)) + _st_global_u32(addr + 8, Uint32(0)) + _st_global_u32(addr + 12, Uint32(0)) + if e0 % HEAD_DIM == 0: + _st_global_f32( + l_base + (dd * l_sd + tok * l_st + hh * l_sh) * 4, + -Float32.inf, + ) + + # ---------------------------------------------------------------- device kernel + @cute.kernel + def kernel( + self, + mma_qk: cute.TiledMma, + mma_pv: cute.TiledMma, + tma_k: cute.CopyAtom, + gK_tma: cute.Tensor, + tma_v: cute.CopyAtom, + gV_tma: cute.Tensor, + tma_q: cute.CopyAtom, + gQ_tma: cute.Tensor, + mTab: cute.Tensor, + mSeq: cute.Tensor, + cp_rank: Int32, + mBmm1: cute.Tensor, + bmm1_log2: Float32, + mBmm2: cute.Tensor, + bmm2: Float32, + mO: cute.Tensor, + o_sd: Int32, + o_st: Int32, + o_sh: Int32, + mLse: cute.Tensor, + l_sd: Int32, + l_st: Int32, + l_sh: Int32, + mWs: cute.Tensor, + mCnt: cute.Tensor, + mDbg: cute.Tensor, + sQ_l: cute.ComposedLayout, + sK_l: cute.ComposedLayout, + sV_l: cute.ComposedLayout, + sP_l: cute.ComposedLayout, + sK_tma_l: cute.ComposedLayout, + sV_tma_l: cute.ComposedLayout, + sP_mk_l: cute.ComposedLayout, + q_bytes: cutlass.Constexpr, + ): + S: cutlass.Constexpr = KV_STAGES + R: cutlass.Constexpr = self.rows + QL: cutlass.Constexpr = self.q_len + GH: cutlass.Constexpr = self.gh + CP: cutlass.Constexpr = self.cp + NS: cutlass.Constexpr = self.n_splits + AW: cutlass.Constexpr = self.act_warps + BAL: cutlass.Constexpr = self.balanced + chunk: cutlass.Constexpr = self.chunk + chunk_pad: cutlass.Constexpr = self.chunk_pad + lse_stride: cutlass.Constexpr = self.lse_stride + tidx, _, _ = cute.arch.thread_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + split, req, _ = cute.arch.block_idx() + cta = split # balanced: the CTA index; its request / index in the request come from the plan + if const_expr(BAL): + dbg_row = cta + else: + dbg_row = req * NS + split + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 0] = _globaltimer() + t_entry = Int64(0) + if const_expr(self.debug_ring > 0): + t_entry = _globaltimer() + + # ---- split geometry: seq_lens (and below, the block tables) are host-uploaded or written before the q + # all-gather (non-PDL torch ops / kernels that completed before it); read before the dependency wait ---- + if const_expr(not BAL): + L = mSeq[req] + if const_expr(self.len_sub > 0): + L = cutlass.max(L - self.len_sub, 0) + n_pages = (L + PAGE - 1) // PAGE + n_own = ( + n_pages - cp_rank + CP - 1 + ) // CP # pages p < n_pages with p % CP == cp_rank (>= 0) + n_tiles = ( + (n_own + 1) // 2 + ) # tile k = owned pages 2k, 2k + 1 (the last one may hold a single page) + # blocked: CTA `split` takes tiles [k_lo, k_hi) (fewer, fully busy CTAs finish a ragged split count sooner + # than a round robin that spreads the last round over every SM) + ts = cutlass.max((n_tiles + NS - 1) // NS, 1) + n_valid = ( + n_tiles + ts - 1 + ) // ts # splits with work (the combine reads only these) + k_lo = cutlass.min(split * ts, n_tiles) + k_hi = cutlass.min(k_lo + ts, n_tiles) + n_t = k_hi - k_lo + prefix = L - QL # keys < prefix are visible to every row + # pages >= fresh_page are read only after griddepcontrol.wait (see fresh_margin_pages) + fresh_page = cutlass.max(prefix // PAGE - self.fresh_margin_pages, 0) + ws_base = ( + mWs.iterator.toint() + req * self.ws_req_bytes + ) # [n_chunks][NS][chunk_pad] FP16 partials + ws_lse = ws_base + self.ws_o_elems * 2 # [R][lse_stride] FP32 + o_base = mO.iterator.toint() + l_base = mLse.iterator.toint() + + # ---- shared memory (K ring then V ring first: [kv_base, kv_base + KV_REGION) is one contiguous region) ---- + smem = SmemAllocator() + sK = smem.allocate_tensor(FP8, sK_l.outer, 1024, sK_l.inner) + sV = smem.allocate_tensor(FP8, sV_l.outer, 1024, sV_l.inner) + sQ = smem.allocate_tensor(FP8, sQ_l.outer, 1024, sQ_l.inner) + sP = smem.allocate_tensor(FP8, sP_l.outer, 1024, sP_l.inner) + bars = smem.allocate_array(Int64, 4 * S + 12) + full_k = bars + empty_k = bars + S + full_v = bars + 2 * S + empty_v = bars + 3 * S + q_full = bars + 4 * S + s_full = bars + 4 * S + 1 # [2] + s_empty = bars + 4 * S + 3 # [2] + p_full = bars + 4 * S + 5 # [2] + p_empty = bars + 4 * S + 7 # [2] + o_done = bars + 4 * S + 9 + o_full = bars + 4 * S + 10 + comb_bar = bars + 4 * S + 11 + tmem_holder = smem.allocate_array(Int32, 4) + sTab = smem.allocate_tensor( + Int32, cute.make_layout((2, TAB_MAX), stride=(TAB_MAX, 1)), 16 + ) + sTarget = smem.allocate_tensor(Int64, cute.make_layout((2,)), 16) + sScal = smem.allocate_tensor(Float32, cute.make_layout((4,)), 16) + sDbgT = smem.allocate_tensor( + Int64, cute.make_layout((2,)), 16 + ) # debug ring: wait-done stamp + sK_tma = cute.make_tensor( + cute.recast_ptr(sK.iterator, sK_tma_l.inner, dtype=FP8), sK_tma_l.outer + ) + sV_tma = cute.make_tensor( + cute.recast_ptr(sV.iterator, sV_tma_l.inner, dtype=FP8), sV_tma_l.outer + ) + kv_base = sK.iterator.toint() + sP_base = sP.iterator.toint() + if const_expr(BAL): + # this CTA's share of the batch (warp 2, published by the CTA barrier after the TMEM allocation) + sPlan = smem.allocate_tensor(Int32, cute.make_layout((PLAN_WORDS,)), 16) + if warp == 2: + if const_expr(self.quad_plan): + self._plan_quad(mSeq, cp_rank, cta, lane, sPlan) + elif const_expr(self.wide_plan): + self._plan_wide(mSeq, cp_rank, cta, lane, sPlan) + else: + self._plan(mSeq, cp_rank, cta, lane, sPlan) + + if warp == 0: + cpasync.prefetch_descriptor(tma_k) + cpasync.prefetch_descriptor(tma_v) + cpasync.prefetch_descriptor(tma_q) + if tidx == 0: + for i in cutlass.range_constexpr(4 * S + 1): + cute.arch.mbarrier_init(bars + i, 1) + for i in cutlass.range_constexpr(2): + cute.arch.mbarrier_init(s_full + i, 1) + cute.arch.mbarrier_init(s_empty + i, 32 * AW) + cute.arch.mbarrier_init(p_full + i, 32 * AW) + cute.arch.mbarrier_init(p_empty + i, 1) + cute.arch.mbarrier_init(o_done, 1) + cute.arch.mbarrier_init(o_full, 1) + cute.arch.mbarrier_init(comb_bar, 1) + cute.arch.mbarrier_init_fence() + if const_expr(AW < 4): + # P rows of the softmax warps without real rows: zero once (rows 8a..8a+7 = bytes [1024a, 1024a + 1024) + # of a 16 KB buffer, any swizzle) + n_zero: cutlass.Constexpr = ( + (M_ROWS - 32 * AW) * TILE_KEYS // 16 + ) # 16-byte words per buffer + for z in cutlass.range(tidx, 2 * n_zero, NUM_THREADS): + zb = z // n_zero + zw = z - zb * n_zero + _st_shared_v4( + sP_base + zb * (M_ROWS * TILE_KEYS) + 32 * AW * TILE_KEYS + zw * 16, + Uint32(0), + Uint32(0), + Uint32(0), + Uint32(0), + ) + cute.arch.fence_proxy("async.shared", space="cta") + if warp == 1: + cute.arch.alloc_tmem(TMEM_COLS, tmem_holder) + cute.arch.sync_threads() + if const_expr(BAL): + req = cute.arch.make_warp_uniform(sPlan[0]) + split = cute.arch.make_warp_uniform( + sPlan[1] + ) # index of this CTA among its request's CTAs + c_req = cute.arch.make_warp_uniform(sPlan[2]) # CTAs of the request + k_lo = cute.arch.make_warp_uniform(sPlan[3]) + n_t = cute.arch.make_warp_uniform(sPlan[4]) + first = cute.arch.make_warp_uniform( + sPlan[5] + ) # the request's first CTA (= its first workspace slot) + L = cute.arch.make_warp_uniform(sPlan[6]) + n_valid = cute.arch.make_warp_uniform( + sPlan[7] + ) # CTAs with work: the first n_valid of the request + n_pages = (L + PAGE - 1) // PAGE + n_own = (n_pages - cp_rank + CP - 1) // CP + prefix = L - QL + fresh_page = cutlass.max(prefix // PAGE - self.fresh_margin_pages, 0) + # the request's combine geometry (cp_partial_geometry with c_req CTAs) and workspace: partials + # [n_chunks][c_req][chunk_pad] FP16 in the slots of its CTAs, LSEs [R][lse_stride] FP32 + bchunk = 8 * _div_floor_s32(R * HEAD_DIM + 8 * c_req - 1, 8 * c_req) + bchunk_g = bchunk // 8 + bchunk_pad = bchunk + if bchunk_g % 2 == 0: + bchunk_pad = ( + bchunk + 8 + ) # odd granules per row: consecutive partials of one granule in other banks + blse_stride = 4 * ((c_req + 3) // 4) + ws_req = mWs.iterator.toint() + first * self.slot_o_bytes + ws_lse = ( + mWs.iterator.toint() + self.lse_area_off + first * self.lse_cta_bytes + ) + if const_expr(self.debug_ts): + if tidx == 0: + # the plan, for the tests: ((request * 256 + index) * 256 + CTAs) * 256 + first CTA; first tile; + # tiles + plan_id = (Int64(req) * Int64(256) + Int64(split)) * Int64( + 256 + ) + Int64(c_req) + mDbg[dbg_row, 13] = plan_id * Int64(256) + Int64(first) + mDbg[dbg_row, 14] = Int64(k_lo) + mDbg[dbg_row, 15] = Int64(n_t) + t_setup = Int64(0) + if const_expr(self.debug_ring > 0): + t_setup = _globaltimer() # TMEM allocated, barriers initialized + tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, tmem_holder) + tStS = mma_qk.make_fragment_C(mma_qk.partition_shape_C((M_ROWS, TILE_KEYS, 2))) + tStS = cute.make_tensor(tmem_ptr, tStS.layout) + s_cols = tcgen05.find_tmem_tensor_col_offset(tStS) + tOtO = mma_pv.make_fragment_C(mma_pv.partition_shape_C((M_ROWS, HEAD_DIM))) + tOtO = cute.make_tensor(tmem_ptr + s_cols, tOtO.layout) + + if warp == 0: + # ======================= TMA producer ======================= + # the split's block-table entries -> SMEM: sTab[kv][j] = owned page 2 k_lo + j (a duplicate second slot + # re-reads the last owned page); balanced: TAB_MAX entries at a time (refilled below) + if const_expr(BAL): + n_stage = cutlass.min(2 * n_t, TAB_MAX) + else: + n_stage = 2 * n_t + for j in cutlass.range(lane, n_stage, 32): + oi = cutlass.min(2 * k_lo + j, n_own - 1) + pg = cp_rank + oi * CP + sTab[0, j] = mTab[req, 0, pg] + sTab[1, j] = mTab[req, 1, pg] + cute.arch.sync_warp() + if const_expr(self.debug_ts): + if lane == 0: + mDbg[dbg_row, 1] = _globaltimer() + gK = cute.tiled_divide(gK_tma, (PAGE, SW_BYTES)) # ((64, 128), blocks, 2) + gV = cute.tiled_divide(gV_tma, (SW_BYTES, PAGE)) # ((128, 64), 2, blocks) + tKs, tKg = cpasync.tma_partition(tma_k, 0, cute.make_layout(1), sK_tma, gK) + tVs, tVg = cpasync.tma_partition(tma_v, 0, cute.make_layout(1), sV_tma, gV) + waited = Int32(0) + if const_expr(not self.kv_pre_wait): + if const_expr(self.use_pdl and not self.unsafe_skip_wait): + _griddep_wait() + waited = Int32(1) + for i in cutlass.range(n_t): + if waited == 0: + if ( + cp_rank + cutlass.min(2 * (k_lo + i) + 1, n_own - 1) * CP + >= fresh_page + ): + # this tile may hold positions written earlier in the step (QKV append, MTP passes) + if const_expr(self.use_pdl and not self.unsafe_skip_wait): + _griddep_wait() + waited = Int32(1) + if const_expr(BAL): + ti = i % (TAB_MAX // 2) + if ti == 0: + if i > 0: + # next TAB_MAX table entries (tiles beyond TAB_MAX / 2 per CTA: > ~490k-token contexts) + cute.arch.sync_warp() + for j in cutlass.range( + lane, cutlass.min(2 * (n_t - i), TAB_MAX), 32 + ): + oi = cutlass.min(2 * (k_lo + i) + j, n_own - 1) + pg = cp_rank + oi * CP + sTab[0, j] = mTab[req, 0, pg] + sTab[1, j] = mTab[req, 1, pg] + cute.arch.sync_warp() + self._produce( + i, + sTab[0, 2 * ti], + sTab[0, 2 * ti + 1], + sTab[1, 2 * ti], + sTab[1, 2 * ti + 1], + tma_k, + tKg, + tKs, + tma_v, + tVg, + tVs, + full_k, + empty_k, + full_v, + empty_v, + ) + else: + self._produce( + i, + sTab[0, 2 * i], + sTab[0, 2 * i + 1], + sTab[1, 2 * i], + sTab[1, 2 * i + 1], + tma_k, + tKg, + tKs, + tma_v, + tVg, + tVs, + full_k, + empty_k, + full_v, + empty_v, + ) + elif warp == 1: + # ======================= MMA issuer ======================= + tCrQ = mma_qk.make_fragment_A(sQ) + tCrK = mma_qk.make_fragment_B(sK) + tCrP = mma_pv.make_fragment_A(sP) + tCrV = mma_pv.make_fragment_B(sV) + if n_t > 0: + cute.arch.mbarrier_wait(q_full, 0) + for i in cutlass.range(n_t + 1): + if i < n_t: + s = i % S + sb = i % 2 + cute.arch.mbarrier_wait(full_k + s, (i // S) & 1) + cute.arch.mbarrier_wait(s_empty + sb, ((i // 2) & 1) ^ 1) + for kk in cutlass.range_constexpr(HEAD_DIM // 32): + mma_qk.set(tcgen05.Field.ACCUMULATE, kk != 0) + cute.gemm( + mma_qk, + tStS[(None, 0, 0, sb)], + tCrQ[(None, 0, kk, 0)], + tCrK[(None, 0, kk, s)], + tStS[(None, 0, 0, sb)], + ) + with cute.arch.elect_one(): + tcgen05.commit(s_full + sb) + tcgen05.commit(empty_k + s) + if i >= 1: + mma_pv = self._gemm2( + i - 1, + mma_pv, + tOtO, + tCrP, + tCrV, + full_v, + p_full, + empty_v, + p_empty, + o_done, + ) + with cute.arch.elect_one(): + tcgen05.commit(o_full) + elif warp >= SOFTMAX_WARP0: + # ======================= softmax (thread st = q row st = TMEM lane st) ======================= + st = tidx - SOFTMAX_WARP0 * 32 + sw = warp - SOFTMAX_WARP0 + if const_expr(self.use_pdl and not self.unsafe_skip_wait): + _griddep_wait() # the scales may come from the preceding kernels + if sw < AW: + scale = bmm1_log2 + if const_expr(self.bmm1_from_ptr): + scale = mBmm1[1] + atom_ld = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), Float32 + ) + tmpl = cute.local_tile( + tStS[(None, None), 0, 0, 0], (M_ROWS, 32), (0, 0) + ) + t2r = tcgen05.make_tmem_copy(atom_ld, tmpl) + thr_ld = t2r.get_slice(st) + frag = thr_ld.partition_D(cute.make_identity_tensor((M_ROWS, 32))) + ob = cute.make_rmem_tensor_like( + frag, Float32 + ) # epilogue chunk (32 columns) + # 64-column TMEM round trips for the lazy O rescale + atom_ld64 = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(64)), Float32 + ) + atom_st64 = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(64)), Float32 + ) + tmpl64 = cute.local_tile(tOtO[(None, None), 0, 0], (M_ROWS, 64), (0, 0)) + t2r64 = tcgen05.make_tmem_copy(atom_ld64, tmpl64) + r2t64 = tcgen05.make_tmem_copy(atom_st64, tmpl64) + thr_ld64 = t2r64.get_slice(st) + thr_st64 = r2t64.get_slice(st) + ob64 = cute.make_rmem_tensor_like( + thr_ld64.partition_D(cute.make_identity_tensor((M_ROWS, 64))), + Float32, + ) + x0 = cute.make_rmem_tensor_like(frag, Float32) + x1 = cute.make_rmem_tensor_like(frag, Float32) + x2 = cute.make_rmem_tensor_like(frag, Float32) + x3 = cute.make_rmem_tensor_like(frag, Float32) + xs = (x0, x1, x2, x3) + stt = cute.make_rmem_tensor((2,), Float32) + stt[0] = -Float32.inf + stt[1] = Float32(0.0) + lim = prefix + cutlass.min( + st // GH, QL - 1 + ) # last visible key position of this row + bars_sm = (s_full, s_empty, p_full, p_empty, o_done) + for i in cutlass.range(n_t): + # tile k = owned pages 2k (global pg_a) and 2k + 1 (pg_b, or a duplicate of pg_a past the end) + oa = 2 * (k_lo + i) + pg_a = cp_rank + oa * CP + pg_b = pg_a + CP + cnt_a = cutlass.min(cutlass.max(lim - pg_a * PAGE + 1, 0), PAGE) + cnt_b = cutlass.min(cutlass.max(lim - pg_b * PAGE + 1, 0), PAGE) + masked = Int32(0) + if (pg_b + 1) * PAGE > prefix: + masked = Int32(1) # holds new tokens or the request tail + if oa + 1 >= n_own: + masked = Int32(1) + cnt_b = Int32(0) # duplicate second slot + if masked != 0: + self._softmax_tile( + i, + xs, + ob64, + stt, + tStS, + tOtO, + t2r, + thr_ld, + t2r64, + thr_ld64, + r2t64, + thr_st64, + bars_sm, + sP_base, + sP_mk_l, + st, + cnt_a, + cnt_b, + scale, + mDbg, + dbg_row, + True, + ) + else: + self._softmax_tile( + i, + xs, + ob64, + stt, + tStS, + tOtO, + t2r, + thr_ld, + t2r64, + thr_ld64, + r2t64, + thr_st64, + bars_sm, + sP_base, + sP_mk_l, + st, + cnt_a, + cnt_b, + scale, + mDbg, + dbg_row, + False, + ) + if const_expr(self.debug_ts): + if st == 0: + mDbg[dbg_row, 3] = _globaltimer() + if n_t > 0: + # ---- normalized FP16 partial row -> staging (the K / V rings are free once o_full fired) ---- + l_tot = stt[1] + inv_l = Float32(0.0) + lse = -Float32.inf + if l_tot > Float32(0.0): + inv_l = Float32(1.0) / l_tot + lse = stt[0] + _lg2(l_tot) - self.p_shift + cute.arch.mbarrier_wait(o_full, 0) + if const_expr(self.debug_ts): + if st == 0: + mDbg[dbg_row, 9] = _globaltimer() + stage = kv_base + st * STAGE_ROW + for h in cutlass.range_constexpr(2): + accO = tOtO[(None, None), 0, h] + for c in cutlass.range_constexpr(4): + cute.copy( + t2r, + thr_ld.partition_S( + cute.local_tile(accO, (M_ROWS, 32), (0, c)) + ), + ob, + ) + cute.arch.fence_view_async_tmem_load() + for q8 in cutlass.range_constexpr(4): + _st_shared_v4( + stage + (h * 128 + c * 32 + 8 * q8) * 2, + _pack_f16x2( + ob[8 * q8] * inv_l, ob[8 * q8 + 1] * inv_l + ), + _pack_f16x2( + ob[8 * q8 + 2] * inv_l, ob[8 * q8 + 3] * inv_l + ), + _pack_f16x2( + ob[8 * q8 + 4] * inv_l, ob[8 * q8 + 5] * inv_l + ), + _pack_f16x2( + ob[8 * q8 + 6] * inv_l, ob[8 * q8 + 7] * inv_l + ), + ) + if st < R: + if const_expr(BAL): + _st_global_f32(ws_lse + (st * blse_stride + split) * 4, lse) + else: + _st_global_f32(ws_lse + (st * lse_stride + split) * 4, lse) + if const_expr(self.debug_ts): + if st == 0: + mDbg[dbg_row, 10] = _globaltimer() + elif warp == 2: + if const_expr(self.use_pdl and not self.unsafe_skip_wait): + _griddep_wait() # the epoch counter of the previous launch, bmm2 + # epoch target of the grid barrier: the counter grows by NS per launch and never resets; fewer than NS CTAs + # of this launch can have arrived yet, so floor(c / NS) is the previous launches' count. + if const_expr(self.debug_ring > 0): + if lane == 0: + sDbgT[0] = _globaltimer() + if lane == 0: + cval = _ld_relaxed_gpu_u64(mCnt.iterator.toint() + req * 8) + ns64 = Int64(NS) + sTarget[0] = (cval // ns64 + 1) * ns64 + b2v = bmm2 + if const_expr(self.bmm2_from_ptr): + b2v = mBmm2[0] + sScal[0] = b2v + elif warp == 3: + # q is the all-gather output: requested right after the dependency wait + if const_expr(self.use_pdl and not self.unsafe_skip_wait): + _griddep_wait() + if const_expr(self.early_trigger): + if lane == 0: + _griddep_launch() + if n_t > 0: + gQ = cute.local_tile( + gQ_tma, + cute.slice_((M_ROWS, TILE_KEYS, HEAD_DIM), (None, 0, None)), + (None, None, None), + ) + tCgQ = mma_qk.get_slice(0).partition_A(gQ) + tQs, tQg = cpasync.tma_partition( + tma_q, + 0, + cute.make_layout(1), + cute.group_modes(sQ, 0, 3), + cute.group_modes(tCgQ, 0, 3), + ) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(q_full, q_bytes) + cute.copy( + tma_q, tQg[(None, 0, 0, req)], tQs[(None, 0)], tma_bar_ptr=q_full + ) + + # ======================= partial -> workspace (coalesced 16-byte stores from the staging rows) ========== + cute.arch.sync_threads() + if n_t > 0: + for g in cutlass.range(tidx, R * (HEAD_DIM // 8), NUM_THREADS): + srow = g // (HEAD_DIM // 8) + sgr = g - srow * (HEAD_DIM // 8) + w4 = _ld_shared_v4_u32(kv_base + srow * STAGE_ROW + sgr * 16) + if const_expr(BAL): + ch = _div_floor_s32(g, bchunk_g) + _st_global_v4_u32( + ws_req + + ((ch * c_req + split) * bchunk_pad + (g - ch * bchunk_g) * 8) + * 2, + w4[0], + w4[1], + w4[2], + w4[3], + ) + else: + e = g * 8 + ch = e // chunk + off = e - ch * chunk + _st_global_v4_u32( + ws_base + ((ch * NS + split) * chunk_pad + off) * 2, + w4[0], + w4[1], + w4[2], + w4[3], + ) + # ======================= grid barrier over the CTAs of this request (epoch counter) ======================= + cute.arch.sync_threads() # partial written; epoch target and bmm2 in SMEM + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 4] = _globaltimer() + t_arrive = Int64(0) + t_release = Int64(0) + if tidx == 0: + cnt = mCnt.iterator.toint() + req * 8 + if const_expr(self.debug_ring > 0): + t_arrive = _globaltimer() + # red.release is cumulative over the CTA's partial stores ordered before it by the bar.sync above + if const_expr(BAL): + # every launch adds NS per counter: the request's first CTA also adds the NS - c_req absent CTAs + add = Int64(1) + if split == 0: + add = Int64(NS + 1) - Int64(c_req) + _red_release_add_u64(cnt, add) + else: + _red_release_add_u64(cnt, 1) + target = sTarget[0] + seen = _ld_acquire_gpu_u64(cnt) + while seen < target: + _nanosleep(32) + seen = _ld_acquire_gpu_u64(cnt) + if const_expr(self.debug_ring > 0): + t_release = _globaltimer() + cute.arch.sync_threads() + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 5] = _globaltimer() + + # ======================= combine: CTA `split` finalizes elements [o_lo, o_hi) of the request ============== + # gps lanes (a power of two inside one warp) share one 8-element granule g of row r and take splits j = jg, + # jg + gps, ...: one pass for the row max of the split LSEs, one for the weights 2^(lse_j - M), the weighted + # sum of the FP16 partials and the denominator, then butterflies over the gps lanes (no weights phase, no + # CTA barrier). The granule holding (r, d = 0) also writes the row's LSE. + if const_expr(BAL): + self._combine_balanced( + tidx, + warp, + lane, + req, + split, + c_req, + n_valid, + bchunk, + bchunk_pad, + blse_stride, + ws_req, + ws_lse, + kv_base, + comb_bar, + o_base, + o_sd, + o_st, + o_sh, + l_base, + l_sd, + l_st, + l_sh, + sScal, + mDbg, + dbg_row, + ) + else: + o_lo = split * chunk + if o_lo < R * HEAD_DIM: + o_hi = cutlass.min(o_lo + chunk, R * HEAD_DIM) + row0 = o_lo // HEAD_DIM + nrows = (o_hi - 1) // HEAD_DIM - row0 + 1 + n_gran = (o_hi - o_lo) // 8 + s_comb = kv_base # [n_valid][chunk_pad] FP16 + s_lse = kv_base + self.comb_o_bytes # [nrows][lse_stride] FP32 + jg = tidx % self.gps + if n_valid > 0: + if warp == 0: + cute.arch.fence_proxy( + "async.global" + ) # acquire (generic proxy) -> bulk reads (async proxy) + if lane == 0: + nb_o = n_valid * (chunk_pad * 2) + nb_l = nrows * (lse_stride * 4) + cute.arch.mbarrier_arrive_and_expect_tx( + comb_bar, nb_o + nb_l + ) + _bulk_g2s( + s_comb, + ws_base + (split * NS * chunk_pad) * 2, + nb_o, + comb_bar.toint(), + ) + _bulk_g2s( + s_lse, + ws_lse + (row0 * lse_stride) * 4, + nb_l, + comb_bar.toint(), + ) + # spin (no suspend hint): every lane of a warp must be awake for the shuffles below + landed = cute.arch.mbarrier_try_wait(comb_bar, 0) + while not landed: + landed = cute.arch.mbarrier_try_wait(comb_bar, 0) + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 6] = _globaltimer() + # granule passes (gran_iters == 1 up to ~8 requests: the original single pass) + for gi in cutlass.range_constexpr(self.gran_iters): + g = gi * (NUM_THREADS // self.gps) + tidx // self.gps + gc = cutlass.min(g, n_gran - 1) + e0 = o_lo + gc * 8 + orow = e0 // HEAD_DIM + dd, tok, hh = self._out_coords(orow, req) + addr = ( + o_base + + (dd * o_sd + tok * o_st + hh * o_sh + e0 % HEAD_DIM) * 2 + ) + if n_valid > 0: + lrow = s_lse + (orow - row0) * (lse_stride * 4) + mrow = -Float32.inf + for j in cutlass.range(jg, n_valid, self.gps): + mrow = cute.arch.fmax(mrow, _ld_shared_f32(lrow + j * 4)) + for lg in cutlass.range_constexpr(self.gps.bit_length() - 1): + mrow = cute.arch.fmax( + mrow, + cute.arch.shuffle_sync_bfly( + mrow, offset=self.gps >> (lg + 1) + ), + ) + m_use = mrow + if mrow == -Float32.inf: + m_use = Float32( + 0.0 + ) # no visible key in any split: every weight is exp2(-inf) = 0 + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 11] = _globaltimer() + den = Float32(0.0) + a0 = Float32(0.0) + a1 = Float32(0.0) + a2 = Float32(0.0) + a3 = Float32(0.0) + a4 = Float32(0.0) + a5 = Float32(0.0) + a6 = Float32(0.0) + a7 = Float32(0.0) + for j in cutlass.range(jg, n_valid, self.gps): + wj = _ex2(_ld_shared_f32(lrow + j * 4) - m_use) + den = den + wj + u4 = _ld_shared_v4_u32( + s_comb + (j * chunk_pad + gc * 8) * 2 + ) + f0, f1 = _unpack_f16x2(u4[0]) + f2, f3 = _unpack_f16x2(u4[1]) + f4, f5 = _unpack_f16x2(u4[2]) + f6, f7 = _unpack_f16x2(u4[3]) + a0 = _fma(wj, f0, a0) + a1 = _fma(wj, f1, a1) + a2 = _fma(wj, f2, a2) + a3 = _fma(wj, f3, a3) + a4 = _fma(wj, f4, a4) + a5 = _fma(wj, f5, a5) + a6 = _fma(wj, f6, a6) + a7 = _fma(wj, f7, a7) + for lg in cutlass.range_constexpr(self.gps.bit_length() - 1): + sh_off = self.gps >> (lg + 1) + den = den + cute.arch.shuffle_sync_bfly(den, offset=sh_off) + a0 = a0 + cute.arch.shuffle_sync_bfly(a0, offset=sh_off) + a1 = a1 + cute.arch.shuffle_sync_bfly(a1, offset=sh_off) + a2 = a2 + cute.arch.shuffle_sync_bfly(a2, offset=sh_off) + a3 = a3 + cute.arch.shuffle_sync_bfly(a3, offset=sh_off) + a4 = a4 + cute.arch.shuffle_sync_bfly(a4, offset=sh_off) + a5 = a5 + cute.arch.shuffle_sync_bfly(a5, offset=sh_off) + a6 = a6 + cute.arch.shuffle_sync_bfly(a6, offset=sh_off) + a7 = a7 + cute.arch.shuffle_sync_bfly(a7, offset=sh_off) + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 12] = _globaltimer() + if jg == 0: + if g < n_gran: + scl = Float32( + 0.0 + ) # a row without any visible key gives 0, not NaN + if den > Float32(0.0): + scl = sScal[0] / den + w0 = _pack_bf16x2(a0 * scl, a1 * scl) + w1 = _pack_bf16x2(a2 * scl, a3 * scl) + w2 = _pack_bf16x2(a4 * scl, a5 * scl) + w3 = _pack_bf16x2(a6 * scl, a7 * scl) + if const_expr(self.o_vec16): + _st_global_v4_u32(addr, w0, w1, w2, w3) + else: + _st_global_u32(addr, w0) + _st_global_u32(addr + 4, w1) + _st_global_u32(addr + 8, w2) + _st_global_u32(addr + 12, w3) + if e0 % HEAD_DIM == 0: + lse_out = -Float32.inf + if den > Float32(0.0): + lse_out = mrow + _lg2(den) + _st_global_f32( + l_base + + (dd * l_sd + tok * l_st + hh * l_sh) * 4, + lse_out, + ) + else: + # this rank sees no key of the request: O = 0, LSE = -inf + if jg == 0: + if g < n_gran: + if const_expr(self.o_vec16): + _st_global_v4_u32( + addr, Uint32(0), Uint32(0), Uint32(0), Uint32(0) + ) + else: + _st_global_u32(addr, Uint32(0)) + _st_global_u32(addr + 4, Uint32(0)) + _st_global_u32(addr + 8, Uint32(0)) + _st_global_u32(addr + 12, Uint32(0)) + if e0 % HEAD_DIM == 0: + _st_global_f32( + l_base + + (dd * l_sd + tok * l_st + hh * l_sh) * 4, + -Float32.inf, + ) + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 7] = _globaltimer() + # ======================= release the dependents: after a CTA barrier that follows the last global store ===== + if const_expr(self.debug_ring > 0): + if tidx == 0: + launch = sTarget[0] // Int64(NS) - Int64(1) + slot = launch % Int64(self.debug_ring) + if const_expr(not BAL): + cta = req * NS + split + mDbg[slot, cta, 0] = t_entry + mDbg[slot, cta, 1] = sDbgT[0] + mDbg[slot, cta, 2] = t_arrive + mDbg[slot, cta, 3] = t_release + mDbg[slot, cta, 5] = Int64(_smid()) + mDbg[slot, cta, 6] = launch * Int64(64) + Int64(mSeq.shape[0]) + mDbg[slot, cta, 7] = t_setup + mDbg[slot, cta, 4] = _globaltimer() + cute.arch.sync_threads() + if const_expr(self.use_pdl): + if tidx == 0: + _griddep_launch() + if warp == 1: + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS) + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 8] = _globaltimer() + + # ---------------------------------------------------------------- device kernel, two-slot plan + @cute.kernel + def kernel_two( + self, + mma_qk: cute.TiledMma, + mma_pv: cute.TiledMma, + tma_k: cute.CopyAtom, + gK_tma: cute.Tensor, + tma_v: cute.CopyAtom, + gV_tma: cute.Tensor, + tma_q: cute.CopyAtom, + gQ_tma: cute.Tensor, + mTab: cute.Tensor, + mSeq: cute.Tensor, + cp_rank: Int32, + mBmm1: cute.Tensor, + bmm1_log2: Float32, + mBmm2: cute.Tensor, + bmm2: Float32, + mO: cute.Tensor, + o_sd: Int32, + o_st: Int32, + o_sh: Int32, + mLse: cute.Tensor, + l_sd: Int32, + l_st: Int32, + l_sh: Int32, + mWs: cute.Tensor, + mCnt: cute.Tensor, + mDbg: cute.Tensor, + sQ_l: cute.ComposedLayout, + sK_l: cute.ComposedLayout, + sV_l: cute.ComposedLayout, + sP_l: cute.ComposedLayout, + sK_tma_l: cute.ComposedLayout, + sV_tma_l: cute.ComposedLayout, + sP_mk_l: cute.ComposedLayout, + q_bytes: cutlass.Constexpr, + ): + """``kernel`` on the two-slot plan (_plan_two): the CTA computes up to two requests back to back (slot 0, then + slot 1) through one K / V ring and one O accumulator, then runs the grid barrier and the combine of each slot's + request. Ring and barrier phases count the CTA's tiles across both slots; the q buffer (q_empty: MMA -> q + loader, after slot 0's last GEMM1) and O (o_empty: softmax -> MMA) change hands between the slots. The K / V + ring the one-slot kernel stages the FP16 partials through holds slot 1's tiles by then, so the softmax rows + store theirs straight to the workspace (staging them through each warp's own P rows measured 0.6-2.3 us per + call slower at b72-128).""" + S: cutlass.Constexpr = KV_STAGES + R: cutlass.Constexpr = self.rows + QL: cutlass.Constexpr = self.q_len + GH: cutlass.Constexpr = self.gh + CP: cutlass.Constexpr = self.cp + NS: cutlass.Constexpr = self.n_splits + AW: cutlass.Constexpr = self.act_warps + TOT: cutlass.Constexpr = self.rows * HEAD_DIM + tidx, _, _ = cute.arch.thread_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + cta, _, _ = cute.arch.block_idx() + dbg_row = cta + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 0] = _globaltimer() + o_base = mO.iterator.toint() + l_base = mLse.iterator.toint() + + # ---- shared memory (as ``kernel``, + q_empty / o_empty and a second plan) ---- + smem = SmemAllocator() + sK = smem.allocate_tensor(FP8, sK_l.outer, 1024, sK_l.inner) + sV = smem.allocate_tensor(FP8, sV_l.outer, 1024, sV_l.inner) + sQ = smem.allocate_tensor(FP8, sQ_l.outer, 1024, sQ_l.inner) + sP = smem.allocate_tensor(FP8, sP_l.outer, 1024, sP_l.inner) + bars = smem.allocate_array(Int64, 4 * S + 14) + full_k = bars + empty_k = bars + S + full_v = bars + 2 * S + empty_v = bars + 3 * S + q_full = bars + 4 * S + s_full = bars + 4 * S + 1 # [2] + s_empty = bars + 4 * S + 3 # [2] + p_full = bars + 4 * S + 5 # [2] + p_empty = bars + 4 * S + 7 # [2] + o_done = bars + 4 * S + 9 + o_full = bars + 4 * S + 10 + comb_bar = bars + 4 * S + 11 + q_empty = ( + bars + 4 * S + 12 + ) # slot 0's GEMM1s done: q of slot 1 may overwrite sQ + o_empty = ( + bars + 4 * S + 13 + ) # slot 0's O read by the softmax rows: slot 1's first GEMM2 may overwrite it + tmem_holder = smem.allocate_array(Int32, 4) + sTab = smem.allocate_tensor( + Int32, cute.make_layout((2, TAB_MAX), stride=(TAB_MAX, 1)), 16 + ) + sTarget = smem.allocate_tensor(Int64, cute.make_layout((2,)), 16) + sScal = smem.allocate_tensor(Float32, cute.make_layout((4,)), 16) + sK_tma = cute.make_tensor( + cute.recast_ptr(sK.iterator, sK_tma_l.inner, dtype=FP8), sK_tma_l.outer + ) + sV_tma = cute.make_tensor( + cute.recast_ptr(sV.iterator, sV_tma_l.inner, dtype=FP8), sV_tma_l.outer + ) + kv_base = sK.iterator.toint() + sP_base = sP.iterator.toint() + sPlan = smem.allocate_tensor(Int32, cute.make_layout((2 * PLAN_WORDS,)), 16) + if warp == 2: + if const_expr(self.debug_ts): + # the plan's cost: %globaltimer before / after (19, 20), %clock64 cycles (21), its T (22) + t_plan = _globaltimer() + c_plan = _clock64() + T_plan = self._plan_two(mSeq, cp_rank, cta, lane, sPlan, kv_base) + cute.arch.sync_warp() + c_plan = _clock64() - c_plan + if lane == 0: + mDbg[dbg_row, 19] = t_plan + mDbg[dbg_row, 20] = _globaltimer() + mDbg[dbg_row, 21] = c_plan + mDbg[dbg_row, 22] = Int64(T_plan) + else: + self._plan_two(mSeq, cp_rank, cta, lane, sPlan, kv_base) + + if warp == 0: + cpasync.prefetch_descriptor(tma_k) + cpasync.prefetch_descriptor(tma_v) + cpasync.prefetch_descriptor(tma_q) + if tidx == 0: + for i in cutlass.range_constexpr(4 * S + 1): + cute.arch.mbarrier_init(bars + i, 1) + for i in cutlass.range_constexpr(2): + cute.arch.mbarrier_init(s_full + i, 1) + cute.arch.mbarrier_init(s_empty + i, 32 * AW) + cute.arch.mbarrier_init(p_full + i, 32 * AW) + cute.arch.mbarrier_init(p_empty + i, 1) + cute.arch.mbarrier_init(o_done, 1) + cute.arch.mbarrier_init(o_full, 1) + cute.arch.mbarrier_init(comb_bar, 1) + cute.arch.mbarrier_init(q_empty, 1) + cute.arch.mbarrier_init(o_empty, 32 * AW) + cute.arch.mbarrier_init_fence() + if const_expr(AW < 4): + n_zero: cutlass.Constexpr = ( + (M_ROWS - 32 * AW) * TILE_KEYS // 16 + ) # 16-byte words per buffer + for z in cutlass.range(tidx, 2 * n_zero, NUM_THREADS): + zb = z // n_zero + zw = z - zb * n_zero + _st_shared_v4( + sP_base + zb * (M_ROWS * TILE_KEYS) + 32 * AW * TILE_KEYS + zw * 16, + Uint32(0), + Uint32(0), + Uint32(0), + Uint32(0), + ) + cute.arch.fence_proxy("async.shared", space="cta") + if warp == 1: + cute.arch.alloc_tmem(TMEM_COLS, tmem_holder) + cute.arch.sync_threads() + + # ---- the two slots (request -1: no slot; tiles 0 then) ---- + req0 = cute.arch.make_warp_uniform(sPlan[0]) + split0 = cute.arch.make_warp_uniform(sPlan[1]) + c0 = cute.arch.make_warp_uniform(sPlan[2]) + klo0 = cute.arch.make_warp_uniform(sPlan[3]) + nt0 = cute.arch.make_warp_uniform(sPlan[4]) + first0 = cute.arch.make_warp_uniform(sPlan[5]) + L0 = cute.arch.make_warp_uniform(sPlan[6]) + nv0 = cute.arch.make_warp_uniform(sPlan[7]) + req1 = cute.arch.make_warp_uniform(sPlan[8]) + split1 = cute.arch.make_warp_uniform(sPlan[9]) + c1 = cute.arch.make_warp_uniform(sPlan[10]) + klo1 = cute.arch.make_warp_uniform(sPlan[11]) + nt1 = cute.arch.make_warp_uniform(sPlan[12]) + first1 = cute.arch.make_warp_uniform(sPlan[13]) + L1 = cute.arch.make_warp_uniform(sPlan[14]) + nv1 = cute.arch.make_warp_uniform(sPlan[15]) + a0 = Int32( + 0 + ) # slot 0 has tiles (loads q, runs the MMA / softmax pipeline, writes a partial) + if nt0 > 0: + a0 = Int32(1) + own0 = (((L0 + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + own1 = (((L1 + PAGE - 1) // PAGE) - cp_rank + CP - 1) // CP + prefix0 = L0 - QL + prefix1 = L1 - QL + fresh0 = cutlass.max(prefix0 // PAGE - self.fresh_margin_pages, 0) + fresh1 = cutlass.max(prefix1 // PAGE - self.fresh_margin_pages, 0) + # each request's combine geometry (cp_partial_geometry with its CTA count) and workspace slots + cc0 = cutlass.max(c0, 1) + cc1 = cutlass.max(c1, 1) + bch0 = 8 * _div_floor_s32(R * HEAD_DIM + 8 * cc0 - 1, 8 * cc0) + bch1 = 8 * _div_floor_s32(R * HEAD_DIM + 8 * cc1 - 1, 8 * cc1) + bpad0 = bch0 + if (bch0 // 8) % 2 == 0: + bpad0 = bch0 + 8 + bpad1 = bch1 + if (bch1 // 8) % 2 == 0: + bpad1 = bch1 + 8 + blse0 = 4 * ((cc0 + 3) // 4) + blse1 = 4 * ((cc1 + 3) // 4) + ws_req0 = mWs.iterator.toint() + first0 * self.slot_o_bytes + ws_req1 = mWs.iterator.toint() + first1 * self.slot_o_bytes + ws_lse0 = mWs.iterator.toint() + self.lse_area_off + first0 * self.lse_cta_bytes + ws_lse1 = mWs.iterator.toint() + self.lse_area_off + first1 * self.lse_cta_bytes + if const_expr(self.debug_ts): + if tidx == 0: + # the plan, for the tests: per slot (request * 65536 + index, CTAs * 65536 + first slot, + # first tile * 65536 + tiles), request -1 = no slot + mDbg[dbg_row, 13] = Int64(-1) + mDbg[dbg_row, 16] = Int64(-1) + if req0 >= 0: + mDbg[dbg_row, 13] = Int64(req0) * Int64(65536) + Int64(split0) + mDbg[dbg_row, 14] = Int64(c0) * Int64(65536) + Int64(first0) + mDbg[dbg_row, 15] = Int64(klo0) * Int64(65536) + Int64(nt0) + if req1 >= 0: + mDbg[dbg_row, 16] = Int64(req1) * Int64(65536) + Int64(split1) + mDbg[dbg_row, 17] = Int64(c1) * Int64(65536) + Int64(first1) + mDbg[dbg_row, 18] = Int64(klo1) * Int64(65536) + Int64(nt1) + tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, tmem_holder) + tStS = mma_qk.make_fragment_C(mma_qk.partition_shape_C((M_ROWS, TILE_KEYS, 2))) + tStS = cute.make_tensor(tmem_ptr, tStS.layout) + s_cols = tcgen05.find_tmem_tensor_col_offset(tStS) + tOtO = mma_pv.make_fragment_C(mma_pv.partition_shape_C((M_ROWS, HEAD_DIM))) + tOtO = cute.make_tensor(tmem_ptr + s_cols, tOtO.layout) + + if warp == 0: + # ======================= TMA producer (tile index i of the CTA runs on across the slots) ============ + gK = cute.tiled_divide(gK_tma, (PAGE, SW_BYTES)) + gV = cute.tiled_divide(gV_tma, (SW_BYTES, PAGE)) + tKs, tKg = cpasync.tma_partition(tma_k, 0, cute.make_layout(1), sK_tma, gK) + tVs, tVg = cpasync.tma_partition(tma_v, 0, cute.make_layout(1), sV_tma, gV) + waited = Int32(0) + if const_expr(not self.kv_pre_wait): + if const_expr(self.use_pdl and not self.unsafe_skip_wait): + _griddep_wait() + waited = Int32(1) + for k in cutlass.range_constexpr(2): + rq = (req0, req1)[k] + klo = (klo0, klo1)[k] + nt = (nt0, nt1)[k] + n_own = (own0, own1)[k] + fresh = (fresh0, fresh1)[k] + gb = (Int32(0), nt0)[k] # tiles of the slots before this one + if nt > 0: + cute.arch.sync_warp() # every lane is past the previous slot's table reads + for j in cutlass.range(lane, cutlass.min(2 * nt, TAB_MAX), 32): + oi = cutlass.min(2 * klo + j, n_own - 1) + pg = cp_rank + oi * CP + sTab[0, j] = mTab[rq, 0, pg] + sTab[1, j] = mTab[rq, 1, pg] + cute.arch.sync_warp() + for i in cutlass.range(nt): + if waited == 0: + if ( + cp_rank + cutlass.min(2 * (klo + i) + 1, n_own - 1) * CP + >= fresh + ): + if const_expr( + self.use_pdl and not self.unsafe_skip_wait + ): + _griddep_wait() + waited = Int32(1) + ti = i % (TAB_MAX // 2) + if ti == 0: + if i > 0: + cute.arch.sync_warp() + for j in cutlass.range( + lane, cutlass.min(2 * (nt - i), TAB_MAX), 32 + ): + oi = cutlass.min(2 * (klo + i) + j, n_own - 1) + pg = cp_rank + oi * CP + sTab[0, j] = mTab[rq, 0, pg] + sTab[1, j] = mTab[rq, 1, pg] + cute.arch.sync_warp() + self._produce( + gb + i, + sTab[0, 2 * ti], + sTab[0, 2 * ti + 1], + sTab[1, 2 * ti], + sTab[1, 2 * ti + 1], + tma_k, + tKg, + tKs, + tma_v, + tVg, + tVs, + full_k, + empty_k, + full_v, + empty_v, + ) + elif warp == 1: + # ======================= MMA issuer ======================= + tCrQ = mma_qk.make_fragment_A(sQ) + tCrK = mma_qk.make_fragment_B(sK) + tCrP = mma_pv.make_fragment_A(sP) + tCrV = mma_pv.make_fragment_B(sV) + for k in cutlass.range_constexpr(2): + nt = (nt0, nt1)[k] + gb = (Int32(0), nt0)[k] + qph = (Int32(0), a0)[k] # q loads before this slot's + if nt > 0: + cute.arch.mbarrier_wait(q_full, qph) + for i in cutlass.range(nt + 1): + if i < nt: + g = gb + i + s = g % S + sb = g % 2 + cute.arch.mbarrier_wait(full_k + s, (g // S) & 1) + cute.arch.mbarrier_wait(s_empty + sb, ((g // 2) & 1) ^ 1) + for kk in cutlass.range_constexpr(HEAD_DIM // 32): + mma_qk.set(tcgen05.Field.ACCUMULATE, kk != 0) + cute.gemm( + mma_qk, + tStS[(None, 0, 0, sb)], + tCrQ[(None, 0, kk, 0)], + tCrK[(None, 0, kk, s)], + tStS[(None, 0, 0, sb)], + ) + with cute.arch.elect_one(): + tcgen05.commit(s_full + sb) + tcgen05.commit(empty_k + s) + if const_expr(k == 0): + if i == nt - 1: + tcgen05.commit( + q_empty + ) # slot 0's last GEMM1: sQ may take slot 1's q + if i >= 1: + if const_expr(k == 1): + if i == 1: + if a0 != 0: + cute.arch.mbarrier_wait( + o_empty, 0 + ) # slot 0's O has been read out + mma_pv = self._gemm2( + gb + i - 1, + mma_pv, + tOtO, + tCrP, + tCrV, + full_v, + p_full, + empty_v, + p_empty, + o_done, + j_loc=i - 1, + ) + with cute.arch.elect_one(): + tcgen05.commit(o_full) + elif warp >= SOFTMAX_WARP0: + # ======================= softmax (thread st = q row st = TMEM lane st) ======================= + st = tidx - SOFTMAX_WARP0 * 32 + sw = warp - SOFTMAX_WARP0 + if const_expr(self.use_pdl and not self.unsafe_skip_wait): + _griddep_wait() # the scales may come from the preceding kernels + if sw < AW: + scale = bmm1_log2 + if const_expr(self.bmm1_from_ptr): + scale = mBmm1[1] + atom_ld = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), Float32 + ) + tmpl = cute.local_tile( + tStS[(None, None), 0, 0, 0], (M_ROWS, 32), (0, 0) + ) + t2r = tcgen05.make_tmem_copy(atom_ld, tmpl) + thr_ld = t2r.get_slice(st) + frag = thr_ld.partition_D(cute.make_identity_tensor((M_ROWS, 32))) + ob = cute.make_rmem_tensor_like(frag, Float32) + atom_ld64 = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(64)), Float32 + ) + atom_st64 = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(64)), Float32 + ) + tmpl64 = cute.local_tile(tOtO[(None, None), 0, 0], (M_ROWS, 64), (0, 0)) + t2r64 = tcgen05.make_tmem_copy(atom_ld64, tmpl64) + r2t64 = tcgen05.make_tmem_copy(atom_st64, tmpl64) + thr_ld64 = t2r64.get_slice(st) + thr_st64 = r2t64.get_slice(st) + ob64 = cute.make_rmem_tensor_like( + thr_ld64.partition_D(cute.make_identity_tensor((M_ROWS, 64))), + Float32, + ) + x0 = cute.make_rmem_tensor_like(frag, Float32) + x1 = cute.make_rmem_tensor_like(frag, Float32) + x2 = cute.make_rmem_tensor_like(frag, Float32) + x3 = cute.make_rmem_tensor_like(frag, Float32) + xs = (x0, x1, x2, x3) + stt = cute.make_rmem_tensor((2,), Float32) + bars_sm = (s_full, s_empty, p_full, p_empty, o_done) + for k in cutlass.range_constexpr(2): + nt = (nt0, nt1)[k] + klo = (klo0, klo1)[k] + n_own = (own0, own1)[k] + prefix = (prefix0, prefix1)[k] + gb = (Int32(0), nt0)[k] + oph = (Int32(0), a0)[k] # o_full completions before this slot's + split_k = (split0, split1)[k] + c_k = (cc0, cc1)[k] + bch = (bch0, bch1)[k] + bpad = (bpad0, bpad1)[k] + blse = (blse0, blse1)[k] + ws_req = (ws_req0, ws_req1)[k] + ws_lse = (ws_lse0, ws_lse1)[k] + if nt > 0: + stt[0] = -Float32.inf + stt[1] = Float32(0.0) + lim = prefix + cutlass.min( + st // GH, QL - 1 + ) # last visible key position of this row + for i in cutlass.range(nt): + oa = 2 * (klo + i) + pg_a = cp_rank + oa * CP + pg_b = pg_a + CP + cnt_a = cutlass.min( + cutlass.max(lim - pg_a * PAGE + 1, 0), PAGE + ) + cnt_b = cutlass.min( + cutlass.max(lim - pg_b * PAGE + 1, 0), PAGE + ) + masked = Int32(0) + if (pg_b + 1) * PAGE > prefix: + masked = Int32( + 1 + ) # holds new tokens or the request tail + if oa + 1 >= n_own: + masked = Int32(1) + cnt_b = Int32(0) # duplicate second slot + if masked != 0: + self._softmax_tile( + gb + i, + xs, + ob64, + stt, + tStS, + tOtO, + t2r, + thr_ld, + t2r64, + thr_ld64, + r2t64, + thr_st64, + bars_sm, + sP_base, + sP_mk_l, + st, + cnt_a, + cnt_b, + scale, + mDbg, + dbg_row, + True, + i_loc=i, + ) + else: + self._softmax_tile( + gb + i, + xs, + ob64, + stt, + tStS, + tOtO, + t2r, + thr_ld, + t2r64, + thr_ld64, + r2t64, + thr_st64, + bars_sm, + sP_base, + sP_mk_l, + st, + cnt_a, + cnt_b, + scale, + mDbg, + dbg_row, + False, + i_loc=i, + ) + # ---- normalized FP16 partial row -> the workspace (chunk-major, see _combine_balanced) ---- + l_tot = stt[1] + inv_l = Float32(0.0) + lse = -Float32.inf + if l_tot > Float32(0.0): + inv_l = Float32(1.0) / l_tot + lse = stt[0] + _lg2(l_tot) - self.p_shift + cute.arch.mbarrier_wait(o_full, oph) + for h in cutlass.range_constexpr(2): + accO = tOtO[(None, None), 0, h] + for c in cutlass.range_constexpr(4): + cute.copy( + t2r, + thr_ld.partition_S( + cute.local_tile(accO, (M_ROWS, 32), (0, c)) + ), + ob, + ) + cute.arch.fence_view_async_tmem_load() + if st < R: + for q8 in cutlass.range_constexpr(4): + e = st * HEAD_DIM + (h * 128 + c * 32 + 8 * q8) + ch = _div_floor_s32(e, bch) + _st_global_v4_u32( + ws_req + + ( + (ch * c_k + split_k) * bpad + + e + - ch * bch + ) + * 2, + _pack_f16x2( + ob[8 * q8] * inv_l, + ob[8 * q8 + 1] * inv_l, + ), + _pack_f16x2( + ob[8 * q8 + 2] * inv_l, + ob[8 * q8 + 3] * inv_l, + ), + _pack_f16x2( + ob[8 * q8 + 4] * inv_l, + ob[8 * q8 + 5] * inv_l, + ), + _pack_f16x2( + ob[8 * q8 + 6] * inv_l, + ob[8 * q8 + 7] * inv_l, + ), + ) + if st < R: + _st_global_f32(ws_lse + (st * blse + split_k) * 4, lse) + if const_expr(k == 0): + cute.arch.mbarrier_arrive( + o_empty + ) # this row's share of O is read out + elif warp == 2: + if const_expr(self.use_pdl and not self.unsafe_skip_wait): + _griddep_wait() # the epoch counters of the previous launch, bmm2 + if lane == 0: + ns64 = Int64(NS) + if req0 >= 0: + cv0 = _ld_relaxed_gpu_u64(mCnt.iterator.toint() + req0 * 8) + sTarget[0] = (cv0 // ns64 + 1) * ns64 + if req1 >= 0: + cv1 = _ld_relaxed_gpu_u64(mCnt.iterator.toint() + req1 * 8) + sTarget[1] = (cv1 // ns64 + 1) * ns64 + b2v = bmm2 + if const_expr(self.bmm2_from_ptr): + b2v = mBmm2[0] + sScal[0] = b2v + elif warp == 3: + # q of each slot's request after the dependency wait; slot 1's once slot 0's GEMM1s are done with sQ + if const_expr(self.use_pdl and not self.unsafe_skip_wait): + _griddep_wait() + if nt0 + nt1 > 0: + gQ = cute.local_tile( + gQ_tma, + cute.slice_((M_ROWS, TILE_KEYS, HEAD_DIM), (None, 0, None)), + (None, None, None), + ) + tCgQ = mma_qk.get_slice(0).partition_A(gQ) + tQs, tQg = cpasync.tma_partition( + tma_q, + 0, + cute.make_layout(1), + cute.group_modes(sQ, 0, 3), + cute.group_modes(tCgQ, 0, 3), + ) + if nt0 > 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(q_full, q_bytes) + cute.copy( + tma_q, + tQg[(None, 0, 0, req0)], + tQs[(None, 0)], + tma_bar_ptr=q_full, + ) + if nt1 > 0: + if nt0 > 0: + cute.arch.mbarrier_wait(q_empty, 0) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(q_full, q_bytes) + cute.copy( + tma_q, + tQg[(None, 0, 0, req1)], + tQs[(None, 0)], + tma_bar_ptr=q_full, + ) + + # ======================= grid barrier and combine of each slot's request ======================= + cute.arch.sync_threads() # both partials written (softmax rows); epoch targets and bmm2 in SMEM + if tidx == 0: + # red.release is cumulative over the CTA's partial stores ordered before it by the bar.sync above; every + # launch adds NS per counter: the request's first CTA also adds the NS - c absent CTAs + if req0 >= 0: + add0 = Int64(1) + if split0 == 0: + add0 = Int64(NS + 1) - Int64(c0) + _red_release_add_u64(mCnt.iterator.toint() + req0 * 8, add0) + if req1 >= 0: + add1 = Int64(1) + if split1 == 0: + add1 = Int64(NS + 1) - Int64(c1) + _red_release_add_u64(mCnt.iterator.toint() + req1 * 8, add1) + cph1 = Int32( + 0 + ) # comb_bar phase of slot 1: 1 when slot 0's combine loaded partials + if req0 >= 0: + if nv0 > 0: + if split0 * bch0 < TOT: + cph1 = Int32(1) + for k in cutlass.range_constexpr(2): + rq = (req0, req1)[k] + if rq >= 0: + if tidx == 0: + cnt = mCnt.iterator.toint() + rq * 8 + target = sTarget[k] + seen = _ld_acquire_gpu_u64(cnt) + while seen < target: + _nanosleep(32) + seen = _ld_acquire_gpu_u64(cnt) + cute.arch.sync_threads() + self._combine_balanced( + tidx, + warp, + lane, + rq, + (split0, split1)[k], + (cc0, cc1)[k], + (nv0, nv1)[k], + (bch0, bch1)[k], + (bpad0, bpad1)[k], + (blse0, blse1)[k], + (ws_req0, ws_req1)[k], + (ws_lse0, ws_lse1)[k], + kv_base, + comb_bar, + o_base, + o_sd, + o_st, + o_sh, + l_base, + l_sd, + l_st, + l_sh, + sScal, + mDbg, + dbg_row, + comb_phase=(Int32(0), cph1)[k], + ) + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 7] = _globaltimer() + # ======================= release the dependents: after a CTA barrier that follows the last global store ===== + cute.arch.sync_threads() + if const_expr(self.use_pdl): + if tidx == 0: + _griddep_launch() + if warp == 1: + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS) + if const_expr(self.debug_ts): + if tidx == 0: + mDbg[dbg_row, 8] = _globaltimer() diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index f531c525..fa28d83c 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -153,7 +153,7 @@ }, # SKV2 against SGLang's TRT-LLM MHA decode (test_skv2_is_within_twice_sglangs_error) "skv2": {"skv2.attention.rel_l2_vs_stock": 2.0}, - # CP-lite's exact path (SGLang's q35_cp_lite) at one request, as skv2's bound + # CP-lite's exact plugin path at one request, as skv2's bound "skv2_cp": {"skv2_cp.attention.rel_l2_vs_stock": 2.0}, "draft_q1_cp": {"draft_q1_cp.attention.rel_l2_vs_stock": 2.0}, "draft_skv2_cp": {"draft_skv2_cp.attention.rel_l2_vs_stock": 2.0}, diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 0f800724..99ff26af 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -20,8 +20,13 @@ @pytest.fixture(autouse=True) def registered(monkeypatch): + from sglang.srt.kernel_plugins.nvidia import binding as plugin_binding + monkeypatch.setattr(registry, "_REGISTRY", {}) monkeypatch.delenv(plugin.PARITY_SORT_ENV, raising=False) + monkeypatch.setattr( + plugin_binding, "megakernel_agree", lambda *args, **kwargs: None + ) plugin.register() # the target and draft builds read moe_ep_size (expert parallelism); no process groups here monkeypatch.setattr( @@ -77,6 +82,113 @@ def test_an_explicit_empty_policy_keeps_every_slot_off(monkeypatch): assert runtime._requested_slots() == frozenset() +def test_cp_slots_build_the_plugin_state_before_validation(monkeypatch): + from sglang.srt.kernel_plugins.nvidia import binding as plugin_binding + + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + built = object() + calls = [] + agreements = [] + monkeypatch.setattr(runtime, "SLOTS", frozenset({"draft_q1_cp"})) + monkeypatch.setattr( + plugin_binding, + "megakernel_agree", + lambda value, *, what: agreements.append((value, what)), + ) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 3) + monkeypatch.setattr( + CpLite, + "maybe_create", + lambda binding, device: calls.append((binding, device)) or built, + ) + + binding = SimpleNamespace(page_size=64, tp_size=8) + runtime.build_cp_lite(binding) + + assert agreements == [ + (("draft_q1_cp",), "Qwen3.5 CP slots requested"), + ((64, 8), "Qwen3.5 CP page size and TP size"), + ] + assert calls == [(binding, torch.device("cuda", 3))] + + +def test_cp_slots_refuse_an_incompatible_page_or_topology(monkeypatch): + from sglang.srt.kernel_plugins.nvidia import binding as plugin_binding + + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + monkeypatch.setattr(runtime, "SLOTS", frozenset({"skv2_cp"})) + monkeypatch.setattr( + plugin_binding, "megakernel_agree", lambda *args, **kwargs: None + ) + with pytest.raises(RuntimeError, match="page size 64"): + runtime.build_cp_lite(SimpleNamespace(page_size=32, tp_size=8)) + + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(CpLite, "maybe_create", lambda binding, device: None) + with pytest.raises(RuntimeError, match="Q/KV-head topology"): + runtime.build_cp_lite(SimpleNamespace(page_size=64, tp_size=8)) + + +@pytest.mark.parametrize( + "tp, rank, want", + [(4, 3, (3, 4, 2, 8)), (8, 7, (7, 8, 4, 4))], +) +def test_cp_state_uses_the_pinned_qwen35_head_topology(monkeypatch, tp, rank, want): + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + calls = [] + group = object() + binding = SimpleNamespace( + tp_size=tp, + tp_rank=rank, + tp_device_group=group, + ) + monkeypatch.setattr(CpLite, "_instance", None) + monkeypatch.setattr( + CpLite, + "__init__", + lambda self, got_group, got_rank, got_tp, cp, hq, device: calls.append( + (got_group, got_rank, got_tp, cp, hq, device) + ), + ) + + made = CpLite.maybe_create(binding, torch.device("cuda", 0)) + + assert isinstance(made, CpLite) + assert calls == [(group, *want, torch.device("cuda", 0))] + + +def test_requested_cp_state_does_not_restore_a_device_length_gate(): + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + cp = CpLite.__new__(CpLite) + cp.min_kv = {} + cp.gate_max = False + + assert [cp._gate_thr(batch) for batch in (2, 4, 16)] == [None, None, None] + + +def test_cp_gate_epilogue_offer_is_consumed_once(): + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + cp = CpLite.__new__(CpLite) + cp._gate_epilogue = None + cp._gate_epilogue_done = False + gate, scale, out = object(), object(), object() + + cp.offer_gate_epilogue(gate, scale, out) + assert cp._gate_epilogue == (gate, scale, out) + assert not cp.take_gate_epilogue() + assert cp._gate_epilogue is None + + cp.offer_gate_epilogue(gate, scale, out) + cp._gate_epilogue_done = True + assert cp.take_gate_epilogue() + assert not cp.take_gate_epilogue() + + def test_runtime_only_slots_are_accepted_but_not_built_by_the_plugin(monkeypatch): requested = frozenset({"gdn_commit", "sharded_gumbel"}) monkeypatch.setattr(runtime, "SLOTS", requested) @@ -233,6 +345,7 @@ def test_slot_skv2_cp_builds_as_a_graph_variant_of_slot_skv2(monkeypatch): from sgl_nvidia_mega.qwen3_5.attention import split_kv monkeypatch.setattr(runtime, "SLOTS", frozenset({"skv2", "skv2_cp"})) + monkeypatch.setattr(runtime, "build_cp_lite", lambda binding: None) monkeypatch.setattr(split_kv, "cp_reason", lambda: "") model = SimpleNamespace( model=SimpleNamespace(start_layer=0, end_layer=0, layers=[]) @@ -399,6 +512,7 @@ def test_the_cp_slots_offer_the_draft_runners_a_cp_variant(monkeypatch): slots = {"draft_loop", "draft_q1", "draft_q1_cp", "draft_extend", "draft_skv2"} monkeypatch.setattr(runtime, "SLOTS", frozenset(slots | {"draft_skv2_cp"})) + monkeypatch.setattr(runtime, "build_cp_lite", lambda binding: None) assert runtime.slot_reasons() == [] built = draft.SPEC.build("lm", SimpleNamespace(context_len=4096), _draft_ctx()) loop = built.graph_variants["draft_loop"] diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 7e09a281..28cb16d8 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -29,8 +29,13 @@ @pytest.fixture(autouse=True) def registered(monkeypatch): + from sglang.srt.kernel_plugins.nvidia import binding as plugin_binding + monkeypatch.setattr(registry, "_REGISTRY", {}) monkeypatch.delenv(plugin.PARITY_SORT_ENV, raising=False) + monkeypatch.setattr( + plugin_binding, "megakernel_agree", lambda *args, **kwargs: None + ) plugin.register() @@ -637,7 +642,12 @@ def test_the_route_slot_is_not_built_for_a_layer_it_cannot_serve(monkeypatch): "sglang.srt.runtime_context.get_parallel", lambda: SimpleNamespace(tp_size=8, moe_ep_size=1), ) - ctx = SimpleNamespace(widths=(7,), tokens_per_request=7, deadline_s=60.0) + ctx = SimpleNamespace( + widths=(7,), + tokens_per_request=7, + deadline_s=60.0, + binding="EagleBinding", + ) want = "layer 0: router gate torch.float16 (512, 4096)" assert route.build(model, 7, fused=False).reason == want with pytest.raises( @@ -1672,7 +1682,7 @@ def _attention_layer(log, x): qwen3_5.Qwen3_5AttentionDecoderLayer ) torch.nn.Module.__init__(layer) - q, k, v, gate, o = (torch.zeros(7, 4) for _ in range(5)) + q, k, v, gate, o = (torch.zeros(x.shape[0], 4) for _ in range(5)) log.known.update(q=q, k=k, v=v, gate=gate, o=o) layer.__dict__.update( attn_output_gate=True, @@ -1772,7 +1782,7 @@ def fused(**kw): yield -def _skv2_setup(log, k_scale=None, **fb): +def _skv2_setup(log, k_scale=None, *, batch_size=1, rows=7, **fb): """Layer 7's RadixAttention (TP8: 4 query heads, 1 KV head of 256), the forward's view of a pool of 3 pages for layers 3 and 7, and the verify step's state.""" from sgl_nvidia_mega.qwen3_5.attention import split_kv @@ -1803,12 +1813,14 @@ def _skv2_setup(log, k_scale=None, **fb): v_scale=(1.0, 0.5), ), attn=SimpleNamespace( - page_table=torch.zeros(1, 4, dtype=torch.int32), - cache_seqlens=torch.full((1,), 130, dtype=torch.int32), + page_table=torch.zeros(batch_size, 4, dtype=torch.int32), + cache_seqlens=torch.full((batch_size,), 130, dtype=torch.int32), ), ) batch = SimpleNamespace( - out_cache_loc=torch.zeros(7, dtype=torch.int64), batch_size=1, **fb + out_cache_loc=torch.zeros(rows, dtype=torch.int64), + batch_size=batch_size, + **fb, ) log.known.update( pool, @@ -2012,6 +2024,62 @@ def attend( assert log[-1][0] == "skv2" +@pytest.mark.parametrize("served", [True, False]) +def test_cp_slots_keep_multi_request_target_verify_in_the_plugin(monkeypatch, served): + from sgl_nvidia_mega.qwen3_5 import runtime + from sgl_nvidia_mega.qwen3_5.attention import split_kv + + requests, rows = 8, 8 * 7 + log = _log() + attn, st = _skv2_setup(log, batch_size=requests, rows=rows) + q8 = torch.zeros(rows, 1024, dtype=torch.float8_e4m3fn) + out = torch.ones(rows, 4, 256, dtype=torch.bfloat16) + log.known["q8"] = q8 + calls = [] + + class CpLite: + def attend( + self, q, kv, pages, lens, q_len, bmm1, bmm2, ws, layer, fb, min_bs=None + ): + calls.append( + ( + tuple(q.shape), + tuple(pages.shape), + tuple(lens.shape), + q_len, + fb is st.fb, + min_bs, + ) + ) + return out if served else None + + monkeypatch.setattr(split_kv, "cp_lite", CpLite) + monkeypatch.setattr(runtime, "SLOTS", frozenset({"skv2_cp"})) + with ( + _skv2_kernels(log, label="split_kv_stock"), + patch("sglang.srt.runtime_context.get_buffer", lambda name, make: None), + ): + got = split_kv.core(attn, q8, st) + + assert calls == [((rows, 4, 256), (requests, 4), (requests,), 7, True, None)] + if served: + assert got.shape == (rows, 1024) and bool((got == 1).all()) + assert not log + else: + assert got == "stock_o" and log[-1][0] == "attention" + + +def test_draft_only_cp_slots_do_not_enable_target_multi_request(monkeypatch): + from sgl_nvidia_mega.qwen3_5 import runtime + from sgl_nvidia_mega.qwen3_5.attention import split_kv + + monkeypatch.setattr(split_kv, "cp_lite", lambda: object()) + monkeypatch.setattr(runtime, "SLOTS", frozenset({"draft_q1_cp", "draft_skv2_cp"})) + assert not split_kv.cp_multi(56) + monkeypatch.setattr(runtime, "SLOTS", frozenset({"skv2_cp"})) + assert split_kv.cp_multi(56) + + @pytest.mark.parametrize( "kw, want", [ @@ -2096,6 +2164,29 @@ def test_the_skv2_slot_runs_the_attention_step_through_the_slot(): assert slot == ["prepare", "skv2_attend", "sigmoid_mul", "o_proj"] +def test_multi_request_cp_routes_the_stock_attention_step_through_the_plugin(): + x = torch.zeros(56, 4) + log = _log(x=x, positions="pos", batch="fb") + layer = _attention_layer(log, x) + mul = log.op("sigmoid_mul", "gated") + st = SimpleNamespace(positions="pos", fb="fb", slots=plan.Slots()) + with ( + patch.object(qwen3_5, "_is_cuda", True), + patch("sglang.kernels.ops.elementwise.elementwise.fused_sigmoid_mul", mul), + patch( + "sgl_nvidia_mega.qwen3_5.attention.split_kv.cp_multi", + return_value=True, + ), + patch( + "sgl_nvidia_mega.qwen3_5.attention.split_kv.attend", + log.op("cp_attend", log.known["o"]), + ), + ): + attention.step(layer, x, st) + + assert [e[0] for e in log] == ["prepare", "cp_attend", "sigmoid_mul", "o_proj"] + + def test_the_executor_offers_the_verify_split_kv_variants_with_slot_skv2(): from sgl_nvidia_mega.qwen3_5.attention import route @@ -2255,17 +2346,17 @@ def symmetric(group, disabled=False): @pytest.mark.parametrize("consumed", [False, True]) def test_the_mka_wide_path_offers_gate_quantization_to_cp_lite(consumed): - from sglang.srt.layers.attention.q35_cp_lite import CpLite from sglang.srt.layers.quantization import fp8_utils - from sgl_nvidia_mega.qwen3_5 import lowm - from sgl_nvidia_mega.qwen3_5.attention import mka + from sgl_nvidia_mega.qwen3_5 import lowm, runtime + from sgl_nvidia_mega.qwen3_5.attention import mka, split_kv from sgl_nvidia_mega.qwen3_5.sm103.attention import gate_quant as gate_module + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite from sgl_nvidia_mega.qwen3_5.sm103.draft import front rows = 21 log = _log() - attn, st = _skv2_setup(log) + attn, st = _skv2_setup(log, batch_size=3, rows=rows) layer = _mka_layer(log, attn) x = torch.zeros(rows, 4096, dtype=torch.bfloat16) qkv = torch.zeros(rows, 2560, dtype=torch.bfloat16) @@ -2301,6 +2392,7 @@ def row_parallel(linear, o_proj_input, dtype, what, hop): with ( patch.object(CpLite, "_instance", GateEpilogue()), + patch.object(runtime, "SLOTS", frozenset({"skv2_cp"})), patch.object( front, "fused_qk_gemma_rmsnorm_rope_gate", @@ -2310,6 +2402,7 @@ def row_parallel(linear, o_proj_input, dtype, what, hop): patch.object(fp8_utils, "flashinfer_bmm_fp8", create=True), patch.object(lowm, "take_fp8", return_value=None), patch.object(lowm, "row_parallel", row_parallel), + patch.object(split_kv, "core", return_value=attention_out), ): out = mka.step(layer, x, st) From c7ea36c349586748b90de8b21e6483596ca33e5a Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 27 Oct 2026 21:57:45 -0700 Subject: [PATCH 30/62] fix: keep CP build compatible with binding paths --- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 35812268..15453ca6 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -126,7 +126,10 @@ def build_cp_lite(binding) -> None: rank, so this runs once from each Qwen3.5 plugin build, before any package checks ``CpLite._instance`` or CUDA graph variant is created. """ - from sglang.srt.kernel_plugins.nvidia.binding import megakernel_agree + try: + from sglang.srt.kernel_plugins.nvidia.binding import megakernel_agree + except ImportError: # pragma: no cover - the pre-kernel_plugins SGLang path + from sglang.srt.megakernels.runtime import megakernel_agree requested = tuple(sorted(SLOTS & CP_LITE_SLOTS)) megakernel_agree( @@ -365,7 +368,9 @@ def slot_reasons() -> list[str]: "tp_size", "tp_device_group", ), - "srt/kernel_plugins/nvidia/binding.py": ("megakernel_agree",), + "srt/kernel_plugins/nvidia/binding.py|srt/megakernels/runtime.py": ( + "megakernel_agree", + ), } From 3953297d57ff34a14c6248e8410b2e31e0fb982c Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 27 Oct 2026 22:01:42 -0700 Subject: [PATCH 31/62] test: pin runtime CP removal mirror --- sgl_nvidia_mega/qwen3_5/attention/split_kv.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py index 5aaf9df4..ccbd333c 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py +++ b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py @@ -85,6 +85,7 @@ "1445b85387928caf", "d7e7c1fe125739e6", "066a1e0aa5bd6e5c", + "10bd9120d713848b", ), "srt/layers/attention/trtllm_mha_backend.py:TRTLLMHAAttnBackend._should_use_fused_fp8_path": "e0b14791d523a019", "srt/layers/attention/trtllm_mha_backend.py:TRTLLMHAAttnBackend._fused_fp8_qkv_kv_cache": ( From 38b8a0d85506a58612db832a44caef7b00243ffa Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 27 Oct 2026 22:40:21 -0700 Subject: [PATCH 32/62] test: pin CP mirror digest for serving Python --- sgl_nvidia_mega/qwen3_5/attention/split_kv.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py index ccbd333c..d1859149 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py +++ b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py @@ -86,6 +86,7 @@ "d7e7c1fe125739e6", "066a1e0aa5bd6e5c", "10bd9120d713848b", + "37c7c1fa55956740", ), "srt/layers/attention/trtllm_mha_backend.py:TRTLLMHAAttnBackend._should_use_fused_fp8_path": "e0b14791d523a019", "srt/layers/attention/trtllm_mha_backend.py:TRTLLMHAAttnBackend._fused_fp8_qkv_kv_cache": ( From a2880f1f64ae1c1dde0d4b94edf478e955fd0ed3 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 27 Oct 2026 23:10:44 -0700 Subject: [PATCH 33/62] refactor: move qwen35 verify and moe policies to plugin --- sgl_nvidia_mega/qwen3_5/draft/__init__.py | 1 + sgl_nvidia_mega/qwen3_5/draft/finalize.py | 69 +++ sgl_nvidia_mega/qwen3_5/draft/parts.py | 11 + sgl_nvidia_mega/qwen3_5/moe/route.py | 99 +++- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 10 +- .../qwen3_5/sm103/moe/shared_act_quant.py | 100 ++++ .../qwen3_5/sm103/sharded_gumbel.py | 247 ++++++++++ sgl_nvidia_mega/qwen3_5/spec.py | 6 + test/kernels/test_qwen3_5_moe.py | 96 ++++ test/kernels/test_qwen3_5_sharded_gumbel.py | 433 ++++++++++++++++++ test/unit/test_qwen3_5_runtime.py | 66 ++- test/unit/test_qwen3_5_spec.py | 16 + test/unit/test_qwen3_5_steps.py | 257 ++++++++++- 13 files changed, 1380 insertions(+), 31 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/draft/finalize.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/moe/shared_act_quant.py create mode 100644 sgl_nvidia_mega/qwen3_5/sm103/sharded_gumbel.py create mode 100644 test/kernels/test_qwen3_5_sharded_gumbel.py diff --git a/sgl_nvidia_mega/qwen3_5/draft/__init__.py b/sgl_nvidia_mega/qwen3_5/draft/__init__.py index 46c38fe8..15883065 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/draft/__init__.py @@ -86,6 +86,7 @@ def numerics(slots: frozenset) -> Numerics: "draft_q1", "draft_xchg", "draft_hop", + "draft_fin", "draft_front", "draft_skv2", }: diff --git a/sgl_nvidia_mega/qwen3_5/draft/finalize.py b/sgl_nvidia_mega/qwen3_5/draft/finalize.py new file mode 100644 index 00000000..20fa5b5a --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/draft/finalize.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Plugin-owned eligibility for the Qwen3.5 MTP draft's deferred MoE finalize. + +SGLang owns the model-independent deferred-finalize protocol. Slot ``draft_fin`` +only selects that protocol for the one Qwen3.5 draft shape whose producer and +consumer have been validated. Keeping the selection here avoids a Qwen-specific +policy branch in SGLang's stock MoE layer. +""" + +from __future__ import annotations + +import torch + +# The SGLang names this module reads (test/unit/test_contract.py pins them). +SGLANG_READS = { + "srt/layers/quantization/unquant.py": ("UnquantizedFusedMoEMethod",), + "srt/layers/moe/fused_moe_triton/layer.py": ( + "is_draft", + "quant_method", + "params_dtype", + "hidden_size", + "num_experts", + "top_k", + "runner_backend", + "supports_deferred_finalize", + ), + "srt/models/qwen2_moe.py": ( + "experts", + "shared_expert", + ), +} + + +def enable(mlp) -> str: + """Enable the generic producer on this exact draft block. + + Returns an empty string when the block can serve ``draft_fin`` and a stable + refusal reason otherwise. The mutation is local to the loaded draft model's + routed-expert module; it changes no process-wide or stock-layer policy. + """ + from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod + + experts = mlp.experts + runner = getattr(experts, "runner_backend", None) + checks = ( + (not getattr(experts, "is_draft", False), "not a draft MoE"), + ( + not isinstance(experts.quant_method, UnquantizedFusedMoEMethod), + f"quant method {type(experts.quant_method).__name__}", + ), + ( + experts.params_dtype != torch.bfloat16, + f"params dtype {experts.params_dtype}", + ), + (experts.hidden_size != 4096, f"hidden size {experts.hidden_size}"), + (experts.num_experts != 512, f"experts {experts.num_experts}"), + (experts.top_k != 10, f"top-k {experts.top_k}"), + (mlp.shared_expert is None, "no separate shared expert"), + ( + runner is None or not runner.is_flashinfer_trtllm(), + f"MoE runner {runner}", + ), + ) + for refused, reason in checks: + if refused: + return reason + experts.supports_deferred_finalize = True + return "" diff --git a/sgl_nvidia_mega/qwen3_5/draft/parts.py b/sgl_nvidia_mega/qwen3_5/draft/parts.py index 9effb0b5..66061faf 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/parts.py +++ b/sgl_nvidia_mega/qwen3_5/draft/parts.py @@ -90,6 +90,17 @@ def build(cls, draft) -> Parts: parts = cls(draft, layers[0]) parts.attention = Attention(parts) slots = parts.slots + if "draft_fin" in slots: + from sgl_nvidia_mega.qwen3_5.draft import finalize + + reason = finalize.enable(parts.layer.mlp) + if reason: + logger.warning( + "NVIDIA megakernel: slot draft_fin cannot run (%s); " + "SGLang finalizes the draft MoE.", + reason, + ) + slots = parts.slots = slots - {"draft_fin"} if "draft_moe" in slots: parts.moe = parts._build_moe() if parts.moe is None: diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index 2431657c..0e4e4841 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -68,8 +68,6 @@ "act_fn", "_enable_silu_fp4_quant_fusion", "_silu_fp4_quant_fused", - "_silu_fp8_quant_fusion", - "_down_takes_static_fp8", ), "srt/models/qwen3_5.py": ("layers", "start_layer", "end_layer", "mlp"), "srt/layers/moe/topk.py": ( @@ -115,7 +113,6 @@ '"weight_scale"', ), "srt/layers/activation.py": ("SiluAndMul",), - "kernels/ops/quantization/fp8_kernel.py": ("silu_and_mul_static_quant_fp8",), } # The SGLang functions whose body this module mirrors, by the digest of their code @@ -123,7 +120,10 @@ SGLANG_MIRRORS = { "srt/layers/moe/fused_moe_triton/layer.py:FusedMoE.forward_deferred_finalize": "af4b474cbaf33493", "srt/layers/moe/fused_moe_triton/layer.py:FusedMoE._dispatch_with_pre_quant": "dcdefc28cf643081", - "srt/models/qwen2_moe.py:Qwen2MoeMLP.forward": "bb4c9fb2d7994e97", + "srt/models/qwen2_moe.py:Qwen2MoeMLP.forward": ( + "bb4c9fb2d7994e97", # runtime v2: temporary core shared-act fast path + "6091ffb613568a10", # runtime v3: the policy and kernel live in this slot + ), } @@ -187,6 +187,13 @@ def build( return Bound(reason=f"route_shared: layer {i}: {reason}") device = torch.cuda.current_device() router.build(device, max_rows, shared=shared) + if shared: + from sgl_nvidia_mega.qwen3_5.sm103.moe import shared_act_quant + + for columns in sorted( + {int(mlp.shared_expert.gate_up_proj.weight.shape[1]) for mlp in layers} + ): + shared_act_quant.build(device, columns) if not fused: return Bound(route=True, shared=shared) for i, mlp in zip(range(m.start_layer, m.end_layer), layers): @@ -369,23 +376,40 @@ def shared_unsupported(mlp) -> str: if not isinstance(expert, Qwen2MoeMLP) or mlp.shared_expert_gate is None: return "no gated Qwen2MoeMLP shared expert" gu = expert.gate_up_proj - method = gu.quant_method + method = getattr(gu, "quant_method", None) + weight = getattr(gu, "weight", None) + input_scale = getattr(gu, "input_scale", None) + weight_scale = getattr(gu, "weight_scale", None) if ( not isinstance(method, ModelOptFp8LinearMethod) or method.use_marlin or method.use_sm120_fp8 - or gu.bias is not None - or gu.weight.dtype != torch.float8_e4m3fn - or gu.weight.shape[0] != router.HIDDEN - or gu.weight.shape[1] % 8 - or gu.weight.shape[1] > router.SHARED_MAX_N - or not gu.weight.t().is_contiguous() - or gu.input_scale.numel() != 1 - or gu.weight_scale.numel() != 1 + or getattr(gu, "bias", None) is not None + or weight is None + or not weight.is_cuda + or weight.dtype != torch.float8_e4m3fn + or weight.shape[0] != router.HIDDEN + or weight.shape[1] % 8 + or weight.shape[1] > router.SHARED_MAX_N + or not weight.t().is_contiguous() + or input_scale is None + or not input_scale.is_cuda + or input_scale.device != weight.device + or input_scale.dtype != torch.float32 + or tuple(input_scale.shape) != (1,) + or not input_scale.is_contiguous() + or weight_scale is None + or not weight_scale.is_cuda + or weight_scale.device != weight.device + or weight_scale.dtype != torch.float32 + or tuple(weight_scale.shape) != (1,) + or not weight_scale.is_contiguous() ): return "the shared gate/up is not a static per-tensor FP8 linear of <= 512 rows" if not isinstance(expert.act_fn, SiluAndMul): return f"shared expert activation {type(expert.act_fn).__name__}" + if not _down_takes_static_fp8(expert): + return "the shared down projection does not accept static per-tensor FP8" return "" @@ -394,22 +418,57 @@ def shared_expert_tail(expert, gate_up: torch.Tensor) -> torch.Tensor: if expert._enable_silu_fp4_quant_fusion and not isinstance(gate_up, tuple): x, _ = expert.down_proj(expert._silu_fp4_quant_fused(gate_up)) return x - if expert._silu_fp8_quant_fusion is None: - expert._silu_fp8_quant_fusion = expert._down_takes_static_fp8() - if expert._silu_fp8_quant_fusion and gate_up.dtype == torch.bfloat16: - from sglang.kernels.ops.quantization.fp8_kernel import ( + if _down_takes_static_fp8(expert) and gate_up.dtype == torch.bfloat16: + from sgl_nvidia_mega.qwen3_5.sm103.moe.shared_act_quant import ( silu_and_mul_static_quant_fp8, ) - x, _ = expert.down_proj( - silu_and_mul_static_quant_fp8(gate_up, expert.down_proj.input_scale) - ) + quantized = silu_and_mul_static_quant_fp8(gate_up, expert.down_proj.input_scale) + key = ("route_shared_act_quant", gate_up.shape[0]) + if key not in _ENGAGED: + _ENGAGED.add(key) + logger.info( + "NVIDIA megakernel: Qwen3.5 slot route_shared shared_act_quant " + "runs at %d rows", + gate_up.shape[0], + ) + x, _ = expert.down_proj(quantized) return x x = expert.act_fn(gate_up) x, _ = expert.down_proj(x) return x +def _down_takes_static_fp8(expert) -> bool: + """Whether this shared expert's down projection accepts static E4M3 input. + + This is the Qwen3.5 slot policy that runtime-v2 temporarily placed on + ``Qwen2MoeMLP``. Keeping it beside the plugin-owned activation kernel lets + runtime-v3 restore the stock model class without leaving a private attribute + dependency here. + """ + from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp8LinearMethod + + down = expert.down_proj + method = getattr(down, "quant_method", None) + weight = getattr(down, "weight", None) + scale = getattr(down, "input_scale", None) + return ( + isinstance(method, ModelOptFp8LinearMethod) + and not method.use_marlin + and not method.use_sm120_fp8 + and weight is not None + and weight.is_cuda + and weight.dtype == torch.float8_e4m3fn + and scale is not None + and scale.is_cuda + and scale.device == weight.device + and scale.dtype == torch.float32 + and tuple(scale.shape) == (1,) + and scale.is_contiguous() + ) + + def routed( mlp, x: torch.Tensor, diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 15453ca6..882516e5 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -86,9 +86,9 @@ "draft_skv2_cp", ) CP_LITE_SLOTS = frozenset({"skv2_cp", "draft_q1_cp", "draft_skv2_cp"}) -# SGLang consumes these names from the shared slot policy itself. They are valid -# requests, but are intentionally not built or dispatched by this plugin. -RUNTIME_ONLY_SLOTS: tuple[str, ...] = ("gdn_commit", "sharded_gumbel") +# Provider callbacks consume these names from the shared slot policy. They do not +# build a target-verify width of their own. +HOOK_SLOTS: tuple[str, ...] = ("sharded_gumbel",) SLOTS_ENV = "SGLANG_NVIDIA_MEGAKERNEL_SLOTS" # These slots run all 512 experts themselves. Under expert parallelism, keep the # route-only and boundary slots but let SGLang's EP path run the routed experts. @@ -183,7 +183,7 @@ def slot_reasons() -> list[str]: *OPT_IN_VARIANTS, *DRAFT_ENTRIES, *DRAFT_SLOTS, - *RUNTIME_ONLY_SLOTS, + *HOOK_SLOTS, ) ) ) @@ -479,7 +479,7 @@ def build(model, ctx): - set(DRAFT_SLOTS) - set(OPT_IN_COVERAGE) - set(OPT_IN_VARIANTS) - - set(RUNTIME_ONLY_SLOTS) + - set(HOOK_SLOTS) ) # mkm's 16-token engine has nothing to serve where no covered step is wider than 7 rows vacuous = ( diff --git a/sgl_nvidia_mega/qwen3_5/sm103/moe/shared_act_quant.py b/sgl_nvidia_mega/qwen3_5/sm103/moe/shared_act_quant.py new file mode 100644 index 00000000..467f572b --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/moe/shared_act_quant.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Adapted from SGLang (Apache-2.0, SGLang Team): +# python/sglang/kernels/ops/quantization/fp8_kernel.py +"""The shared expert's BF16 SwiGLU boundary and static FP8 quantization. + +This is the activation tail of slot ``route_shared``. Keeping it in the plugin +means the slot no longer calls a Qwen3.5-specific SGLang kernel after the fused +router/shared projection. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _silu_and_mul_static_quant_fp8( + x_ptr, + y_q_ptr, + y_s_ptr, + width, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + cols = tl.arange(0, BLOCK) + mask = cols < width + tl.extra.cuda.gdc_wait() + gate = tl.load(x_ptr + row * 2 * width + cols, mask=mask, other=0.0).to(tl.float32) + up = tl.load(x_ptr + row * 2 * width + width + cols, mask=mask, other=0.0).to( + tl.float32 + ) + scale = tl.load(y_s_ptr).to(tl.float32) + tl.extra.cuda.gdc_launch_dependents() + # Match the BF16 activation boundary before the static FP8 quantization. + y = (gate / (1.0 + tl.exp(-gate)) * up).to(tl.bfloat16).to(tl.float32) + q = tl.clamp(y * (1.0 / scale), -448.0, 448.0).to(tl.float8e4nv) + tl.store( + y_q_ptr + row * width + cols, + q.to(tl.uint8, bitcast=True), + mask=mask, + ) + + +def build(device_index: int, gate_up_columns: int) -> None: + """Compile the exact shared-expert tail before CUDA-graph capture.""" + if gate_up_columns <= 0 or gate_up_columns % 2: + raise ValueError( + f"shared gate/up columns must be positive and even, got {gate_up_columns}" + ) + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("compiling shared act+quant inside CUDA-graph capture") + with torch.cuda.device(device_index): + gate_up = torch.empty((1, gate_up_columns), dtype=torch.bfloat16, device="cuda") + scale = torch.ones(1, dtype=torch.float32, device="cuda") + silu_and_mul_static_quant_fp8(gate_up, scale) + torch.cuda.synchronize(device_index) + + +def silu_and_mul_static_quant_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Return the shared expert activation as static per-tensor E4M3 rows.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_contiguous() + or x.shape[1] <= 0 + or x.shape[1] % 2 + ): + raise ValueError( + "shared act+quant input must be contiguous CUDA BF16 [rows, 2 * width]" + ) + if ( + scale.device != x.device + or scale.dtype != torch.float32 + or tuple(scale.shape) != (1,) + or not scale.is_contiguous() + ): + raise ValueError( + "shared act+quant scale must be contiguous CUDA FP32 [1] " + "on the input device" + ) + + rows, gate_up_columns = x.shape + width = gate_up_columns // 2 + out = torch.empty((rows, width), device=x.device, dtype=torch.float8_e4m3fn) + block = triton.next_power_of_2(width) + _silu_and_mul_static_quant_fp8[(rows,)]( + x, + out.view(torch.uint8), + scale, + width, + BLOCK=block, + num_warps=min(max(block // 256, 1), 8), + num_stages=1, + launch_pdl=True, + ) + return out diff --git a/sgl_nvidia_mega/qwen3_5/sm103/sharded_gumbel.py b/sgl_nvidia_mega/qwen3_5/sm103/sharded_gumbel.py new file mode 100644 index 00000000..256eafb3 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/sm103/sharded_gumbel.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Adapted from SGLang (Apache-2.0, SGLang Team): +# python/sglang/srt/speculative/sharded_verify.py (_shard_key_kernel, shard_keys) +"""Gumbel-max keys for one rank's vocabulary shard. + +The generic runtime owns the shard descriptor, cross-rank key exchange, greedy path, +and stock full-logits fallback. This module owns only the numerical opt-in: for each +row and local vocabulary range it writes the best 64-bit key after temperature scaling +and deterministic Gumbel noise. Noise is keyed by ``(seed, global token id)``, so a +row has the same winner whether its vocabulary is evaluated whole, in TP shards, or in +multiple programs. + +The key layout is the one consumed by SGLang's generic sharded-verify reducer:: + + ordered_fp32_bits << 32 | (0xffffffff - global_token_id) + +Thus an int64 maximum chooses the largest sampled value and breaks exact ties toward +the smaller token id. This is a direct Triton migration of the v2 kernel, including +its RNG stream and FP32 operation order. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +__all__ = ["build", "gumbel_shard_keys", "key_splits"] + +# Key of (-inf, token id 0xffffffff), below every real column's key. Keep this +# signed representation bit-for-bit identical to runtime-v2's generic reducer. +_KEY_FLOOR = -2139095041 << 32 +_KEY_BLOCK = 2048 + + +@triton.jit +def _gumbel_shard_key_kernel( + logits_ptr, + stride_row, + temperature_ptr, + seed_ptr, + keys_ptr, + vocab_start, + valid, + vocab_size, + span, + KEY_FLOOR: tl.constexpr, + BLOCK: tl.constexpr, +): + # Program (row, part) covers [part * span, (part + 1) * span) of this shard. + row = tl.program_id(0).to(tl.int64) + part = tl.program_id(1) + lo = part * span + hi = tl.minimum(lo + span, valid) + best = tl.full([BLOCK], KEY_FLOOR, tl.int64) + temperature = tl.load(temperature_ptr + row) + inverse_temperature = 1.0 / tl.maximum(temperature, 1e-20) + seed = tl.load(seed_ptr + row) + for offset in range(lo, hi, BLOCK): + columns = offset + tl.arange(0, BLOCK) + token_ids = vocab_start + columns + mask = (columns < hi) & (token_ids < vocab_size) + values = tl.load( + logits_ptr + row * stride_row + columns, + mask=mask, + other=float("-inf"), + ).to(tl.float32) + # Mixed batches retain runtime-v2's behavior: T <= 0 rows use their plain + # logits while positive-temperature rows consume the deterministic stream. + if temperature > 0.0: + uniform = tl.rand(seed, token_ids) + values = values * inverse_temperature - tl.log(-tl.log(uniform)) + # NaNs never beat a finite value. Adding +0 canonicalizes -0 for argmax ties. + not_nan = values == values # noqa: PLR0124 -- intentional IEEE NaN test + values = tl.where(not_nan, values, float("-inf")) + 0.0 + bits = values.to(tl.int32, bitcast=True) + ordered = tl.where(bits < 0, bits ^ 0x7FFFFFFF, bits).to(tl.int64) + keys = (ordered << 32) | (4294967295 - token_ids.to(tl.int64)) + keys = tl.where(mask, keys, KEY_FLOOR) + best = tl.maximum(best, keys) + tl.store( + keys_ptr + row * tl.num_programs(1) + part, + tl.max(best, axis=0), + ) + + +def key_splits(valid: int) -> int: + """Programs per row, padded even for the runtime's multimem chunk layout.""" + if valid <= 0: + raise ValueError(f"valid vocabulary columns must be positive, got {valid}") + count = -(-valid // _KEY_BLOCK) + return count + (count & 1) + + +def _launch( + logits: torch.Tensor, + temperatures: torch.Tensor, + seeds: torch.Tensor, + out: torch.Tensor, + *, + vocab_start: int, + valid: int, + vocab_size: int, + splits: int, +) -> None: + rows = logits.shape[0] + if rows == 0: + return + _gumbel_shard_key_kernel[(rows, splits)]( + logits, + logits.stride(0), + temperatures, + seeds, + out, + vocab_start, + valid, + vocab_size, + -(-valid // splits), + KEY_FLOOR=_KEY_FLOOR, + BLOCK=_KEY_BLOCK, + num_warps=8, + ) + + +def build( + device_index: int, + *, + columns: int, + vocab_start: int, + valid: int, + vocab_size: int, + splits: int, +) -> None: + """Compile both dtypes for the rank's exact production shard metadata.""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("compiling sharded Gumbel inside CUDA-graph capture") + if not 0 < valid <= columns: + raise ValueError( + f"valid vocabulary columns must be in [1, {columns}], got {valid}" + ) + if vocab_start < 0 or vocab_start + valid > vocab_size: + raise ValueError( + f"invalid vocabulary shard: start={vocab_start}, valid={valid}, " + f"size={vocab_size}" + ) + if splits <= 0: + raise ValueError(f"splits must be positive, got {splits}") + with torch.cuda.device(device_index): + temperatures = torch.ones(2, dtype=torch.float32, device="cuda") + seeds = torch.arange(2, dtype=torch.int64, device="cuda") + shape = (2, splits) if splits > 1 else (2,) + keys = torch.empty(shape, dtype=torch.int64, device="cuda") + for dtype in (torch.bfloat16, torch.float32): + logits = torch.empty((2, columns), dtype=dtype, device="cuda") + _launch( + logits, + temperatures, + seeds, + keys, + vocab_start=vocab_start, + valid=valid, + vocab_size=vocab_size, + splits=splits, + ) + torch.cuda.synchronize(device_index) + + +def gumbel_shard_keys( + logits: torch.Tensor, + temperatures: torch.Tensor, + seeds: torch.Tensor, + *, + vocab_start: int, + valid: int, + vocab_size: int, + splits: int = 1, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Return this rank's best Gumbel key for every row and split. + + ``logits`` is contiguous by column and may contain padding after ``valid``. + ``temperatures`` is FP32 and ``seeds`` is int64, both one value per row. Positive + temperatures sample ``logits / T``; nonpositive temperatures retain the plain + argmax key, which supports mixed graph batches without another launch. Columns + at or beyond ``valid`` or ``vocab_size`` are masked even when their padding holds a + large value. Pass a preallocated ``out`` to avoid an allocation in a captured + graph; call :func:`build` before the first capture. + """ + if not logits.is_cuda or logits.dim() != 2 or logits.stride(1) != 1: + raise ValueError( + "sharded Gumbel logits must be a CUDA [rows, columns] tensor with " + "unit column stride" + ) + if logits.dtype not in (torch.bfloat16, torch.float32): + raise ValueError( + f"sharded Gumbel logits must be BF16 or FP32, got {logits.dtype}" + ) + rows, columns = logits.shape + for name, tensor, dtype in ( + ("temperatures", temperatures, torch.float32), + ("seeds", seeds, torch.int64), + ): + if ( + tensor.device != logits.device + or tensor.dtype != dtype + or tuple(tensor.shape) != (rows,) + or not tensor.is_contiguous() + ): + raise ValueError( + f"sharded Gumbel {name} must be contiguous {dtype} [{rows}] " + f"on {logits.device}, got {tensor.dtype} {tuple(tensor.shape)} " + f"on {tensor.device}" + ) + if not 0 < valid <= columns: + raise ValueError( + f"valid vocabulary columns must be in [1, {columns}], got {valid}" + ) + if vocab_start < 0 or vocab_size <= vocab_start: + raise ValueError( + f"invalid vocabulary range: start={vocab_start}, size={vocab_size}" + ) + if splits <= 0: + raise ValueError(f"splits must be positive, got {splits}") + shape = (rows, splits) if splits > 1 else (rows,) + if out is None: + out = torch.empty(shape, dtype=torch.int64, device=logits.device) + elif ( + out.device != logits.device + or out.dtype != torch.int64 + or tuple(out.shape) != shape + or not out.is_contiguous() + ): + raise ValueError( + f"sharded Gumbel out must be contiguous int64 {shape} on {logits.device}" + ) + _launch( + logits, + temperatures, + seeds, + out, + vocab_start=vocab_start, + valid=valid, + vocab_size=vocab_size, + splits=splits, + ) + return out diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index fa28d83c..45860645 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -112,6 +112,12 @@ # outputs are held to stock or fp64 on the same quantized operands (test/kernels), at # most these multiples of the error of SGLang's own op. NUMERIC_TOLERANCES = { + # Gumbel-max consumes a deterministic per-token stream different from the stock + # sampler. It is categorical-distribution equivalent, not tokenwise close; the + # migrated chi-square test pools low-expectation bins and enforces this p-value. + "sharded_gumbel": { + "sharded_gumbel.categorical.chi_square_p_value_min": 1e-4, + }, # Four router GEMV chains sum in another order than SGLang's cuBLAS GEMM. "route": {"route.router_logits.bf16_ulp_vs_stock": 1.0}, # The tiled router front stays within one BF16 ULP of SGLang's whole-M GEMM. diff --git a/test/kernels/test_qwen3_5_moe.py b/test/kernels/test_qwen3_5_moe.py index 6507e939..afb2f272 100644 --- a/test/kernels/test_qwen3_5_moe.py +++ b/test/kernels/test_qwen3_5_moe.py @@ -28,6 +28,8 @@ (MoE TP2 x EP8: 64 local experts, intermediate 512 a rank, ranks 0 and 5): the weights, the slots routed outside the rank (-1) and every local slot's FC2 row bit for bit.""" +from types import SimpleNamespace + import pytest import torch @@ -434,6 +436,100 @@ def test_router_gemm_quant_shared_replays_from_a_cuda_graph(): assert all(torch.equal(a, b) for a, b in zip(eager, captured)) +@pytest.mark.parametrize("m", [1, 7, 8]) +@pytest.mark.parametrize("width", [128, 256]) +def test_shared_swiglu_static_quant_is_sglangs_bytes(m, width): + from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 + from sglang.srt.layers.activation import SiluAndMul + + from sgl_nvidia_mega.qwen3_5.sm103.moe.shared_act_quant import ( + silu_and_mul_static_quant_fp8, + ) + + g = torch.Generator(device="cuda").manual_seed(73 + m + width) + gate_up = torch.randn( + m, 2 * width, device="cuda", dtype=torch.bfloat16, generator=g + ) + # Exercise both saturation edges as well as ordinary activation values. + gate_up[0, :4] = torch.tensor( + [-80.0, -8.0, 8.0, 80.0], device="cuda", dtype=torch.bfloat16 + ) + scale = torch.tensor([0.003], device="cuda") + # Keep the reference independent of the fused migration, which runtime-v3 + # removes: this is the stock BF16 activation boundary followed by its static + # per-tensor FP8 quantizer. + want, _ = static_quant_fp8(SiluAndMul()(gate_up), scale) + got = silu_and_mul_static_quant_fp8(gate_up, scale) + assert torch.equal(got.view(torch.uint8), want.view(torch.uint8)) + + +def test_shared_swiglu_static_quant_replays_from_a_cuda_graph(): + from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 + from sglang.srt.layers.activation import SiluAndMul + + from sgl_nvidia_mega.qwen3_5.sm103.moe.shared_act_quant import ( + build, + silu_and_mul_static_quant_fp8, + ) + + gate_up = torch.randn(8, 512, device="cuda", dtype=torch.bfloat16) + scale = torch.tensor([0.003], device="cuda") + with pytest.raises(ValueError, match=r"FP32 \[1\]"): + silu_and_mul_static_quant_fp8(gate_up, scale.view(1, 1)) + with pytest.raises(ValueError, match="input device"): + silu_and_mul_static_quant_fp8(gate_up, scale.cpu()) + with pytest.raises(ValueError, match=r"\[rows, 2 \* width\]"): + silu_and_mul_static_quant_fp8(gate_up.view(2, 4, 512), scale) + build(gate_up.device.index, gate_up.shape[1]) + expected, _ = static_quant_fp8(SiluAndMul()(gate_up), scale) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = silu_and_mul_static_quant_fp8(gate_up, scale) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(expected.view(torch.uint8), captured.view(torch.uint8)) + + +def test_shared_swiglu_static_quant_feeds_modelopt_flashinfer_down_projection(): + from sglang.srt.layers.activation import SiluAndMul + from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp8LinearMethod + + from sgl_nvidia_mega.qwen3_5.sm103.moe.shared_act_quant import ( + silu_and_mul_static_quant_fp8, + ) + + generator = torch.Generator(device="cuda").manual_seed(91) + rows, intermediate, hidden = 8, 256, 4096 + gate_up = ( + torch.randn( + rows, + 2 * intermediate, + device="cuda", + dtype=torch.bfloat16, + generator=generator, + ) + * 2 + ) + weight = ( + torch.randn(hidden, intermediate, device="cuda", generator=generator) * 0.05 + ).to(torch.float8_e4m3fn) + input_scale = torch.tensor([0.02], device="cuda") + layer = SimpleNamespace( + use_flashinfer_bmm=True, + weight=weight.t(), + weight_scale=torch.tensor([0.01], device="cuda"), + input_scale=input_scale, + ) + method = object.__new__(ModelOptFp8LinearMethod) + method.use_marlin = method.use_sm120_fp8 = False + + expected = method.apply(layer, SiluAndMul()(gate_up)) + actual = method.apply(layer, silu_and_mul_static_quant_fp8(gate_up, input_scale)) + + assert actual.dtype == torch.bfloat16 + assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)) + + # ---------------------------------------------------------------- slot route_topk diff --git a/test/kernels/test_qwen3_5_sharded_gumbel.py b/test/kernels/test_qwen3_5_sharded_gumbel.py new file mode 100644 index 00000000..3549304c --- /dev/null +++ b/test/kernels/test_qwen3_5_sharded_gumbel.py @@ -0,0 +1,433 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Direct tests of the Qwen3.5 sharded-Gumbel numerical slot.""" + +from __future__ import annotations + +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch +import triton +import triton.language as tl + +from sgl_nvidia_mega.qwen3_5.sm103.sharded_gumbel import ( + build, + gumbel_shard_keys, + key_splits, +) + +_KEY_FLOOR = -2139095041 << 32 + + +@triton.jit +def _per_token_key_oracle( + logits, + stride_row, + temperatures, + seeds, + keys, + columns, + vocab_start, + valid, + vocab_size, + KEY_FLOOR: tl.constexpr, + BLOCK: tl.constexpr, +): + """Independent elementwise oracle: no producer-sized blocks or device reduction.""" + row = tl.program_id(0).to(tl.int64) + column = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + token_id = vocab_start + column + in_tensor = column < columns + in_vocab = (column < valid) & (token_id < vocab_size) + value = tl.load( + logits + row * stride_row + column, + mask=in_tensor & in_vocab, + other=float("-inf"), + ).to(tl.float32) + temperature = tl.load(temperatures + row) + if temperature > 0.0: + uniform = tl.rand(tl.load(seeds + row), token_id) + value = value * (1.0 / tl.maximum(temperature, 1e-20)) - tl.log( + -tl.log(uniform) + ) + value = ( + tl.where( + value == value, # noqa: PLR0124 -- IEEE NaN test + value, + float("-inf"), + ) + + 0.0 + ) + bits = value.to(tl.int32, bitcast=True) + ordered = tl.where(bits < 0, bits ^ 0x7FFFFFFF, bits).to(tl.int64) + key = (ordered << 32) | (4294967295 - token_id.to(tl.int64)) + key = tl.where(in_vocab, key, KEY_FLOOR) + tl.store(keys + row * columns + column, key, mask=in_tensor) + + +def _oracle_keys( + logits, + temperatures, + seeds, + *, + vocab_start, + valid, + vocab_size, + splits, +): + rows, columns = logits.shape + per_token = torch.empty((rows, columns), dtype=torch.int64, device=logits.device) + block = 256 + _per_token_key_oracle[(rows, triton.cdiv(columns, block))]( + logits, + logits.stride(0), + temperatures, + seeds, + per_token, + columns, + vocab_start, + valid, + vocab_size, + KEY_FLOOR=_KEY_FLOOR, + BLOCK=block, + num_warps=4, + ) + if splits == 1: + return per_token.amax(dim=1) + span = -(-valid // splits) + parts = [] + for part in range(splits): + lo, hi = part * span, min((part + 1) * span, valid) + parts.append( + per_token[:, lo:hi].amax(dim=1) + if lo < hi + else torch.full( + (rows,), _KEY_FLOOR, dtype=torch.int64, device=logits.device + ) + ) + return torch.stack(parts, dim=1) + + +def _tokens(keys: torch.Tensor) -> torch.Tensor: + if keys.dim() == 2: + keys = keys.amax(dim=1) + return 0xFFFFFFFF - (keys & 0xFFFFFFFF) + + +def test_key_splits_keeps_multimem_rows_even(): + assert [key_splits(n) for n in (1, 2048, 2049, 4096, 4097, 8192)] == [ + 2, + 2, + 2, + 2, + 4, + 4, + ] + with pytest.raises(ValueError, match="positive"): + key_splits(0) + + +def test_build_prewarms_bf16_and_fp32(monkeypatch): + launches = [] + synchronized = [] + + def fake_tensor(*args, dtype, device, **kwargs): + assert device == "cuda" + return SimpleNamespace(dtype=dtype) + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr(torch.cuda, "device", lambda device: nullcontext()) + monkeypatch.setattr(torch.cuda, "synchronize", synchronized.append) + monkeypatch.setattr(torch, "empty", fake_tensor) + monkeypatch.setattr(torch, "ones", fake_tensor) + monkeypatch.setattr(torch, "arange", fake_tensor) + monkeypatch.setattr( + "sgl_nvidia_mega.qwen3_5.sm103.sharded_gumbel._launch", + lambda logits, *args, **kwargs: launches.append(logits.dtype), + ) + + build( + 3, + columns=31040, + vocab_start=31040, + valid=31040, + vocab_size=248320, + splits=16, + ) + + assert launches == [torch.bfloat16, torch.float32] + assert synchronized == [3] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_padding_nan_and_ties_follow_runtime_v2_argmax_semantics(): + # gid 15 and columns 6: are padding. Their large values must never win. + logits = torch.tensor( + [ + [float("nan"), 5.0, -0.0, 5.0, -2.0, 1000.0, 2000.0, 2000.0], + [float("nan"), 1.0, 3.0, 2.0, 0.0, 1000.0, 2000.0, 2000.0], + [ + float("nan"), + float("nan"), + float("nan"), + float("nan"), + float("nan"), + 1000.0, + 2000.0, + 2000.0, + ], + ], + dtype=torch.bfloat16, + device="cuda", + ) + temperatures = torch.tensor([0.0, -1.0, 0.0], device="cuda") + seeds = torch.tensor([7, 11, 13], dtype=torch.int64, device="cuda") + keys = gumbel_shard_keys( + logits, + temperatures, + seeds, + vocab_start=10, + valid=6, + vocab_size=15, + ) + # Row 0 has a tie at global ids 11 and 13. Row 2 has no finite real + # column, so sanitized -inf values still tie toward its first valid id. + assert torch.equal( + _tokens(keys), torch.tensor([11, 12, 10], device="cuda", dtype=torch.int64) + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_row_metadata_requires_exact_vectors(): + logits = torch.zeros((3, 8), dtype=torch.bfloat16, device="cuda") + temperatures = torch.ones(3, dtype=torch.float32, device="cuda") + seeds = torch.arange(3, dtype=torch.int64, device="cuda") + kwargs = {"vocab_start": 0, "valid": 8, "vocab_size": 8} + + with pytest.raises(ValueError, match=r"temperatures.*\[3\]"): + gumbel_shard_keys(logits, temperatures.view(3, 1), seeds, **kwargs) + with pytest.raises(ValueError, match=r"seeds.*\[3\]"): + gumbel_shard_keys(logits, temperatures, seeds.view(3, 1), **kwargs) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_seeded_keys_are_repeatable_and_independent_of_partitioning(): + generator = torch.Generator(device="cuda").manual_seed(17) + rows, vocabulary = 9, 5003 + logits = torch.randn( + rows, + vocabulary, + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + temperatures = torch.rand(rows, device="cuda", generator=generator) + 0.2 + temperatures[::4] = 0.0 + seeds = torch.arange(rows, device="cuda", dtype=torch.int64) * 7919 + 1 + + whole = gumbel_shard_keys( + logits, + temperatures, + seeds, + vocab_start=0, + valid=vocabulary, + vocab_size=vocabulary, + ) + repeated = gumbel_shard_keys( + logits, + temperatures, + seeds, + vocab_start=0, + valid=vocabulary, + vocab_size=vocabulary, + ) + split_programs = gumbel_shard_keys( + logits, + temperatures, + seeds, + vocab_start=0, + valid=vocabulary, + vocab_size=vocabulary, + splits=key_splits(vocabulary), + ) + + cut = 2501 + left = gumbel_shard_keys( + logits[:, :cut], + temperatures, + seeds, + vocab_start=0, + valid=cut, + vocab_size=vocabulary, + ) + right = gumbel_shard_keys( + logits[:, cut:].contiguous(), + temperatures, + seeds, + vocab_start=cut, + valid=vocabulary - cut, + vocab_size=vocabulary, + ) + sharded = torch.maximum(left, right) + + assert torch.equal(repeated, whole) + assert torch.equal(split_programs.amax(dim=1), whole) + assert torch.equal(sharded, whole) + + other_seeds = seeds + 104729 + changed = gumbel_shard_keys( + logits, + temperatures, + other_seeds, + vocab_start=0, + valid=vocabulary, + vocab_size=vocabulary, + ) + # Greedy rows ignore the seed; every positive-temperature row should consume it. + assert torch.equal(changed[::4], whole[::4]) + assert torch.any(changed[temperatures > 0] != whole[temperatures > 0]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_keys_match_an_independent_per_token_oracle(dtype): + generator = torch.Generator(device="cuda").manual_seed(23) + rows, columns = 11, 4097 + logits = torch.randn( + rows, + columns, + dtype=dtype, + device="cuda", + generator=generator, + ) + logits[0, 0] = float("nan") + # Greedy row 5 has an exact tie; columns outside either validity bound hold + # tempting values and must still reduce to the key floor. + logits[5, 4] = logits[5, 3011] + temperatures = torch.rand(rows, device="cuda", generator=generator) + 0.1 + temperatures[::5] = 0.0 + seeds = torch.randint( + 0, + 2**31 - 1, + (rows,), + dtype=torch.int64, + device="cuda", + generator=generator, + ) + vocab_start, valid, vocab_size = 31040, 4093, 35129 + logits[:, vocab_size - vocab_start :] = 1000 + for splits in (1, key_splits(valid)): + got = gumbel_shard_keys( + logits, + temperatures, + seeds, + vocab_start=vocab_start, + valid=valid, + vocab_size=vocab_size, + splits=splits, + ) + want = _oracle_keys( + logits, + temperatures, + seeds, + vocab_start=vocab_start, + valid=valid, + vocab_size=vocab_size, + splits=splits, + ) + assert torch.equal(got, want) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("temperature", [1.0, 0.6]) +def test_gumbel_max_matches_the_softmax_distribution(temperature): + from scipy.stats import chi2 + + samples, vocabulary = 200_000, 4096 + generator = torch.Generator(device="cuda").manual_seed(1) + logits = ( + torch.randn(vocabulary, device="cuda", generator=generator) * 3 + ).bfloat16() + # A stride-0 expansion gives every sample the same categorical distribution; + # its seed supplies the independent Gumbel stream. + expanded = logits.view(1, -1).expand(samples, -1) + temperatures = torch.full((samples,), temperature, device="cuda") + seeds = torch.arange(samples, device="cuda", dtype=torch.int64) * 7919 + 1 + tokens = _tokens( + gumbel_shard_keys( + expanded, + temperatures, + seeds, + vocab_start=0, + valid=vocabulary, + vocab_size=vocabulary, + ) + ) + counts = torch.bincount(tokens, minlength=vocabulary).double() + expected = torch.softmax(logits.double() / temperature, dim=-1) * samples + # Pool the tail so every Pearson bin has an expected count of at least five. + large = expected >= 5 + observed_bins = torch.cat([counts[large], counts[~large].sum().view(1)]) + expected_bins = torch.cat([expected[large], expected[~large].sum().view(1)]) + statistic = ( + (observed_bins - expected_bins).square() / expected_bins.clamp_min(1e-12) + ).sum() + p_value = chi2.sf(statistic.item(), observed_bins.numel() - 1) + assert p_value > 1e-4, f"chi-square p-value {p_value:.3g}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_prebuilt_kernel_captures_with_a_preallocated_output(dtype): + device = torch.cuda.current_device() + columns, vocab_start, valid, vocab_size, splits = ( + 31040, + 31040, + 31040, + 248320, + 16, + ) + build( + device, + columns=columns, + vocab_start=vocab_start, + valid=valid, + vocab_size=vocab_size, + splits=splits, + ) + logits = torch.randn(4, columns, dtype=dtype, device="cuda") + temperatures = torch.ones(4, dtype=torch.float32, device="cuda") + seeds = torch.arange(4, dtype=torch.int64, device="cuda") + 1 + captured = torch.empty((4, splits), dtype=torch.int64, device="cuda") + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + returned = gumbel_shard_keys( + logits, + temperatures, + seeds, + vocab_start=vocab_start, + valid=valid, + vocab_size=vocab_size, + splits=splits, + out=captured, + ) + assert returned.data_ptr() == captured.data_ptr() + + # Exercise the captured mixed-temperature branch with data written after capture. + temperatures.zero_() + logits.zero_() + logits[:, 31] = 4.0 + graph.replay() + expected = gumbel_shard_keys( + logits, + temperatures, + seeds, + vocab_start=vocab_start, + valid=valid, + vocab_size=vocab_size, + splits=splits, + ) + assert torch.equal(captured, expected) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 99ff26af..b5cf8628 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -189,8 +189,8 @@ def test_cp_gate_epilogue_offer_is_consumed_once(): assert not cp.take_gate_epilogue() -def test_runtime_only_slots_are_accepted_but_not_built_by_the_plugin(monkeypatch): - requested = frozenset({"gdn_commit", "sharded_gumbel"}) +def test_hook_slots_are_accepted_but_build_no_verify_width(monkeypatch): + requested = frozenset({"sharded_gumbel"}) monkeypatch.setattr(runtime, "SLOTS", requested) assert runtime.slot_reasons() == [] @@ -203,6 +203,68 @@ def test_runtime_only_slots_are_accepted_but_not_built_by_the_plugin(monkeypatch assert runtime.build(model, ctx).summary() == "slots none (stock)" +def test_draft_fin_owns_the_exact_bf16_finalize_policy(): + from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod + + from sgl_nvidia_mega.qwen3_5.draft import finalize + + method = object.__new__(UnquantizedFusedMoEMethod) + experts = SimpleNamespace( + is_draft=True, + quant_method=method, + params_dtype=torch.bfloat16, + hidden_size=4096, + num_experts=512, + top_k=10, + runner_backend=SimpleNamespace(is_flashinfer_trtllm=lambda: True), + supports_deferred_finalize=False, + ) + assert ( + finalize.enable(SimpleNamespace(experts=experts, shared_expert=object())) == "" + ) + assert experts.supports_deferred_finalize + + +@pytest.mark.parametrize( + "change, reason", + [ + ({"is_draft": False}, "not a draft MoE"), + ({"quant_method": object()}, "quant method object"), + ({"params_dtype": torch.float16}, "params dtype torch.float16"), + ({"hidden_size": 2048}, "hidden size 2048"), + ({"num_experts": 256}, "experts 256"), + ({"top_k": 8}, "top-k 8"), + ({"shared_expert": None}, "no separate shared expert"), + ( + {"runner_backend": SimpleNamespace(is_flashinfer_trtllm=lambda: False)}, + "MoE runner", + ), + ], +) +def test_draft_fin_refuses_every_unvalidated_policy_dimension(change, reason): + from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod + + from sgl_nvidia_mega.qwen3_5.draft import finalize + + change = dict(change) + experts = SimpleNamespace( + is_draft=True, + quant_method=object.__new__(UnquantizedFusedMoEMethod), + params_dtype=torch.bfloat16, + hidden_size=4096, + num_experts=512, + top_k=10, + runner_backend=SimpleNamespace(is_flashinfer_trtllm=lambda: True), + supports_deferred_finalize=False, + ) + shared_expert = change.pop("shared_expert", object()) + for name, value in change.items(): + setattr(experts, name, value) + got = finalize.enable(SimpleNamespace(experts=experts, shared_expert=shared_expert)) + assert reason in got + assert not experts.supports_deferred_finalize + + def test_multi_request_coverage_has_no_kernel_build_result(monkeypatch): monkeypatch.setattr(runtime, "SLOTS", frozenset({"multi_request"})) model = SimpleNamespace( diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index fcdbab79..47f8a102 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -611,6 +611,22 @@ def test_a_numeric_slot_turns_the_claim_into_its_tolerances(): } +def test_plugin_runtime_slot_declares_its_actual_numerics(): + from sglang.kernels.plugin_kernels.nvidia.spec import NumericsKind + + from sgl_nvidia_mega.qwen3_5 import draft + + sampled = spec.numerics(frozenset({"sharded_gumbel"})) + assert sampled.kind is NumericsKind.TOLERANCE + assert dict(sampled.tolerances) == spec.NUMERIC_TOLERANCES["sharded_gumbel"] + + # The deferred producer changes the draft's finalize/reduction order. It uses + # the draft spec's existing final-logit envelope, not a target-model tolerance. + finalize = draft.numerics(frozenset({"draft_fin"})) + assert finalize.kind is NumericsKind.TOLERANCE + assert dict(finalize.tolerances) == {"draft_logits.rel": 2**-7} + + def test_the_shipped_policy_claims_bitwise_numerics(): from sglang.kernels.plugin_kernels.nvidia.spec import NumericsKind diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 28cb16d8..31fbb09e 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -27,6 +27,38 @@ from sgl_nvidia_mega.qwen3_5.runtime import executor +class _CudaMeta: + """Tensor metadata stand-in for CPU capability-policy tests.""" + + def __init__( + self, + shape, + dtype, + *, + transpose_contiguous=True, + contiguous=True, + is_cuda=True, + device="cuda:0", + ): + self.shape = tuple(shape) + self.dtype = dtype + self.transpose_contiguous = transpose_contiguous + self.contiguous = contiguous + self.is_cuda = is_cuda + self.device = torch.device(device) + + def is_contiguous(self): + return self.contiguous + + def dim(self): + return len(self.shape) + + def t(self): + return SimpleNamespace( + is_contiguous=lambda: self.transpose_contiguous, + ) + + @pytest.fixture(autouse=True) def registered(monkeypatch): from sglang.srt.kernel_plugins.nvidia import binding as plugin_binding @@ -313,6 +345,7 @@ class _GateUp: def __init__(self, log): self.log = log + self.dtype = torch.bfloat16 def record_stream(self, stream): self.log.append(("record_stream", stream.name)) @@ -343,7 +376,6 @@ def test_the_route_shared_slot_starts_the_shared_expert_after_the_router_launch( weight=k["sh_w"], input_scale=k["sh_in"], weight_scale=k["sh_ws"] ), _enable_silu_fp4_quant_fusion=False, - _silu_fp8_quant_fusion=False, act_fn=log.op("act", "act_out"), down_proj=log.op("down_proj", ("down_out", None)), ) @@ -438,17 +470,234 @@ def test_the_route_shared_slot_names_what_a_layer_lacks(change, want): gate_up_proj=SimpleNamespace( quant_method=method, bias=change.get("bias"), - weight=torch.zeros(rows, 4096, dtype=torch.float8_e4m3fn).t(), - input_scale=torch.ones(1), - weight_scale=torch.ones(1), + weight=_CudaMeta((4096, rows), torch.float8_e4m3fn), + input_scale=_CudaMeta((1,), torch.float32), + weight_scale=_CudaMeta((1,), torch.float32), ), act_fn=change.get("act_fn", object.__new__(SiluAndMul)), + down_proj=SimpleNamespace( + quant_method=method, + weight=_CudaMeta((256, 4096), torch.float8_e4m3fn), + input_scale=_CudaMeta((1,), torch.float32), + ), ) mlp = SimpleNamespace(shared_expert=expert, shared_expert_gate="gate") got = route.shared_unsupported(mlp) assert (want in got) if want else got == "" +def test_route_shared_missing_quant_scales_falls_back_instead_of_raising(): + from sgl_nvidia_mega.qwen3_5.moe import route + + expert = qwen2_moe.Qwen2MoeMLP.__new__(qwen2_moe.Qwen2MoeMLP) + torch.nn.Module.__init__(expert) + expert.__dict__.update( + gate_up_proj=SimpleNamespace( + quant_method=object(), + bias=None, + weight=_CudaMeta((4096, 256), torch.bfloat16), + ), + act_fn=object(), + down_proj=object(), + ) + + reason = route.shared_unsupported( + SimpleNamespace(shared_expert=expert, shared_expert_gate="gate") + ) + + assert "not a static per-tensor FP8 linear" in reason + + +def test_route_shared_owns_the_static_fp8_down_projection_policy(): + from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp8LinearMethod + + from sgl_nvidia_mega.qwen3_5.moe import route + + method = object.__new__(ModelOptFp8LinearMethod) + method.use_marlin = method.use_sm120_fp8 = False + weight = _CudaMeta((256, 4096), torch.float8_e4m3fn) + scale = _CudaMeta((1,), torch.float32) + expert = SimpleNamespace( + down_proj=SimpleNamespace(quant_method=method, weight=weight, input_scale=scale) + ) + + assert route._down_takes_static_fp8(expert) + method.use_marlin = True + assert not route._down_takes_static_fp8(expert) + method.use_marlin = False + method.use_sm120_fp8 = True + assert not route._down_takes_static_fp8(expert) + method.use_sm120_fp8 = False + expert.down_proj.input_scale = _CudaMeta((2,), torch.float32) + assert not route._down_takes_static_fp8(expert) + expert.down_proj.input_scale = scale + expert.down_proj.weight = _CudaMeta((256, 4096), torch.bfloat16) + assert not route._down_takes_static_fp8(expert) + expert.down_proj = SimpleNamespace(weight=weight, input_scale=scale) + assert not route._down_takes_static_fp8(expert) + + +@pytest.mark.parametrize("scale_name", ["input_scale", "weight_scale"]) +@pytest.mark.parametrize( + "scale", + [ + _CudaMeta((1,), torch.bfloat16), + _CudaMeta((1,), torch.float32, is_cuda=False, device="cpu"), + _CudaMeta((1,), torch.float32, contiguous=False), + _CudaMeta((2,), torch.float32), + _CudaMeta((1,), torch.float32, device="cuda:1"), + ], + ids=["wrong-dtype", "not-cuda", "non-contiguous", "wrong-shape", "device"], +) +def test_route_shared_rejects_invalid_gate_up_scales(scale_name, scale): + from sglang.srt.layers.activation import SiluAndMul + from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp8LinearMethod + + from sgl_nvidia_mega.qwen3_5.moe import route + + method = object.__new__(ModelOptFp8LinearMethod) + method.use_marlin = method.use_sm120_fp8 = False + gate_up = SimpleNamespace( + quant_method=method, + bias=None, + weight=_CudaMeta((4096, 256), torch.float8_e4m3fn), + input_scale=_CudaMeta((1,), torch.float32), + weight_scale=_CudaMeta((1,), torch.float32), + ) + setattr(gate_up, scale_name, scale) + expert = qwen2_moe.Qwen2MoeMLP.__new__(qwen2_moe.Qwen2MoeMLP) + torch.nn.Module.__init__(expert) + expert.__dict__.update( + gate_up_proj=gate_up, + act_fn=object.__new__(SiluAndMul), + down_proj=SimpleNamespace( + quant_method=method, + weight=_CudaMeta((256, 4096), torch.float8_e4m3fn), + input_scale=_CudaMeta((1,), torch.float32), + ), + ) + mlp = SimpleNamespace(shared_expert=expert, shared_expert_gate="gate") + + assert "shared gate/up is not a static per-tensor FP8 linear" in ( + route.shared_unsupported(mlp) + ) + + +@pytest.mark.parametrize( + "scale", + [ + _CudaMeta((1,), torch.bfloat16), + _CudaMeta((1,), torch.float32, is_cuda=False, device="cpu"), + _CudaMeta((1,), torch.float32, contiguous=False), + _CudaMeta((2,), torch.float32), + _CudaMeta((1,), torch.float32, device="cuda:1"), + ], + ids=["wrong-dtype", "not-cuda", "non-contiguous", "wrong-shape", "device"], +) +def test_route_shared_rejects_invalid_down_scale(scale): + from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp8LinearMethod + + from sgl_nvidia_mega.qwen3_5.moe import route + + method = object.__new__(ModelOptFp8LinearMethod) + method.use_marlin = method.use_sm120_fp8 = False + expert = SimpleNamespace( + down_proj=SimpleNamespace( + quant_method=method, + weight=_CudaMeta((256, 4096), torch.float8_e4m3fn), + input_scale=scale, + ) + ) + + assert not route._down_takes_static_fp8(expert) + + +@pytest.mark.parametrize( + "x", + [ + _CudaMeta((2, 4), torch.bfloat16, is_cuda=False, device="cpu"), + _CudaMeta((2, 4), torch.float32), + _CudaMeta((4,), torch.bfloat16), + _CudaMeta((2, 4), torch.bfloat16, contiguous=False), + _CudaMeta((2, 0), torch.bfloat16), + _CudaMeta((2, 3), torch.bfloat16), + ], + ids=["not-cuda", "wrong-dtype", "wrong-rank", "non-contiguous", "empty", "odd"], +) +def test_route_shared_act_quant_rejects_invalid_input_metadata(x): + from sgl_nvidia_mega.qwen3_5.sm103.moe.shared_act_quant import ( + silu_and_mul_static_quant_fp8, + ) + + with pytest.raises(ValueError, match="shared act\\+quant input"): + silu_and_mul_static_quant_fp8(x, _CudaMeta((1,), torch.float32)) + + +@pytest.mark.parametrize( + "scale", + [ + _CudaMeta((1,), torch.float32, is_cuda=False, device="cpu"), + _CudaMeta((1,), torch.bfloat16), + _CudaMeta((2,), torch.float32), + _CudaMeta((1,), torch.float32, contiguous=False), + _CudaMeta((1,), torch.float32, device="cuda:1"), + ], + ids=["not-cuda", "wrong-dtype", "wrong-shape", "non-contiguous", "device"], +) +def test_route_shared_act_quant_rejects_invalid_scale_metadata(scale): + from sgl_nvidia_mega.qwen3_5.sm103.moe.shared_act_quant import ( + silu_and_mul_static_quant_fp8, + ) + + with pytest.raises(ValueError, match="shared act\\+quant scale"): + silu_and_mul_static_quant_fp8( + _CudaMeta((2, 4), torch.bfloat16), + scale, + ) + + +def test_route_shared_prebuilds_the_plugin_owned_act_quant(monkeypatch): + from sgl_nvidia_mega.qwen3_5.moe import route + from sgl_nvidia_mega.qwen3_5.sm103.moe import router, shared_act_quant + + layers = [ + SimpleNamespace( + mlp=SimpleNamespace( + shared_expert=SimpleNamespace( + gate_up_proj=SimpleNamespace(weight=torch.empty(4096, columns)) + ) + ) + ) + for columns in (512, 256, 256) + ] + model = SimpleNamespace( + model=SimpleNamespace(start_layer=0, end_layer=len(layers), layers=layers) + ) + calls = [] + monkeypatch.setattr(route, "unsupported", lambda mlp: "") + monkeypatch.setattr(route, "shared_unsupported", lambda mlp: "") + monkeypatch.setattr(torch.cuda, "current_device", lambda: 3) + monkeypatch.setattr( + router, + "build", + lambda device, rows, shared: calls.append(("router", device, rows, shared)), + ) + monkeypatch.setattr( + shared_act_quant, + "build", + lambda device, columns: calls.append(("tail", device, columns)), + ) + + bound = route.build(model, 7, fused=False, shared=True) + + assert bound.route and bound.shared + assert calls == [ + ("router", 3, 7, True), + ("tail", 3, 256), + ("tail", 3, 512), + ] + + @pytest.mark.parametrize("capture", [False, True]) def test_the_mkm_slot_runs_the_routed_experts_on_the_router_fronts_outputs(capture): """Slot mkm: the router front, then one fused launch whose rows and weights reach From a698526f8d9a7d77601457170cc436ba32d16776 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 27 Oct 2026 23:27:47 -0700 Subject: [PATCH 34/62] feat: attach qwen35 speculative local keys provider --- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 6 +- sgl_nvidia_mega/qwen3_5/runtime/executor.py | 4 + sgl_nvidia_mega/qwen3_5/runtime/hooks.py | 254 +++++++++++ test/kernels/test_qwen3_5_sharded_gumbel.py | 67 +++ test/unit/test_qwen3_5_runtime.py | 15 +- test/unit/test_qwen3_5_runtime_hooks.py | 482 ++++++++++++++++++++ 6 files changed, 825 insertions(+), 3 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/runtime/hooks.py create mode 100644 test/unit/test_qwen3_5_runtime_hooks.py diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 882516e5..381bacb1 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -390,7 +390,7 @@ def build(model, ctx): from sgl_nvidia_mega.qwen3_5 import attention, boundary, gdn, lowm, moe from sgl_nvidia_mega.qwen3_5 import plan as slot_plan - from sgl_nvidia_mega.qwen3_5.runtime import parity + from sgl_nvidia_mega.qwen3_5.runtime import hooks, parity from sgl_nvidia_mega.qwen3_5.runtime.executor import Qwen35Verify build_cp_lite(ctx.binding) @@ -533,6 +533,9 @@ def build(model, ctx): if parity.ENABLED else None ) + hook_factory = ( + hooks.build(model, ctx.binding) if "sharded_gumbel" in SLOTS else None + ) on = sorted(name for name, ok in built.items() if ok) return Qwen35Verify( model, @@ -542,4 +545,5 @@ def build(model, ctx): binding=ctx.binding, summary=f"slots {on or 'none (stock)'}{wider}", parity=checker, + speculative_runtime_hooks=hook_factory, ) diff --git a/sgl_nvidia_mega/qwen3_5/runtime/executor.py b/sgl_nvidia_mega/qwen3_5/runtime/executor.py index e848f5ee..758edd2c 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/executor.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/executor.py @@ -111,6 +111,7 @@ def __init__( binding=None, summary: str = "", parity=None, + speculative_runtime_hooks=None, ): from sglang.srt.models.qwen3_5 import Qwen3_5LinearDecoderLayer @@ -119,6 +120,9 @@ def __init__( self.binding = binding # SGLang's EagleBinding (BuildContext.binding) self._summary = summary self.parity = parity # runtime.parity.Parity under the test-only switch + # A factory over SGLang's stock strategy, or None. The binding reads this + # optional BuiltMegakernel member once when it wires speculative verify. + self.speculative_runtime_hooks = speculative_runtime_hooks # the attention graph variants per SGLang runner ("target_verify"), which it # captures once it takes a megakernel's (attention/route.py) self.graph_variants = ( diff --git a/sgl_nvidia_mega/qwen3_5/runtime/hooks.py b/sgl_nvidia_mega/qwen3_5/runtime/hooks.py new file mode 100644 index 00000000..c409387a --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/runtime/hooks.py @@ -0,0 +1,254 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Qwen3.5 providers for SGLang's generic speculative-runtime hooks. + +The runtime owns the verifier, its fallback and every cross-rank exchange. This +module offers only slot ``sharded_gumbel``'s local numerical policy. Its vocabulary +descriptor is validated collectively at build, before graphs are captured; the +per-step callback is local and either declines before RNG or returns reduction keys. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from sglang.kernels.plugin_kernels.nvidia.spec import ( + HooksFactory, + LocalKeysFn, + PostVerifyStrategy, + ShardedLogitsView, + SpeculativeRuntimeHooks, +) + +from sgl_nvidia_mega.qwen3_5.sm103 import sharded_gumbel + +logger = logging.getLogger(__name__) + +# The SGLang names this module reads (test/unit/test_contract.py pins them). +SGLANG_READS = { + "kernels/plugin_kernels/nvidia/spec.py": ( + "HooksFactory", + "LocalKeysFn", + "PostVerifyStrategy", + "ShardedLogitsView", + "SpeculativeRuntimeHooks", + ), + "srt/kernel_plugins/nvidia/eagle_binding.py": ( + "tp_rank", + "tp_size", + "tp_cpu_group", + ), + "srt/layers/logits_processor.py": ("vocab_size",), + "srt/layers/vocab_parallel_embedding.py": ( + "weight", + "shard_indices", + "org_vocab_start_index", + "num_org_elements", + "num_added_elements_padded", + ), +} + + +@dataclass(frozen=True) +class _Descriptor: + columns: int + start: int + valid: int + vocab_size: int + splits: int + device_index: int + + +def _descriptor(model) -> tuple[_Descriptor | None, str]: + """This rank's loaded LM-head shard, represented without any live tensor.""" + try: + head = model.lm_head + weight = head.weight + shard = head.shard_indices + columns = int(weight.shape[0]) + start = int(shard.org_vocab_start_index) + valid = int(shard.num_org_elements) + vocab_size = int(model.logits_processor.vocab_size) + added = int(shard.num_added_elements_padded) + if not weight.is_cuda or weight.dim() != 2: + return None, "LM head weight is not a CUDA matrix" + if added: + return None, f"LM head has {added} padded added-vocabulary columns" + if not 0 < valid <= columns: + return None, f"valid columns {valid} are outside [1, {columns}]" + if start < 0 or start + valid > vocab_size: + return None, ( + f"vocabulary range [{start}, {start + valid}) is outside " + f"[0, {vocab_size})" + ) + device_index = weight.device.index + if device_index is None: + device_index = torch.cuda.current_device() + return ( + _Descriptor( + columns=columns, + start=start, + valid=valid, + vocab_size=vocab_size, + # Runtime owns the partition count and derives it from every + # rank's uniform padded-logit width, not an uneven valid tail. + splits=sharded_gumbel.key_splits(columns), + device_index=int(device_index), + ), + "", + ) + except (AttributeError, IndexError, TypeError, ValueError) as error: + return None, f"LM-head shard metadata is unavailable: {error}" + + +def _gather_descriptors(local, binding) -> list[tuple[_Descriptor | None, str]]: + """Gather build metadata on the CPU group; no provider callback gathers.""" + world = int(binding.tp_size) + if world <= 0: + return [(None, f"invalid TP size {world}")] + gathered: list[tuple[_Descriptor | None, str] | None] = [None] * world + if world == 1: + gathered[0] = local + else: + dist.all_gather_object(gathered, local, group=binding.tp_cpu_group) + return [ + item if item is not None else (None, f"rank {rank} returned no descriptor") + for rank, item in enumerate(gathered) + ] + + +def _topology_reason(gathered: list[tuple[_Descriptor | None, str]], binding) -> str: + """A deterministic decision made identically from every rank's descriptors.""" + failures = [ + f"rank {rank}: {why}" + for rank, (desc, why) in enumerate(gathered) + if desc is None + ] + if failures: + return "; ".join(failures) + descriptors = [desc for desc, _ in gathered if desc is not None] + rank = int(binding.tp_rank) + if rank < 0 or rank >= len(descriptors): + return f"TP rank {rank} is outside [0, {len(descriptors)})" + common = {(desc.columns, desc.vocab_size, desc.splits) for desc in descriptors} + if len(common) != 1: + return "ranks disagree on LM-head columns, vocabulary size or key splits" + ranges = sorted((desc.start, desc.start + desc.valid) for desc in descriptors) + cursor = 0 + for start, end in ranges: + if start != cursor: + return f"vocabulary shards are not contiguous at token {cursor}" + cursor = end + if cursor != descriptors[0].vocab_size: + return ( + f"vocabulary shards end at token {cursor}, expected " + f"{descriptors[0].vocab_size}" + ) + return "" + + +class _LocalKeys: + """The graph-capturable local callback installed into the stock strategy.""" + + def __init__(self, descriptor: _Descriptor, default: LocalKeysFn): + self.descriptor = descriptor + self.default = default + self._engaged = False + + def __call__(self, view: ShardedLogitsView) -> torch.Tensor | None: + # Greedy is the runtime's bitwise policy, including when seed material is + # present. Do not inspect any provider capability before delegating it. + if view.temperatures is None: + return self.default(view) + # A request seed is not a per-row Gumbel seed. Stock seeded sampling stays + # authoritative, and this decline precedes every capability check and RNG. + if view.seeds is not None: + return None + if not self._supports(view): + return None + + values = view.values + rows = int(values.shape[0]) + seeds = torch.randint( + 0, + 2**31 - 1, + (rows,), + dtype=torch.int64, + device=values.device, + ) + # There is deliberately no catch or fallback after the RNG draw: doing so + # would perturb stock sampling's stream. A launch failure is a hard error. + keys = sharded_gumbel.gumbel_shard_keys( + values, + view.temperatures, + seeds, + vocab_start=self.descriptor.start, + valid=self.descriptor.valid, + vocab_size=self.descriptor.vocab_size, + splits=view.splits, + ) + if not self._engaged: + self._engaged = True + logger.info( + "NVIDIA megakernel: Qwen3.5 slot sharded_gumbel runs at %d rows", + rows, + ) + return keys + + def _supports(self, view: ShardedLogitsView) -> bool: + values = view.values + temperatures = view.temperatures + descriptor = self.descriptor + return bool( + isinstance(values, torch.Tensor) + and values.is_cuda + and values.device.index == descriptor.device_index + and values.dim() == 2 + and values.dtype in (torch.bfloat16, torch.float32) + and values.stride(1) == 1 + and int(values.shape[1]) == descriptor.columns + and isinstance(temperatures, torch.Tensor) + and temperatures.device == values.device + and temperatures.dtype == torch.float32 + and tuple(temperatures.shape) == (int(values.shape[0]),) + and temperatures.is_contiguous() + and view.shard.valid == descriptor.valid + and view.shard.vocab_size == descriptor.vocab_size + and view.splits == descriptor.splits + ) + + +def build(model, binding) -> HooksFactory | None: + """Build the provider factory, or return ``None`` uniformly before prewarm.""" + local = _descriptor(model) + gathered = _gather_descriptors(local, binding) + reason = _topology_reason(gathered, binding) + if reason: + logger.info( + "NVIDIA megakernel: Qwen3.5 slot sharded_gumbel uses stock: %s", reason + ) + return None + + descriptor = local[0] + assert descriptor is not None + sharded_gumbel.build( + descriptor.device_index, + columns=descriptor.columns, + vocab_start=descriptor.start, + valid=descriptor.valid, + vocab_size=descriptor.vocab_size, + splits=descriptor.splits, + ) + + def factory(defaults: SpeculativeRuntimeHooks) -> SpeculativeRuntimeHooks: + return SpeculativeRuntimeHooks( + post_verify=PostVerifyStrategy( + run=defaults.post_verify.run, + local_keys=_LocalKeys(descriptor, defaults.post_verify.local_keys), + ) + ) + + return factory diff --git a/test/kernels/test_qwen3_5_sharded_gumbel.py b/test/kernels/test_qwen3_5_sharded_gumbel.py index 3549304c..04964011 100644 --- a/test/kernels/test_qwen3_5_sharded_gumbel.py +++ b/test/kernels/test_qwen3_5_sharded_gumbel.py @@ -11,7 +11,9 @@ import torch import triton import triton.language as tl +from sglang.kernels.plugin_kernels.nvidia import spec +from sgl_nvidia_mega.qwen3_5.runtime import hooks from sgl_nvidia_mega.qwen3_5.sm103.sharded_gumbel import ( build, gumbel_shard_keys, @@ -431,3 +433,68 @@ def test_prebuilt_kernel_captures_with_a_preallocated_output(dtype): splits=splits, ) assert torch.equal(captured, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_provider_factory_draws_seeds_and_replays_from_a_cuda_graph(dtype): + rows, columns = 8, 4097 + splits = key_splits(columns) + weight = torch.empty((columns, 1), dtype=torch.bfloat16, device="cuda") + model = SimpleNamespace( + lm_head=SimpleNamespace( + weight=weight, + shard_indices=SimpleNamespace( + org_vocab_start_index=0, + num_org_elements=columns, + num_added_elements_padded=0, + ), + ), + logits_processor=SimpleNamespace(vocab_size=columns), + ) + binding = SimpleNamespace(tp_size=1, tp_rank=0) + # hooks.build invokes the real two-dtype precompile before this graph exists. + factory = hooks.build(model, binding) + assert callable(factory) + + def stock_run(io, local_keys): + return io, local_keys + + def stock_keys(view): + raise AssertionError("sampled provider delegated to greedy stock keys") + + defaults = spec.SpeculativeRuntimeHooks( + post_verify=spec.PostVerifyStrategy(run=stock_run, local_keys=stock_keys) + ) + selected = factory(defaults) + assert selected.post_verify.run is defaults.post_verify.run + local_keys = selected.post_verify.local_keys + + # Equal logits make the internally generated per-row seed observable in the + # winning key. A graph replay advances PyTorch's captured RNG state. + logits = torch.zeros((rows, columns), dtype=dtype, device="cuda") + temperatures = torch.ones(rows, dtype=torch.float32, device="cuda") + view = spec.ShardedLogitsView( + values=logits, + shard=spec.VocabShardView(start=0, valid=columns, vocab_size=columns), + temperatures=temperatures, + seeds=None, + splits=splits, + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + keys = local_keys(view) + assert keys is not None + assert keys.dtype == torch.int64 + assert keys.device == logits.device + assert keys.is_contiguous() + assert tuple(keys.shape) == (rows, splits) + + first = keys.clone() + graph.replay() + second = keys.clone() + graph.replay() + third = keys.clone() + assert not torch.equal(first, second) + assert not torch.equal(second, third) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index b5cf8628..297269f7 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -14,7 +14,7 @@ from sgl_nvidia_mega.qwen3_5 import ( runtime, ) -from sgl_nvidia_mega.qwen3_5.runtime import executor, parity +from sgl_nvidia_mega.qwen3_5.runtime import executor, hooks, parity from sgl_nvidia_mega.qwen3_5.runtime.state import StepState @@ -194,13 +194,24 @@ def test_hook_slots_are_accepted_but_build_no_verify_width(monkeypatch): monkeypatch.setattr(runtime, "SLOTS", requested) assert runtime.slot_reasons() == [] + factory = object() + calls = [] + monkeypatch.setattr( + hooks, + "build", + lambda model, binding: calls.append((model, binding)) or factory, + ) + model = SimpleNamespace( model=SimpleNamespace(start_layer=0, end_layer=0, layers=[]) ) ctx = SimpleNamespace( widths=(7,), tokens_per_request=7, deadline_s=60.0, binding="EagleBinding" ) - assert runtime.build(model, ctx).summary() == "slots none (stock)" + built = runtime.build(model, ctx) + assert built.summary() == "slots none (stock)" + assert built.speculative_runtime_hooks is factory + assert calls == [(model, "EagleBinding")] def test_draft_fin_owns_the_exact_bf16_finalize_policy(): diff --git a/test/unit/test_qwen3_5_runtime_hooks.py b/test/unit/test_qwen3_5_runtime_hooks.py new file mode 100644 index 00000000..ae31a8bb --- /dev/null +++ b/test/unit/test_qwen3_5_runtime_hooks.py @@ -0,0 +1,482 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""CPU contracts for Qwen3.5's post-verify local-key provider.""" + +from __future__ import annotations + +import logging +from types import SimpleNamespace + +import pytest +import torch +from sglang.kernels.plugin_kernels.nvidia import spec + +from sgl_nvidia_mega.qwen3_5.runtime import hooks + + +class _Tensor: + """CUDA tensor metadata for CPU-only provider policy tests.""" + + def __init__( + self, + shape, + dtype, + *, + device="cuda:0", + is_cuda=True, + stride=None, + contiguous=True, + ): + self.shape = tuple(shape) + self.dtype = dtype + self.device = torch.device(device) + self.is_cuda = is_cuda + self._stride = stride or (self.shape[-1], 1) + self._contiguous = contiguous + + def dim(self): + return len(self.shape) + + def stride(self, dimension): + return self._stride[dimension] + + def is_contiguous(self): + return self._contiguous + + +class _Weight: + def __init__(self, columns=4, *, is_cuda=True, dimensions=2): + self.shape = (columns, 8) if dimensions == 2 else (columns,) + self.is_cuda = is_cuda + self.device = torch.device("cuda:0") + self._dimensions = dimensions + + def dim(self): + return self._dimensions + + +def _descriptor(**changes): + values = { + "columns": 4, + "start": 4, + "valid": 4, + "vocab_size": 8, + "splits": 2, + "device_index": 0, + } + values.update(changes) + return hooks._Descriptor(**values) + + +def _model( + *, + columns=4, + start=4, + valid=4, + vocab_size=8, + added=0, + is_cuda=True, + dimensions=2, +): + return SimpleNamespace( + lm_head=SimpleNamespace( + weight=_Weight(columns, is_cuda=is_cuda, dimensions=dimensions), + shard_indices=SimpleNamespace( + org_vocab_start_index=start, + num_org_elements=valid, + num_added_elements_padded=added, + ), + ), + logits_processor=SimpleNamespace(vocab_size=vocab_size), + ) + + +def _defaults(default_local=None): + if default_local is None: + default_local = lambda view: None + + def run(io, local_keys): + return io, local_keys + + return spec.SpeculativeRuntimeHooks( + post_verify=spec.PostVerifyStrategy(run=run, local_keys=default_local) + ) + + +def _view( + values, + temperatures, + *, + seeds=None, + start=4, + valid=4, + vocab_size=8, + splits=2, +): + return spec.ShardedLogitsView( + values=values, + shard=spec.VocabShardView( + start=start, + valid=valid, + vocab_size=vocab_size, + ), + temperatures=temperatures, + seeds=seeds, + splits=splits, + ) + + +def _local(default_local=None): + defaults = _defaults(default_local) + factory = lambda stock: spec.SpeculativeRuntimeHooks( + post_verify=spec.PostVerifyStrategy( + run=stock.post_verify.run, + local_keys=hooks._LocalKeys(_descriptor(), stock.post_verify.local_keys), + ) + ) + return factory(defaults).post_verify.local_keys, defaults + + +def test_build_collects_descriptors_then_prebuilds_exact_local_shard(monkeypatch): + model = _model() + group = object() + binding = SimpleNamespace(tp_size=2, tp_rank=1, tp_cpu_group=group) + prebuilt = [] + + def gather(output, local, *, group): + assert group is binding.tp_cpu_group + output[:] = [(_descriptor(start=0), ""), local] + + monkeypatch.setattr(hooks.dist, "all_gather_object", gather) + monkeypatch.setattr( + hooks.sharded_gumbel, + "build", + lambda *args, **kwargs: prebuilt.append((args, kwargs)), + ) + + factory = hooks.build(model, binding) + + assert callable(factory) + assert prebuilt == [ + ( + (0,), + { + "columns": 4, + "vocab_start": 4, + "valid": 4, + "vocab_size": 8, + "splits": 2, + }, + ) + ] + defaults = _defaults() + selected = factory(defaults) + assert selected.post_verify.run is defaults.post_verify.run + assert selected.post_verify.local_keys is not defaults.post_verify.local_keys + + +def test_padded_columns_keep_splits_uniform_across_an_uneven_valid_tail(monkeypatch): + columns, tail, vocab_size = 4104, 4041, 8145 + model = _model( + columns=columns, + start=columns, + valid=tail, + vocab_size=vocab_size, + ) + binding = SimpleNamespace(tp_size=2, tp_rank=1, tp_cpu_group=object()) + prebuilt = [] + + def gather(output, local, *, group): + output[:] = [ + ( + hooks._Descriptor( + columns=columns, + start=0, + valid=columns, + vocab_size=vocab_size, + splits=4, + device_index=0, + ), + "", + ), + local, + ] + + monkeypatch.setattr(hooks.dist, "all_gather_object", gather) + monkeypatch.setattr( + hooks.sharded_gumbel, + "build", + lambda *args, **kwargs: prebuilt.append((args, kwargs)), + ) + factory = hooks.build(model, binding) + + assert callable(factory) + assert prebuilt == [ + ( + (0,), + { + "columns": columns, + "vocab_start": columns, + "valid": tail, + "vocab_size": vocab_size, + "splits": 4, + }, + ) + ] + + monkeypatch.setattr(hooks.torch, "Tensor", _Tensor) + generated = _Tensor((2,), torch.int64) + launches = [] + monkeypatch.setattr(hooks.torch, "randint", lambda *args, **kwargs: generated) + monkeypatch.setattr( + hooks.sharded_gumbel, + "gumbel_shard_keys", + lambda *args, **kwargs: launches.append((args, kwargs)) or object(), + ) + local_keys = factory(_defaults()).post_verify.local_keys + values = _Tensor((2, columns), torch.bfloat16) + temperatures = _Tensor((2,), torch.float32) + local_keys( + _view( + values, + temperatures, + start=columns, + valid=tail, + vocab_size=vocab_size, + splits=4, + ) + ) + assert launches == [ + ( + (values, temperatures, generated), + { + "vocab_start": columns, + "valid": tail, + "vocab_size": vocab_size, + "splits": 4, + }, + ) + ] + + +@pytest.mark.parametrize( + "gathered, match", + [ + ( + [(_descriptor(start=0), ""), (None, "missing LM head")], + "rank 1: missing LM head", + ), + ( + [(_descriptor(start=0), ""), (_descriptor(start=5), "")], + "not contiguous", + ), + ( + [(_descriptor(start=0), ""), (_descriptor(start=4, splits=4), "")], + "disagree", + ), + ], +) +def test_unsupported_topology_returns_none_uniformly_without_prebuild( + monkeypatch, caplog, gathered, match +): + binding = SimpleNamespace(tp_size=2, tp_rank=0, tp_cpu_group=object()) + + def gather(output, local, *, group): + output[:] = gathered + + monkeypatch.setattr(hooks.dist, "all_gather_object", gather) + monkeypatch.setattr( + hooks.sharded_gumbel, + "build", + lambda *args, **kwargs: pytest.fail("unsupported provider was prebuilt"), + ) + + with caplog.at_level(logging.INFO, logger=hooks.__name__): + assert hooks.build(_model(start=0), binding) is None + assert match in caplog.text + + +@pytest.mark.parametrize( + "model, match", + [ + (SimpleNamespace(), "metadata is unavailable"), + (_model(is_cuda=False), "not a CUDA matrix"), + (_model(dimensions=1), "not a CUDA matrix"), + (_model(added=1), "added-vocabulary"), + (_model(valid=0), "valid columns"), + (_model(valid=5), "valid columns"), + (_model(start=7, valid=4), "outside"), + ], +) +def test_local_metadata_refusal_is_a_build_result_not_an_exception( + monkeypatch, caplog, model, match +): + binding = SimpleNamespace(tp_size=1, tp_rank=0) + monkeypatch.setattr( + hooks.sharded_gumbel, + "build", + lambda *args, **kwargs: pytest.fail("unsupported provider was prebuilt"), + ) + with caplog.at_level(logging.INFO, logger=hooks.__name__): + assert hooks.build(model, binding) is None + assert match in caplog.text + + +def test_greedy_always_delegates_by_identity_even_with_seed_material(): + sentinel = object() + calls = [] + + def default(view): + calls.append(view) + return sentinel + + local, defaults = _local(default) + view = _view(object(), None, seeds=object()) + + assert local(view) is sentinel + assert calls == [view] + assert local.default is defaults.post_verify.local_keys + + +def test_sampled_request_seeds_decline_before_checks_or_rng(monkeypatch): + local, _ = _local() + monkeypatch.setattr( + hooks.torch, + "randint", + lambda *args, **kwargs: pytest.fail("seeded fallback consumed RNG"), + ) + # Neither object implements tensor metadata; an early capability read fails. + view = _view(object(), object(), seeds=object()) + assert local(view) is None + + +@pytest.mark.parametrize( + "change", + [ + {"values": object()}, + {"values.is_cuda": False}, + {"values.device": "cuda:1"}, + {"values.shape": (2, 4, 1)}, + {"values.dtype": torch.float16}, + {"values.stride": (8, 2)}, + {"values.shape": (2, 5)}, + {"temperatures.device": "cuda:1"}, + {"temperatures.dtype": torch.bfloat16}, + {"temperatures.shape": (2, 1)}, + {"temperatures.contiguous": False}, + {"valid": 3}, + {"vocab_size": 9}, + {"splits": 4}, + ], +) +def test_every_sampled_capability_refusal_precedes_rng(monkeypatch, change): + monkeypatch.setattr(hooks.torch, "Tensor", _Tensor) + values = _Tensor((2, 4), torch.bfloat16) + temperatures = _Tensor((2,), torch.float32) + view_kwargs = {} + if "values" in change: + values = change["values"] + for name, value in change.items(): + if name.startswith("values."): + attribute = name.split(".", 1)[1] + if attribute == "shape": + values.shape = value + elif attribute == "stride": + values._stride = value + elif attribute == "device": + values.device = torch.device(value) + else: + setattr(values, attribute, value) + elif name.startswith("temperatures."): + attribute = name.split(".", 1)[1] + if attribute == "contiguous": + temperatures._contiguous = value + elif attribute == "device": + temperatures.device = torch.device(value) + else: + setattr(temperatures, attribute, value) + elif name != "values": + view_kwargs[name] = value + local, _ = _local() + monkeypatch.setattr( + hooks.torch, + "randint", + lambda *args, **kwargs: pytest.fail("unsupported view consumed RNG"), + ) + + assert local(_view(values, temperatures, **view_kwargs)) is None + + +def test_supported_sample_draws_exact_seed_then_launches_without_collective( + monkeypatch, caplog +): + monkeypatch.setattr(hooks.torch, "Tensor", _Tensor) + values = _Tensor((3, 4), torch.float32) + temperatures = _Tensor((3,), torch.float32) + generated = _Tensor((3,), torch.int64) + keys = object() + events = [] + + def randint(*args, **kwargs): + events.append(("rng", args, kwargs)) + return generated + + def launch(*args, **kwargs): + events.append(("launch", args, kwargs)) + return keys + + monkeypatch.setattr(hooks.torch, "randint", randint) + monkeypatch.setattr(hooks.dist, "all_gather_object", pytest.fail) + monkeypatch.setattr(hooks.sharded_gumbel, "gumbel_shard_keys", launch) + local, _ = _local() + # The runtime descriptor's rank-varying start is not a call-time applicability + # branch. The build-captured start remains the kernel's authoritative value. + view = _view(values, temperatures, start=777) + + with caplog.at_level(logging.INFO, logger=hooks.__name__): + assert local(view) is keys + assert local(view) is keys + + assert events[0] == ( + "rng", + (0, 2**31 - 1, (3,)), + {"dtype": torch.int64, "device": torch.device("cuda:0")}, + ) + assert events[1] == ( + "launch", + (values, temperatures, generated), + { + "vocab_start": 4, + "valid": 4, + "vocab_size": 8, + "splits": 2, + }, + ) + marker = "NVIDIA megakernel: Qwen3.5 slot sharded_gumbel runs at 3 rows" + assert caplog.text.count(marker) == 1 + + +def test_failure_after_rng_is_hard_and_does_not_set_engagement(monkeypatch, caplog): + monkeypatch.setattr(hooks.torch, "Tensor", _Tensor) + values = _Tensor((2, 4), torch.bfloat16) + temperatures = _Tensor((2,), torch.float32) + generated = _Tensor((2,), torch.int64) + monkeypatch.setattr(hooks.torch, "randint", lambda *args, **kwargs: generated) + attempts = [] + + def launch(*args, **kwargs): + attempts.append(None) + if len(attempts) == 1: + raise RuntimeError("launch failed after RNG") + return object() + + monkeypatch.setattr(hooks.sharded_gumbel, "gumbel_shard_keys", launch) + local, _ = _local() + view = _view(values, temperatures) + + with caplog.at_level(logging.INFO, logger=hooks.__name__): + with pytest.raises(RuntimeError, match="launch failed after RNG"): + local(view) + assert "slot sharded_gumbel runs" not in caplog.text + local(view) + assert caplog.text.count("slot sharded_gumbel runs at 2 rows") == 1 From fc50b17ec29e782ad72982e12d2853aae77fc92f Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Wed, 28 Oct 2026 00:28:38 -0700 Subject: [PATCH 35/62] fix: accept modelopt scalar shared-expert scales --- sgl_nvidia_mega/qwen3_5/moe/route.py | 10 +++-- test/kernels/test_qwen3_5_moe.py | 37 +++++++++++------ test/unit/test_qwen3_5_steps.py | 61 +++++++++++++++++++++++++++- 3 files changed, 89 insertions(+), 19 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index 0e4e4841..b7d488ee 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -396,13 +396,13 @@ def shared_unsupported(mlp) -> str: or not input_scale.is_cuda or input_scale.device != weight.device or input_scale.dtype != torch.float32 - or tuple(input_scale.shape) != (1,) + or tuple(input_scale.shape) not in ((), (1,)) or not input_scale.is_contiguous() or weight_scale is None or not weight_scale.is_cuda or weight_scale.device != weight.device or weight_scale.dtype != torch.float32 - or tuple(weight_scale.shape) != (1,) + or tuple(weight_scale.shape) not in ((), (1,)) or not weight_scale.is_contiguous() ): return "the shared gate/up is not a static per-tensor FP8 linear of <= 512 rows" @@ -423,7 +423,9 @@ def shared_expert_tail(expert, gate_up: torch.Tensor) -> torch.Tensor: silu_and_mul_static_quant_fp8, ) - quantized = silu_and_mul_static_quant_fp8(gate_up, expert.down_proj.input_scale) + quantized = silu_and_mul_static_quant_fp8( + gate_up, expert.down_proj.input_scale.reshape(1) + ) key = ("route_shared_act_quant", gate_up.shape[0]) if key not in _ENGAGED: _ENGAGED.add(key) @@ -464,7 +466,7 @@ def _down_takes_static_fp8(expert) -> bool: and scale.is_cuda and scale.device == weight.device and scale.dtype == torch.float32 - and tuple(scale.shape) == (1,) + and tuple(scale.shape) in ((), (1,)) and scale.is_contiguous() ) diff --git a/test/kernels/test_qwen3_5_moe.py b/test/kernels/test_qwen3_5_moe.py index afb2f272..12d0d961 100644 --- a/test/kernels/test_qwen3_5_moe.py +++ b/test/kernels/test_qwen3_5_moe.py @@ -490,13 +490,11 @@ def test_shared_swiglu_static_quant_replays_from_a_cuda_graph(): assert torch.equal(expected.view(torch.uint8), captured.view(torch.uint8)) -def test_shared_swiglu_static_quant_feeds_modelopt_flashinfer_down_projection(): +def test_shared_tail_normalizes_modelopt_scalar_scale_for_flashinfer_down_projection(): from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp8LinearMethod - from sgl_nvidia_mega.qwen3_5.sm103.moe.shared_act_quant import ( - silu_and_mul_static_quant_fp8, - ) + from sgl_nvidia_mega.qwen3_5.moe import route generator = torch.Generator(device="cuda").manual_seed(91) rows, intermediate, hidden = 8, 256, 4096 @@ -513,18 +511,31 @@ def test_shared_swiglu_static_quant_feeds_modelopt_flashinfer_down_projection(): weight = ( torch.randn(hidden, intermediate, device="cuda", generator=generator) * 0.05 ).to(torch.float8_e4m3fn) - input_scale = torch.tensor([0.02], device="cuda") - layer = SimpleNamespace( - use_flashinfer_bmm=True, - weight=weight.t(), - weight_scale=torch.tensor([0.01], device="cuda"), - input_scale=input_scale, - ) + input_scale = torch.tensor(0.02, device="cuda") method = object.__new__(ModelOptFp8LinearMethod) method.use_marlin = method.use_sm120_fp8 = False - expected = method.apply(layer, SiluAndMul()(gate_up)) - actual = method.apply(layer, silu_and_mul_static_quant_fp8(gate_up, input_scale)) + class Down: + use_flashinfer_bmm = True + quant_method = method + + def __init__(self): + self.weight = weight.t() + self.weight_scale = torch.tensor(0.01, device="cuda") + self.input_scale = input_scale + + def __call__(self, value): + return method.apply(self, value), None + + down = Down() + expert = SimpleNamespace( + _enable_silu_fp4_quant_fusion=False, + act_fn=SiluAndMul(), + down_proj=down, + ) + + expected = method.apply(down, SiluAndMul()(gate_up)) + actual = route.shared_expert_tail(expert, gate_up) assert actual.dtype == torch.bfloat16 assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)) diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 31fbb09e..15209a77 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -485,6 +485,14 @@ def test_the_route_shared_slot_names_what_a_layer_lacks(change, want): got = route.shared_unsupported(mlp) assert (want in got) if want else got == "" + if not change: + # ModelOptFp8LinearMethod collapses these scales with ``.max()`` after + # loading, so production modules carry 0-D Parameters rather than [1]. + expert.gate_up_proj.input_scale = _CudaMeta((), torch.float32) + expert.gate_up_proj.weight_scale = _CudaMeta((), torch.float32) + expert.down_proj.input_scale = _CudaMeta((), torch.float32) + assert route.shared_unsupported(mlp) == "" + def test_route_shared_missing_quant_scales_falls_back_instead_of_raising(): from sgl_nvidia_mega.qwen3_5.moe import route @@ -522,6 +530,9 @@ def test_route_shared_owns_the_static_fp8_down_projection_policy(): ) assert route._down_takes_static_fp8(expert) + expert.down_proj.input_scale = _CudaMeta((), torch.float32) + assert route._down_takes_static_fp8(expert) + expert.down_proj.input_scale = scale method.use_marlin = True assert not route._down_takes_static_fp8(expert) method.use_marlin = False @@ -537,6 +548,36 @@ def test_route_shared_owns_the_static_fp8_down_projection_policy(): assert not route._down_takes_static_fp8(expert) +def test_route_shared_normalizes_modelopts_scalar_down_scale(monkeypatch): + from sgl_nvidia_mega.qwen3_5.moe import route + from sgl_nvidia_mega.qwen3_5.sm103.moe import shared_act_quant + + scale = torch.ones((), dtype=torch.float32) + seen = [] + + def quantize(gate_up, got_scale): + seen.append((gate_up, got_scale)) + return "quantized" + + class Down: + input_scale = scale + + def __call__(self, value): + return value, None + + expert = SimpleNamespace( + _enable_silu_fp4_quant_fusion=False, + down_proj=Down(), + ) + monkeypatch.setattr(route, "_down_takes_static_fp8", lambda unused: True) + monkeypatch.setattr(shared_act_quant, "silu_and_mul_static_quant_fp8", quantize) + + gate_up = torch.ones((2, 4), dtype=torch.bfloat16) + assert route.shared_expert_tail(expert, gate_up) == "quantized" + assert seen[0][0] is gate_up + assert tuple(seen[0][1].shape) == (1,) + + @pytest.mark.parametrize("scale_name", ["input_scale", "weight_scale"]) @pytest.mark.parametrize( "scale", @@ -545,9 +586,17 @@ def test_route_shared_owns_the_static_fp8_down_projection_policy(): _CudaMeta((1,), torch.float32, is_cuda=False, device="cpu"), _CudaMeta((1,), torch.float32, contiguous=False), _CudaMeta((2,), torch.float32), + _CudaMeta((1, 1), torch.float32), _CudaMeta((1,), torch.float32, device="cuda:1"), ], - ids=["wrong-dtype", "not-cuda", "non-contiguous", "wrong-shape", "device"], + ids=[ + "wrong-dtype", + "not-cuda", + "non-contiguous", + "wrong-size", + "wrong-rank", + "device", + ], ) def test_route_shared_rejects_invalid_gate_up_scales(scale_name, scale): from sglang.srt.layers.activation import SiluAndMul @@ -590,9 +639,17 @@ def test_route_shared_rejects_invalid_gate_up_scales(scale_name, scale): _CudaMeta((1,), torch.float32, is_cuda=False, device="cpu"), _CudaMeta((1,), torch.float32, contiguous=False), _CudaMeta((2,), torch.float32), + _CudaMeta((1, 1), torch.float32), _CudaMeta((1,), torch.float32, device="cuda:1"), ], - ids=["wrong-dtype", "not-cuda", "non-contiguous", "wrong-shape", "device"], + ids=[ + "wrong-dtype", + "not-cuda", + "non-contiguous", + "wrong-size", + "wrong-rank", + "device", + ], ) def test_route_shared_rejects_invalid_down_scale(scale): from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp8LinearMethod From 1c5e308bf9f4b0facb8d64a77b7e61ea048cc7b3 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Wed, 28 Oct 2026 16:34:51 -0700 Subject: [PATCH 36/62] docs: restore hop-mid in tuned qwen35 policy --- README.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/README.md b/README.md index f6a48dc4..8ce32131 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,7 @@ The remaining kernel families are opt-in through that one slot policy. These cha | --- | --- | --- | | `route`, `route_gemm`, `mkm`, `mkm_wide` | 1–32 | Fused router fronts and the small routed-expert engines | | `gdn_core` | 7–224 | ReplaySSM verify, gated norm and FP8 output in one launch | +| `hop_mid` (with `hop`) | 16–28 | Extends the one-shot boundary mailbox through four requests instead of returning those widths to the stock boundary | | `gdn_push` (with `gdn_verify`, `hop`) | 7 | The GatedDeltaNet `out_proj` as the hop's GEMM, whose epilogue pushes the partial into the hop's mailbox | | `gdn_front`, `gdn_trt_front` | 7–28 | Fused GatedDeltaNet projections and conv before the dependent verify | | `route_shared` (with `route`) | 1–8 | The shared expert's static-FP8 gate/up in the router front's launch, so the shared expert on the alternate stream waits for that launch only | @@ -73,6 +74,11 @@ The remaining kernel families are opt-in through that one slot policy. These cha | `lowm_in`, `lowm_out`, `lowm_se`, `lowm_push` | shape-bounded | Validated low-row projection and shared-expert launches | | `draft_moe`, `draft_q1`, `draft_xchg`, `draft_hop`, `draft_fin`, `draft_front`, `draft_skv2` | shape-bounded | Fused MTP draft kernels and boundaries | +The tuned Qwen3.5 TP8/TEP8 serving policy explicitly includes `hop_mid` whenever it +includes `hop`. It remains opt-in rather than part of the bitwise default policy; +an explicit `SGLANG_NVIDIA_MEGAKERNEL_SLOTS` value replaces, rather than extends, +that default list. + The sampling and acceptance after the verify, and the GatedDeltaNet extend of prefill, stay SGLang's own code. In an SGLang that has the switches, `SGLANG_ENABLE_EAGLE_POST_VERIFY_CUDA_GRAPH=1` runs the sampling and acceptance in CUDA graphs, and `SGLANG_ENABLE_GDN_PREFILL_GRAPH_EXTEND=1` captures the GatedDeltaNet extend inside the breakable prefill CUDA graphs. ### Design From b399c8d7801575fd2e04a407a842f4c12c31324c Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Wed, 28 Oct 2026 17:36:50 -0700 Subject: [PATCH 37/62] perf: widen qwen35 gdn core to high batch --- sgl_nvidia_mega/qwen3_5/plan.py | 4 ++-- test/unit/test_qwen3_5_spec.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index 21cfa4a2..1ca1c3c0 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -16,7 +16,7 @@ the same routed-expert arithmetic on route_topk's 17–32 rows. ``gdn_verify`` (the GatedDeltaNet verify chain with its launch choices changed) covers every width, with ``gdn_core`` (numeric: its ReplaySSM verify, gated norm and FP8 output in one launch) -up to 224 rows, and ``hop`` (the boundary's fused all-reduce points on the decode hop) the +up to 896 rows, and ``hop`` (the boundary's fused all-reduce points on the decode hop) the widths where FlashInfer's low-latency kernels, which it equals bit for bit, run, or with ``hop_mid`` (numeric: its one-shot protocol past FlashInfer's) up to 28 rows. ``gdn_push`` (numeric: out_proj as the hop's push GEMM, on gdn_verify's FP8 norm) runs up @@ -61,7 +61,7 @@ # The router front's 16-row tiles replace SGLang's router GEMM through two # waves; wider route_topk rows keep the stock GEMM. ROUTE_GEMM_MAX_ROWS = 32 -GDN_CORE_MAX_ROWS = 224 # and gdn_core's launch +GDN_CORE_MAX_ROWS = 896 # and gdn_core's launch (TRT-L4: checked at 64 and 128 requests) MKM_MAX_ROWS = 7 MKM_WIDE_MAX_ROWS = 16 MKM_M32_MAX_ROWS = 32 diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 47f8a102..f61f9071 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -335,7 +335,7 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): assert three.hop and three.hop_fp8 and three.mka_push assert not (three.lowm_in or three.lowm_outproj or three.lowm_oproj) # The 28/56-row low-M projection paths stay stock; hop_mid's one-shot path and - # FP8 copy also serve 28 rows, while GDN core and MKA serve through 224. + # FP8 copy also serve 28 rows, while GDN core and MKA serve through 896. for rows in (28, 56, 112, 224): expected = [ "gdn_conv", @@ -363,6 +363,7 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): for rows in (448, 896): expected = [ "gdn_conv", + "gdn_core", "gdn_pdl", "gdn_verify", "mka", From dee1b5f6dca3a670455170e489057fb5318cc625 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Wed, 28 Oct 2026 21:07:45 -0700 Subject: [PATCH 38/62] fix: align qwen35 mirrors and fp8 output dtype --- sgl_nvidia_mega/qwen3_5/attention/split_kv.py | 1 + sgl_nvidia_mega/qwen3_5/gdn/__init__.py | 8 + sgl_nvidia_mega/qwen3_5/gdn/verify.py | 20 ++- sgl_nvidia_mega/qwen3_5/lowm.py | 3 + test/unit/test_qwen3_5_steps.py | 143 ++++++++++++++++++ 5 files changed, 173 insertions(+), 2 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py index d1859149..85638184 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py +++ b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py @@ -87,6 +87,7 @@ "066a1e0aa5bd6e5c", "10bd9120d713848b", "37c7c1fa55956740", + "1445b85387928caf", ), "srt/layers/attention/trtllm_mha_backend.py:TRTLLMHAAttnBackend._should_use_fused_fp8_path": "e0b14791d523a019", "srt/layers/attention/trtllm_mha_backend.py:TRTLLMHAAttnBackend._fused_fp8_qkv_kv_cache": ( diff --git a/sgl_nvidia_mega/qwen3_5/gdn/__init__.py b/sgl_nvidia_mega/qwen3_5/gdn/__init__.py index 6a0e95f5..cef99e4f 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/__init__.py @@ -49,6 +49,7 @@ "srt/layers/attention/hybrid_linear_attn_backend.py": ("linear_attn_backend",), "srt/model_executor/forward_context.py": ("get_attn_backend",), "srt/layers/layer_boundary/fusions/norm.py": ("fused_norm_gamma",), + "srt/layers/quantization/modelopt_quant.py": ("orig_dtype",), "srt/runtime_context.py": ("get_parallel",), } @@ -100,9 +101,16 @@ def build(model, ctx) -> Bound: return Bound(reason=f"layer {i}: {reason}") packed = {} if "gdn_push" in SLOTS: + import torch + from sgl_nvidia_mega.qwen3_5.sm103.hop import pack_out_proj for i, gdn in gdns: + # The hop's out-projection kernel writes BF16. Keep FP16 layers on + # the non-push path, whose FlashInfer BMM honors layer.orig_dtype. + if getattr(gdn.out_proj, "orig_dtype", torch.bfloat16) != torch.bfloat16: + packed = {} + break # the method keeps the weight as the [K, N] transpose of a row-major [N, K] try: packed[gdn.attn.layer_id] = pack_out_proj(gdn.out_proj.weight.t()) diff --git a/sgl_nvidia_mega/qwen3_5/gdn/verify.py b/sgl_nvidia_mega/qwen3_5/gdn/verify.py index ddf6c970..57cd453b 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/verify.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/verify.py @@ -82,6 +82,7 @@ "use_marlin", "use_sm120_fp8", "use_flashinfer_bmm", + "orig_dtype", '"input_scale"', '"weight_scale"', ), @@ -152,8 +153,12 @@ "srt/layers/quantization/fp8_utils.py:apply_fp8_linear_bmm_flashinfer": ( "c66562e1de0008db", "c94a4f1a622c2fe0", + "aaa39b9452d9df32", + ), + "srt/layers/quantization/modelopt_quant.py:ModelOptFp8LinearMethod.apply": ( + "91b563cf074bff42", + "4f99db0f712ad24a", ), - "srt/layers/quantization/modelopt_quant.py:ModelOptFp8LinearMethod.apply": "91b563cf074bff42", } RING_LEN, MAX_DRAFT = 16, 8 # the verify tiles' bounds, as measured @@ -352,7 +357,18 @@ def _out_proj( its own quantization returns), the TP-partial output in SGLang's symmetric pool; with slot lowm_out (``Slots.lowm_outproj``) the low-M launch instead (numeric), pushing into the hop with ``push`` (slot lowm_push).""" - if st is not None and st.slots.lowm_outproj: + # A prequantized E4M3 tensor does not carry the activation dtype that ModelOpt's + # linear must produce. Match ModelOptFp8LinearMethod.apply's v3 contract rather + # than inheriting the GDN recurrence's (potentially different) working dtype. + dtype = out_proj.orig_dtype + if dtype not in (torch.bfloat16, torch.float16): + raise ValueError( + "prequantized FP8 FlashInfer BMM requires an explicit " + "BF16 or FP16 output dtype" + ) + # The low-M kernel writes BF16 unconditionally. An FP16 ModelOpt layer must + # therefore take the stock FlashInfer BMM below even when lowm_out is enabled. + if st is not None and st.slots.lowm_outproj and dtype == torch.bfloat16: from sgl_nvidia_mega.qwen3_5 import lowm hop = st.bound["boundary"].hop if push else None diff --git a/sgl_nvidia_mega/qwen3_5/lowm.py b/sgl_nvidia_mega/qwen3_5/lowm.py index a5b613cd..4887ea38 100644 --- a/sgl_nvidia_mega/qwen3_5/lowm.py +++ b/sgl_nvidia_mega/qwen3_5/lowm.py @@ -68,6 +68,7 @@ "use_marlin", "use_sm120_fp8", "use_flashinfer_bmm", + "orig_dtype", '"input_scale"', '"weight_scale"', ), @@ -192,6 +193,8 @@ def fp8_unsupported(linear) -> str: return "scales are not FP32" if linear.bias is not None: return "a bias" + if getattr(linear, "orig_dtype", torch.bfloat16) != torch.bfloat16: + return f"output dtype {linear.orig_dtype} (low-M writes BF16)" return "" diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 15209a77..20b19b6a 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -1088,6 +1088,7 @@ def _gdn_verify_setup(log, *, ring_len=16, temporal_dtype=torch.bfloat16, draft= weight_scale=k["out_s"], input_scale=k["scale"], tp_group="tp", + orig_dtype=torch.bfloat16, ) gdn = SimpleNamespace( attn=attn, @@ -1300,6 +1301,148 @@ def symmetric(group, disabled=False): ) +@pytest.mark.parametrize( + "model_dtype,recurrence_dtype", + [ + (torch.bfloat16, torch.float16), + (torch.float16, torch.bfloat16), + (torch.float32, torch.bfloat16), + ], +) +def test_gdn_prequantized_out_proj_uses_v3_model_dtype( + monkeypatch, model_dtype, recurrence_dtype +): + """The plugin's prequantized out_proj follows the layer dtype, not the + recurrence's working dtype, exactly as v3 ModelOpt does.""" + from sglang.srt import runtime_context + from sglang.srt.distributed.device_communicators import pynccl_allocator + from sglang.srt.layers import dp_attention + from sglang.srt.layers.quantization import fp8_utils + + from sgl_nvidia_mega.qwen3_5.gdn import verify + + layer = SimpleNamespace( + weight=torch.zeros((4, 3), dtype=torch.float8_e4m3fn), + weight_scale=torch.ones(1), + input_scale=torch.ones(1), + tp_group="tp", + orig_dtype=model_dtype, + ) + fp8 = torch.zeros((2, 4), dtype=torch.float8_e4m3fn) + calls = [] + + def bmm(x, weight, x_scale, weight_scale, output_dtype): + calls.append((x, weight, x_scale, weight_scale, output_dtype)) + return torch.zeros((2, 3), dtype=output_dtype) + + # This symbol is platform-gated in SGLang and is absent in a real CPU-only + # contract process, so install the native BMM stand-in explicitly. + monkeypatch.setattr(fp8_utils, "flashinfer_bmm_fp8", bmm, raising=False) + monkeypatch.setattr( + pynccl_allocator, + "use_symmetric_memory", + lambda *args, **kwargs: contextlib.nullcontext(), + ) + monkeypatch.setattr(dp_attention, "is_allocation_symmetric", lambda: True) + monkeypatch.setattr( + runtime_context, "get_parallel", lambda: SimpleNamespace(tp_group=None) + ) + + if model_dtype == torch.float32: + with pytest.raises(ValueError, match="explicit BF16 or FP16"): + fp8_utils.apply_fp8_linear_bmm_flashinfer( + fp8, + layer.weight, + layer.weight_scale, + layer.input_scale, + pre_quant_output_dtype=model_dtype, + ) + with pytest.raises(ValueError, match="explicit BF16 or FP16"): + verify._out_proj(layer, fp8, recurrence_dtype) + assert calls == [] + return + + expected = fp8_utils.apply_fp8_linear_bmm_flashinfer( + fp8, + layer.weight, + layer.weight_scale, + layer.input_scale, + pre_quant_output_dtype=model_dtype, + ) + actual = verify._out_proj(layer, fp8, recurrence_dtype) + + assert actual.dtype == expected.dtype == model_dtype + assert torch.equal(actual, expected) + assert len(calls) == 2 + assert calls[0][-1] == calls[1][-1] == model_dtype + assert calls[0][0].data_ptr() == calls[1][0].data_ptr() == fp8.data_ptr() + assert all(left is right for left, right in zip(calls[0][1:4], calls[1][1:4])) + + +def test_gdn_bf16_out_proj_can_use_lowm(monkeypatch): + from sgl_nvidia_mega.qwen3_5 import lowm + from sgl_nvidia_mega.qwen3_5.gdn import verify + + layer = SimpleNamespace(orig_dtype=torch.bfloat16) + fp8 = torch.zeros((2, 4), dtype=torch.float8_e4m3fn) + seen = [] + + def row_parallel(linear, x, dtype, what, hop): + seen.append((linear, x, dtype, what, hop)) + return "out" + + monkeypatch.setattr(lowm, "row_parallel", row_parallel) + st = SimpleNamespace(slots=SimpleNamespace(lowm_outproj=True)) + + assert verify._out_proj(layer, fp8, torch.float32, st) == "out" + assert seen == [(layer, fp8, torch.bfloat16, "out_proj", None)] + + +def test_gdn_fp16_out_proj_bypasses_the_bf16_only_lowm(monkeypatch): + from sglang.srt import runtime_context + from sglang.srt.distributed.device_communicators import pynccl_allocator + from sglang.srt.layers import dp_attention + from sglang.srt.layers.quantization import fp8_utils + + from sgl_nvidia_mega.qwen3_5 import lowm + from sgl_nvidia_mega.qwen3_5.gdn import verify + + layer = SimpleNamespace( + weight=torch.zeros((4, 3), dtype=torch.float8_e4m3fn), + weight_scale=torch.ones(1), + input_scale=torch.ones(1), + tp_group="tp", + orig_dtype=torch.float16, + ) + fp8 = torch.zeros((2, 4), dtype=torch.float8_e4m3fn) + seen = [] + + def bmm(x, weight, x_scale, weight_scale, output_dtype): + seen.append(output_dtype) + return torch.zeros((2, 3), dtype=output_dtype) + + monkeypatch.setattr( + lowm, + "row_parallel", + lambda *args, **kwargs: pytest.fail("FP16 entered the BF16-only low-M kernel"), + ) + monkeypatch.setattr(fp8_utils, "flashinfer_bmm_fp8", bmm, raising=False) + monkeypatch.setattr( + pynccl_allocator, + "use_symmetric_memory", + lambda *args, **kwargs: contextlib.nullcontext(), + ) + monkeypatch.setattr(dp_attention, "is_allocation_symmetric", lambda: True) + monkeypatch.setattr( + runtime_context, "get_parallel", lambda: SimpleNamespace(tp_group=None) + ) + st = SimpleNamespace(slots=SimpleNamespace(lowm_outproj=True)) + + out = verify._out_proj(layer, fp8, torch.bfloat16, st) + assert out.dtype == torch.float16 + assert seen == [torch.float16] + + @pytest.mark.parametrize("takes", [False, True]) def test_the_gdn_push_slot_pushes_only_where_the_moe_input_consumes( monkeypatch, caplog, takes From 455d46d0261ed5bae55a36068c78af2d991c65de Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Wed, 28 Oct 2026 21:34:03 -0700 Subject: [PATCH 39/62] fix(qwen3.5): restore R11 CP fused batch cap --- .../qwen3_5/sm103/attention/q35_cp_fused.py | 10 +++++----- test/unit/test_qwen3_5_runtime.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py index d9daa3d1..6d81a690 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py @@ -47,11 +47,11 @@ _LANE_PLAN_MAX_BATCH = 32 -# Verify calls use the wide and two-slot plans through 128 requests; draft calls use -# the wide plan through 64. Batches at or below 32 keep the equal split, so the -# separate balanced-split optimization remains disabled. -MAX_BATCH = 128 -MAX_BATCH_Q1 = 64 +# Keep the production selector on R11's unbalanced lane plan. The wide and +# two-slot kernels remain available to explicit test hooks, but R11 did not +# select them when SGLANG_PLB_CP_BALANCED was disabled. +MAX_BATCH = _LANE_PLAN_MAX_BATCH +MAX_BATCH_Q1 = _LANE_PLAN_MAX_BATCH # Pack every two-slot call; the unpacked form remains available as a test hook. _TWO_SLOT_PACK = True diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 297269f7..04013233 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -170,6 +170,24 @@ def test_requested_cp_state_does_not_restore_a_device_length_gate(): assert [cp._gate_thr(batch) for batch in (2, 4, 16)] == [None, None, None] +def test_cp_exact_selector_keeps_r11s_unbalanced_batch_cap(): + from sgl_nvidia_mega.qwen3_5.sm103.attention import q35_cp_fused as fused + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + cp = CpLite.__new__(CpLite) + cp.cp = 4 + cp.hq = 4 + cp.sms = 132 + + assert (fused.MAX_BATCH, fused.MAX_BATCH_Q1) == (32, 32) + assert [cp._exact_supported(7, batch) for batch in (32, 64, 128)] == [ + True, + False, + False, + ] + assert [cp._exact_supported(1, batch) for batch in (32, 64)] == [True, False] + + def test_cp_gate_epilogue_offer_is_consumed_once(): from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite From fdcce94caaba637395e001bc8ee2a75c7b252c6e Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Wed, 28 Oct 2026 22:18:03 -0700 Subject: [PATCH 40/62] fix(qwen3.5): filter CP slots by attention topology --- sgl_nvidia_mega/qwen3_5/__init__.py | 23 +++++ sgl_nvidia_mega/qwen3_5/draft/__init__.py | 10 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 37 ++++++-- .../qwen3_5/sm103/attention/q35_cp_lite.py | 15 ++- sgl_nvidia_mega/qwen3_5/spec.py | 13 +++ test/unit/test_qwen3_5_runtime.py | 71 +++++++++++++- test/unit/test_tp2_cp_build.py | 94 +++++++++++++++++++ 7 files changed, 244 insertions(+), 19 deletions(-) create mode 100644 test/unit/test_tp2_cp_build.py diff --git a/sgl_nvidia_mega/qwen3_5/__init__.py b/sgl_nvidia_mega/qwen3_5/__init__.py index 9ba3b849..697b402e 100644 --- a/sgl_nvidia_mega/qwen3_5/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/__init__.py @@ -29,3 +29,26 @@ hidden_act="silu", rms_norm_eps=1e-6, ) + +CP_LITE_HEAD_DIM = 256 + + +def cp_lite_topology( + attention_tp_size: int, + *, + num_key_value_heads: int = TEXT_CONFIG["num_key_value_heads"], + num_attention_heads: int = TEXT_CONFIG["num_attention_heads"], + head_dim: int = TEXT_CONFIG["head_dim"], + kernel_head_dim: int = CP_LITE_HEAD_DIM, +) -> tuple[int, int] | None: + """Return ``(CP size, local Q heads)`` when CP-lite supports the topology.""" + tp = attention_tp_size + if tp <= num_key_value_heads or tp % num_key_value_heads: + return None + if head_dim != kernel_head_dim: + return None + cp = tp // num_key_value_heads + q_heads = num_attention_heads // tp + if cp not in (2, 4, 8) or q_heads * cp > 16: + return None + return cp, q_heads diff --git a/sgl_nvidia_mega/qwen3_5/draft/__init__.py b/sgl_nvidia_mega/qwen3_5/draft/__init__.py index 15883065..13473c54 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/draft/__init__.py @@ -70,7 +70,11 @@ "bind_extend_graphs", "begin_draft_round", ), - "srt/runtime_context.py": ("get_parallel", '"moe_ep_size"'), + "srt/runtime_context.py": ( + "get_parallel", + '"attn_tp_size"', + '"moe_ep_size"', + ), } ARCHITECTURE = "Qwen3_5ForCausalLMMTP" @@ -214,7 +218,9 @@ def _build(model, hf_config, ctx) -> MtpDraft: from sgl_nvidia_mega.qwen3_5 import runtime - runtime.apply_expert_parallel(get_parallel().moe_ep_size) + parallel = get_parallel() + runtime.apply_expert_parallel(parallel.moe_ep_size) + runtime.apply_cp_lite_topology(parallel.attn_tp_size) runtime.build_cp_lite(ctx.binding) return MtpDraft( model, diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 381bacb1..c02f0d5d 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -119,6 +119,25 @@ def _requested_slots() -> frozenset: SLOTS = _requested_slots() +def apply_cp_lite_topology(attention_tp_size: int) -> frozenset: + """Drop CP-lite slots when the attention-TP Q/KV topology cannot serve them.""" + from sgl_nvidia_mega.qwen3_5 import cp_lite_topology + + global SLOTS + if cp_lite_topology(attention_tp_size) is not None: + return frozenset() + dropped = SLOTS & CP_LITE_SLOTS + SLOTS = SLOTS - CP_LITE_SLOTS + if dropped: + logger.info( + "NVIDIA megakernel: Qwen3.5 attention-TP %d: CP slots %s use stock " + "(unsupported Q/KV-head topology)", + attention_tp_size, + sorted(dropped), + ) + return dropped + + def build_cp_lite(binding) -> None: """Collectively build the CP slots' process-local state before validation. @@ -361,7 +380,12 @@ def slot_reasons() -> list[str]: "deadline_s", "binding", ), - "srt/runtime_context.py": ("get_parallel", '"tp_size"', '"moe_ep_size"'), + "srt/runtime_context.py": ( + "get_parallel", + '"tp_size"', + '"attn_tp_size"', + '"moe_ep_size"', + ), "srt/kernel_plugins/nvidia/eagle_binding.py|srt/megakernels/eagle_binding.py": ( "page_size", "tp_rank", @@ -393,19 +417,20 @@ def build(model, ctx): from sgl_nvidia_mega.qwen3_5.runtime import hooks, parity from sgl_nvidia_mega.qwen3_5.runtime.executor import Qwen35Verify - build_cp_lite(ctx.binding) - - dropped = apply_expert_parallel(get_parallel().moe_ep_size) + parallel = get_parallel() + dropped = apply_expert_parallel(parallel.moe_ep_size) if dropped: logger.info( "NVIDIA megakernel: Qwen3.5 under expert parallelism (moe_ep_size %d): " "slots %s run SGLang's EP path", - get_parallel().moe_ep_size, + parallel.moe_ep_size, sorted(dropped), ) + apply_cp_lite_topology(parallel.attn_tp_size) + build_cp_lite(ctx.binding) multi = slot_plan.MULTI_REQUEST in SLOTS if multi: - multi = slot_plan.multi_request(get_parallel().tp_size, SLOTS) + multi = slot_plan.multi_request(parallel.tp_size, SLOTS) widths = slot_plan.candidates(ctx.widths, ctx.tokens_per_request, multi) if not widths: raise RuntimeError( diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py index 4e546958..731f162d 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py @@ -29,9 +29,11 @@ import triton import triton.language as tl +from sgl_nvidia_mega.qwen3_5 import CP_LITE_HEAD_DIM, cp_lite_topology + logger = logging.getLogger(__name__) -PAGE, HEAD_DIM = 64, 256 +PAGE, HEAD_DIM = 64, CP_LITE_HEAD_DIM _MAX_ROWS = 896 # widest call: rows = requests * q_len (TP8 verify at B128: 896) _KINDS = 2 # flag regions: q gather, partial exchange # exact path's fused partial: E4M3 P with shift 8 against the exact running max (lazy rescale threshold 0), so the @@ -1143,18 +1145,13 @@ def maybe_create(cls, binding, device) -> CpLite | None: those model-specific facts here avoids a runtime-side prebuild hook. Every rank calls this from plugin build before graph variants are validated. """ - from sgl_nvidia_mega.qwen3_5 import TEXT_CONFIG - if cls._instance is not None: return cls._instance tp, rank = binding.tp_size, binding.tp_rank - kv = TEXT_CONFIG["num_key_value_heads"] - if tp <= kv or tp % kv or TEXT_CONFIG["head_dim"] != HEAD_DIM: - return None - cp = tp // kv - q_heads = TEXT_CONFIG["num_attention_heads"] // tp - if cp not in (2, 4, 8) or q_heads * cp > 16: + topology = cp_lite_topology(tp, kernel_head_dim=HEAD_DIM) + if topology is None: return None + cp, q_heads = topology group = binding.tp_device_group cls._instance = cls(group, rank, tp, cp, q_heads, device) return cls._instance diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index 45860645..580fe4b3 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -61,9 +61,13 @@ "linear_attn_verify_backend", "enable_linear_replayssm_spec", "mamba_ssm_dtype", + "tp_size", + "attn_cp_size", + "attn_dp_size", "startup_weight_load_mode", "disaggregation_mode", ), + "srt/runtime_context.py": ("derive_attn_tp_size",), "srt/configs/qwen3_5.py": ("text_config",), "srt/kernel_plugins/nvidia/eagle_binding.py|srt/megakernels/eagle_binding.py": ( "EagleForwardView", @@ -240,7 +244,16 @@ def pd_reasons(cfg) -> list[str]: def _config_disable_reasons(cfg, hf_config) -> list[str]: + from sglang.srt.runtime_context import derive_attn_tp_size + runtime.apply_expert_parallel(cfg.ep_size or 1) + runtime.apply_cp_lite_topology( + derive_attn_tp_size( + tp_size=cfg.tp_size, + attn_cp_size=cfg.attn_cp_size, + attn_dp_size=cfg.attn_dp_size, + ) + ) return ( text_config_reasons(hf_config) + binding_reasons() diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 04013233..9487e0fe 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -31,7 +31,7 @@ def registered(monkeypatch): # the target and draft builds read moe_ep_size (expert parallelism); no process groups here monkeypatch.setattr( "sglang.srt.runtime_context.get_parallel", - lambda: SimpleNamespace(tp_size=8, moe_ep_size=1), + lambda: SimpleNamespace(tp_size=8, attn_tp_size=8, moe_ep_size=1), ) @@ -131,6 +131,73 @@ def test_cp_slots_refuse_an_incompatible_page_or_topology(monkeypatch): runtime.build_cp_lite(SimpleNamespace(page_size=64, tp_size=8)) +@pytest.mark.parametrize("attention_tp_size", [1, 2, 3, 6, 32]) +def test_unsupported_attention_topology_drops_only_cp_slots( + monkeypatch, caplog, attention_tp_size +): + requested = runtime.CP_LITE_SLOTS | {"skv2"} + monkeypatch.setattr(runtime, "SLOTS", requested) + + with caplog.at_level("INFO"): + dropped = runtime.apply_cp_lite_topology(attention_tp_size) + + assert dropped == runtime.CP_LITE_SLOTS + assert runtime.SLOTS == frozenset({"skv2"}) + assert f"attention-TP {attention_tp_size}" in caplog.text + + +@pytest.mark.parametrize("attention_tp_size", [4, 8, 16]) +def test_supported_attention_topology_keeps_cp_slots(monkeypatch, attention_tp_size): + requested = runtime.CP_LITE_SLOTS | {"skv2"} + monkeypatch.setattr(runtime, "SLOTS", requested) + + assert runtime.apply_cp_lite_topology(attention_tp_size) == frozenset() + assert runtime.SLOTS == requested + + +def test_config_filter_derives_attention_tp_before_publication(monkeypatch): + from sgl_nvidia_mega.qwen3_5 import spec as target + + calls = [] + monkeypatch.setattr( + runtime, + "apply_expert_parallel", + lambda ep_size: calls.append(("ep", ep_size)) or frozenset(), + ) + monkeypatch.setattr( + runtime, + "apply_cp_lite_topology", + lambda attention_tp_size: ( + calls.append(("attention_tp", attention_tp_size)) or frozenset() + ), + ) + monkeypatch.setattr(target, "text_config_reasons", lambda hf_config: []) + monkeypatch.setattr(target, "binding_reasons", lambda: []) + monkeypatch.setattr(runtime, "slot_reasons", lambda: []) + monkeypatch.setattr(target, "pd_reasons", lambda cfg: []) + cfg = SimpleNamespace(tp_size=8, attn_cp_size=1, attn_dp_size=8, ep_size=8) + + assert not hasattr(cfg, "attn_tp_size") + assert target._config_disable_reasons(cfg, SimpleNamespace()) == [] + assert calls == [("ep", 8), ("attention_tp", 1)] + + +@pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 8, 16, 32]) +def test_cp_topology_predicate_matches_state_creation(monkeypatch, tp): + from sgl_nvidia_mega.qwen3_5 import cp_lite_topology + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + made = [] + binding = SimpleNamespace(tp_size=tp, tp_rank=0, tp_device_group=object()) + monkeypatch.setattr(CpLite, "_instance", None) + monkeypatch.setattr(CpLite, "__init__", lambda self, *args: made.append(args)) + + state = CpLite.maybe_create(binding, torch.device("cuda", 0)) + + assert (state is not None) == (cp_lite_topology(tp) is not None) + assert bool(made) == (cp_lite_topology(tp) is not None) + + @pytest.mark.parametrize( "tp, rank, want", [(4, 3, (3, 4, 2, 8)), (8, 7, (7, 8, 4, 4))], @@ -337,7 +404,7 @@ def test_the_build_covers_more_requests_only_where_a_slot_runs( monkeypatch.setattr(runtime, "SLOTS", frozenset()) monkeypatch.setattr( "sglang.srt.runtime_context.get_parallel", - lambda: SimpleNamespace(tp_size=tp_size, moe_ep_size=1), + lambda: SimpleNamespace(tp_size=tp_size, attn_tp_size=tp_size, moe_ep_size=1), ) model = SimpleNamespace( model=SimpleNamespace(start_layer=0, end_layer=0, layers=[]) diff --git a/test/unit/test_tp2_cp_build.py b/test/unit/test_tp2_cp_build.py new file mode 100644 index 00000000..04717844 --- /dev/null +++ b/test/unit/test_tp2_cp_build.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""CPU regression for the actual #4162 CP build layouts; applies to pure v3. + +DEP8's actual EagleBinding exposes global TP8, not attention TP1. Before the +fix its unused CP state could initialize; TP2 instead raised. Topology filtering +must precede collective initialization in either case. +""" + +from types import SimpleNamespace + +import pytest +import torch + + +@pytest.mark.parametrize( + "name,tp,dp,ep,batch,expects_cp", + [ + ("agg-tp2", 2, 1, 1, 1, False), + ("pd-d-tp2-b6", 2, 1, 1, 6, False), + ("pd-d-tp2-b10", 2, 1, 1, 10, False), + ("agg-tp4", 4, 1, 1, 1, True), + ("agg-tp8", 8, 1, 1, 1, True), + ("pd-dep8", 8, 8, 8, 32, False), + ], +) +def test_actual_cp_layout_filters_before_target_build( + monkeypatch, name, tp, dp, ep, batch, expects_cp +): + from sglang.srt.kernel_plugins.nvidia import binding as plugin_binding + + from sgl_nvidia_mega.qwen3_5 import runtime + from sgl_nvidia_mega.qwen3_5 import spec as target + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + # Exact current B env, including CP names. No downstream topology import: + # this file applies without changes to the upstream c9d3 source tree. + requested_csv = ( + "route_topk,mkm_m32,mkm_m64,mkm_m128,mkm_m256,mkm_m512,gdn_verify," + "gdn_pdl,gdn_conv,hop,hop_fp8,multi_request,draft_loop,draft_extend," + "draft_multi,draft_embed,route,route_gemm,mkm,mkm_wide,gdn_core," + "hop_mid,gdn_push,gdn_front,gdn_trt_front,route_shared,route_lin,skv2," + "mka,mka_push,lowm_in,lowm_out,lowm_se,lowm_push,skv2_cp,mkm_wide_bal," + "draft_moe,draft_q1,draft_xchg,draft_hop,draft_fin,draft_front," + "draft_skv2,draft_q1_cp,draft_skv2_cp,sharded_gumbel," + "gdn_prefill_split_conv" + ) + requested = frozenset(requested_csv.split(",")) + monkeypatch.setattr(runtime, "SLOTS", requested) + parallel = SimpleNamespace( + tp_size=tp, + tp_rank=0, + attn_dp_size=dp, + attn_tp_size=tp // dp, + attn_tp_rank=0, + moe_ep_size=ep, + tp_group=SimpleNamespace(device_group="group"), + ) + monkeypatch.setattr("sglang.srt.runtime_context.get_parallel", lambda: parallel) + monkeypatch.setattr(plugin_binding, "megakernel_agree", lambda *a, **k: None) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(CpLite, "_instance", None) + allocations = [] + monkeypatch.setattr( + CpLite, "__init__", lambda self, *args: allocations.append(args) + ) + # Actual EagleBinding.tp_size is global TP, including DEP8. + binding = SimpleNamespace( + page_size=64, tp_size=tp, tp_rank=0, tp_device_group="group" + ) + original = runtime.build_cp_lite + checked = [] + + class SetupChecked(Exception): + pass + + def setup(actual_binding): + assert actual_binding is binding + original(actual_binding) + assert bool(runtime.SLOTS & runtime.CP_LITE_SLOTS) == expects_cp, name + assert bool(allocations) == expects_cp, name + checked.append(True) + raise SetupChecked + + monkeypatch.setattr(runtime, "build_cp_lite", setup) + ctx = SimpleNamespace( + widths=(batch * 7,), + tokens_per_request=7, + deadline_s=60.0, + binding=binding, + ) + with pytest.raises(SetupChecked): + target.SPEC.build(SimpleNamespace(), SimpleNamespace(), ctx) + assert checked == [True] From b1f6e08ef3138b1355553d3a1f721747ca4efc9f Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Wed, 28 Oct 2026 22:36:19 -0700 Subject: [PATCH 41/62] fix(qwen3.5): cap target CP route at fused batch limit --- sgl_nvidia_mega/qwen3_5/attention/split_kv.py | 5 ++++- test/unit/test_qwen3_5_steps.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py index 85638184..382f2e88 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py +++ b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py @@ -241,8 +241,11 @@ def cp_lite(): def cp_multi(rows: int) -> bool: """Whether a multi-request target verify can try the plugin's CP state.""" from sgl_nvidia_mega.qwen3_5.runtime import SLOTS + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_fused import MAX_BATCH - return "skv2_cp" in SLOTS and rows > ROWS and cp_lite() is not None + return ( + "skv2_cp" in SLOTS and ROWS < rows <= ROWS * MAX_BATCH and cp_lite() is not None + ) def cp_reason() -> str: diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 20b19b6a..89eab5d5 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -2529,6 +2529,21 @@ def test_draft_only_cp_slots_do_not_enable_target_multi_request(monkeypatch): assert split_kv.cp_multi(56) +@pytest.mark.parametrize( + ("requests", "eligible"), + [(5, True), (32, True), (33, False), (64, False), (128, False)], +) +def test_target_cp_multi_request_honors_fused_batch_cap( + monkeypatch, requests, eligible +): + from sgl_nvidia_mega.qwen3_5 import runtime + from sgl_nvidia_mega.qwen3_5.attention import split_kv + + monkeypatch.setattr(split_kv, "cp_lite", lambda: object()) + monkeypatch.setattr(runtime, "SLOTS", frozenset({"skv2_cp"})) + assert split_kv.cp_multi(requests * split_kv.ROWS) is eligible + + @pytest.mark.parametrize( "kw, want", [ From a2384abebe2dda839c8dd8fec4f820e46fff5b7b Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Wed, 28 Oct 2026 22:41:19 -0700 Subject: [PATCH 42/62] test(qwen3.5): model attention TP in route fixture --- test/unit/test_qwen3_5_steps.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 89eab5d5..cda1e943 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -946,7 +946,7 @@ def test_the_route_slot_is_not_built_for_a_layer_it_cannot_serve(monkeypatch): # runtime.build reads moe_ep_size (expert parallelism); no process groups here monkeypatch.setattr( "sglang.srt.runtime_context.get_parallel", - lambda: SimpleNamespace(tp_size=8, moe_ep_size=1), + lambda: SimpleNamespace(tp_size=8, attn_tp_size=8, moe_ep_size=1), ) ctx = SimpleNamespace( widths=(7,), From 3440d8173296c8c18b593b52e6dc0d09ffd77417 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Thu, 29 Oct 2026 10:44:37 -0700 Subject: [PATCH 43/62] perf(qwen3.5): use plain CP merge for multi-request draft attention Route multi-request draft CP-lite calls through a standalone plain merge-and-wait kernel equivalent to the legacy R11 implementation. Keep target verify and single-request draft calls on the existing EARLY/gate-epilogue path. The plain exact and fold variants are bitwise identical to R11 on the validated inputs. This removes about 30 us from each TEP8 B32 draft round without regressing verify. --- sgl_nvidia_mega/qwen3_5/attention/split_kv.py | 17 +- sgl_nvidia_mega/qwen3_5/draft/attention.py | 3 + .../qwen3_5/sm103/attention/q35_cp_lite.py | 195 ++++++++++++++- test/kernels/test_qwen3_5_attention.py | 230 ++++++++++++++++++ test/unit/test_qwen3_5_runtime.py | 228 ++++++++++++++++- 5 files changed, 661 insertions(+), 12 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py index 382f2e88..b51bd2f1 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/split_kv.py +++ b/sgl_nvidia_mega/qwen3_5/attention/split_kv.py @@ -253,11 +253,22 @@ def cp_reason() -> str: return "" if cp_lite() is not None else "the plugin did not build CP-lite" -def cp_core(attn, q, kv, meta, forward_batch, q_len: int, min_bs: int | None = 1): +def cp_core( + attn, + q, + kv, + meta, + forward_batch, + q_len: int, + min_bs: int | None = 1, + early_merge: bool | None = None, +): """CP-lite on requests of ``q_len`` FP8 query rows over the pool ``kv``: [rows, heads * head_dim] bf16, or None where CP-lite declines. ``min_bs`` overrides CP-lite's configured minimum for a host-selected one-request graph; ``None`` keeps - the configured gate for a multi-request graph.""" + the configured gate for a multi-request graph. ``early_merge=False`` restores the + R11 plain merge for both exact and fold paths when no fused gate epilogue was + offered; ``None`` preserves CP-lite's default.""" from sglang.kernels.ops.attention.utils import canonicalize_stride from sglang.srt.layers.attention.trtllm_mha_backend import ( DEFAULT_WORKSPACE_SIZE_MB, @@ -281,6 +292,7 @@ def cp_core(attn, q, kv, meta, forward_batch, q_len: int, min_bs: int | None = 1 DEFAULT_WORKSPACE_SIZE_MB * 1024 * 1024, dtype=torch.uint8, device=q.device ), ) + options = {} if early_merge is None else {"early_merge": early_merge} o = cp.attend( q.reshape(-1, attn.tp_q_head_num, attn.head_dim), caches, @@ -293,6 +305,7 @@ def cp_core(attn, q, kv, meta, forward_batch, q_len: int, min_bs: int | None = 1 attn, forward_batch, min_bs=min_bs, + **options, ) return None if o is None else o.view(-1, attn.tp_q_head_num * attn.head_dim) diff --git a/sgl_nvidia_mega/qwen3_5/draft/attention.py b/sgl_nvidia_mega/qwen3_5/draft/attention.py index de3b322a..552ed214 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/attention.py +++ b/sgl_nvidia_mega/qwen3_5/draft/attention.py @@ -467,6 +467,9 @@ def run(self, attention, q8, meta, kv, rows_per_request: int): attention, rows_per_request, min_bs=None if multi else 1, + # R11 orders the draft-multi exact partial and merge after their waits. + # Keep the shared CP-lite default for target and one-request draft calls. + early_merge=False if multi else None, ) if o is not None and torch.cuda.is_current_stream_capturing(): key = (rows_per_request, q8.shape[0]) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py index 731f162d..bf3e6755 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_lite.py @@ -471,6 +471,16 @@ def _ll_wait(src, n: tl.constexpr, e): return v +@triton.jit +def _ll_wait_plain(src, n: tl.constexpr, e): + """R11's volatile-load wait, used only by the R11 plain merge.""" + k = tl.arange(0, n) + v = tl.load(src + k, volatile=True) + while tl.min(((v >> 32) == e).to(tl.int32), axis=0) == 0: + v = tl.load(src + k, volatile=True) + return v + + @triton.jit def _cp_gather_q_ll( q, @@ -988,6 +998,163 @@ def _cp_merge_h_ll( tl.store(cnt + i * HQ + h + one, e) +@triton.jit +def _cp_merge_h_ll_plain( + po, + plse, + po_sr, + po_sc, + pl_sr, + pl_sc, + q, + q_stride, + k_cache, + v_cache, + kv_page_stride, + kv_tok_stride, + table, + t_stride, + seq_lens, + out, + out_stride, + peers, + me, + ll_off, + cnt, + bmm1, + bmm2, + nb, + thr, + QL: tl.constexpr, + HQ: tl.constexpr, + CP: tl.constexpr, + D: tl.constexpr, + PG: tl.constexpr, + NWP: tl.constexpr, + NB: tl.constexpr, + GMAX: tl.constexpr, + WAIT: tl.constexpr, + po_v=None, + pl_v=None, + vstart=None, + vcnt=None, + CHUNK: tl.constexpr = False, +): + """R11's plain non-EARLY merge, kept separate from EARLY/gate epilogues. + + The separate JIT function is intentional: it gives draft calls the same + compile-time program as R11 instead of relying on dead-code elimination in + the feature-rich shared merge. + """ + i = tl.program_id(0) + h = tl.program_id(1) + if WAIT: + _gdc_wait_clobber() + if _gate_on(seq_lens, nb, thr, NB, GMAX).to(tl.int32) == 0: + return + one = tl.arange(0, 1) + e = tl.load(cnt + i * HQ + h) + 1 + flag = e.to(tl.int64) << 32 + DW: tl.constexpr = D // 2 + NO: tl.constexpr = HQ * D // 2 + H: tl.constexpr = QL * CP * HQ + w = tl.arange(0, DW) + if CHUNK: + n_v = tl.load(vcnt + i // QL) + v0 = tl.load(vstart + i // QL) + for c in tl.static_range(CP): + if CHUNK: + if n_v > 0: + hh = (i % QL) * (CP * HQ) + c * HQ + h + mx = tl.full([1], float("-inf"), tl.float32) + sw = tl.zeros([1], tl.float32) + acc_lo = tl.zeros([DW], tl.float32) + acc_hi = tl.zeros([DW], tl.float32) + for j in range(0, n_v): + vr = v0 + j + lj = tl.load(pl_v + vr * H + hh + one) + wd = tl.load(po_v + (vr * H + hh).to(tl.int64) * DW + w).to( + tl.uint32, bitcast=True + ) + mn = tl.maximum(mx, lj) + a = tl.exp2(mx - mn) + b2 = tl.exp2(lj - mn) + acc_lo = acc_lo * a + (wd << 16).to(tl.float32, bitcast=True) * b2 + acc_hi = ( + acc_hi * a + (wd & 0xFFFF0000).to(tl.float32, bitcast=True) * b2 + ) + sw = sw * a + b2 + mx = mn + inv = 1.0 / tl.where(sw > 0, sw, 1.0) + lo16 = (acc_lo * inv).to(tl.bfloat16).to(tl.int16, bitcast=True).to( + tl.int32 + ) & 0xFFFF + hi16 = (acc_hi * inv).to(tl.bfloat16).to(tl.int16, bitcast=True).to( + tl.int32 + ) & 0xFFFF + lse2 = tl.where( + sw > 0, + mx + tl.log2(tl.where(sw > 0, sw, 1.0)), + float("-inf"), + ) + ow = _ll_pack(flag, lo16 | (hi16 << 16)) + lw = _ll_pack(flag, lse2.to(tl.int32, bitcast=True)) + else: + ow = _ll_pack(flag, tl.load(po + i * po_sr + c * po_sc + h * DW + w)) + lw = _ll_pack(flag, tl.load(plse + i * pl_sr + c * pl_sc + h + one)) + else: + ow = _ll_pack(flag, tl.load(po + i * po_sr + c * po_sc + h * DW + w)) + lw = _ll_pack(flag, tl.load(plse + i * pl_sr + c * pl_sc + h + one)) + dst = (tl.load(peers + c) + ll_off + (i * CP + me) * NWP * 8).to( + tl.pointer_type(tl.int64) + ) + tl.store(dst + h * DW + w, ow) + tl.store(dst + NO + h + one, lw) + L = tl.load(seq_lens + i // QL) + acc, m, wsum = _causal_block_h( + q, + q_stride, + i, + h, + L, + k_cache, + v_cache, + kv_page_stride, + kv_tok_stride, + table, + t_stride, + bmm1, + bmm2, + QL, + D, + PG, + ) + n_pages = (tl.maximum(L - QL, 0) + PG - 1) // PG + own = tl.load(peers + me) + ll_off + for c in tl.static_range(CP): + src = (own + (i * CP + c) * NWP * 8).to(tl.pointer_type(tl.int64)) + v = _ll_wait_plain(src + h * DW, DW, e) + lv = _ll_wait_plain(src + NO + h, 1, e) + lc = (lv & 0xFFFFFFFF).to(tl.int32).to( + tl.float32, bitcast=True + ) * 0.6931471805599453 + lc = tl.where(n_pages > c, lc, float("-inf")) + wd = (v & 0xFFFFFFFF).to(tl.uint32) + lo = (wd << 16).to(tl.float32, bitcast=True) + hi = (wd & 0xFFFF0000).to(tl.float32, bitcast=True) + oc = tl.reshape(tl.join(lo, hi), (D,)) + mn = tl.maximum(m, lc) + mn0 = tl.where(mn > float("-inf"), mn, 0.0) + a_old = tl.where(m > float("-inf"), tl.exp(m - mn0), 0.0) + a_new = tl.where(lc > float("-inf"), tl.exp(lc - mn0), 0.0) + acc = acc * a_old + oc * a_new + wsum = wsum * a_old + a_new + m = mn + o = acc / tl.where(wsum > 0, wsum, 1.0) + tl.store(out + i * out_stride + h * D + tl.arange(0, D), o.to(tl.bfloat16)) + tl.store(cnt + i * HQ + h + one, e) + + class CpLite: """One process's exchange buffer and scratch, shared by the Qwen3.5 target and draft slots; built collectively on the plugin binding's TP group.""" @@ -1210,13 +1377,17 @@ def attend( forward_batch, stock=None, min_bs: int | None = None, + early_merge: bool | None = None, ) -> torch.Tensor | None: """CP-lite attention of this rank's heads: [rows, heads, head_dim] bf16, or None where the call is not served (decided by shapes only, so every rank agrees). stock(seq_lens, out): the caller's own kernel call into out; with it and a KV gate for this batch size, both are captured and the device picks one per call. min_bs: override the smallest batch this call serves; a caller that picks CP by - its own host-side KV length can pass 1.""" + its own host-side KV length can pass 1. + early_merge: override the exact ungated partial/merge ordering for this call; + without a fused gate epilogue, False also selects the R11 plain merge for + exact and non-exact fold calls; None keeps the shared default.""" rows = q.shape[0] b = seq_lens.shape[0] if ( @@ -1291,6 +1462,7 @@ def attend( bmm1_scale, bmm2_scale, gate, + early_merge=early_merge, ) if self.ll: _cp_gather_q_ll[(rows,)]( @@ -1391,6 +1563,7 @@ def attend( bmm2_scale, out, fold_gate, + plain=early_merge is False, ) elif self.ll: self._fold_merge( @@ -1795,8 +1968,10 @@ def _ungated_main( bmm1_scale, bmm2_scale, gate, + early_merge=None, ): """Run an ungated exact call with PDL and the per-head merge.""" + early = self.early_merge if early_merge is None else early_merge words = self.hq * HEAD_DIM // 4 _cp_gather_q_ll[(rows,)]( q.view(torch.int32), @@ -1832,7 +2007,7 @@ def _ungated_main( bmm2_scale, pdl=True, len_sub=q_len, - early_trigger=self.early_merge, + early_trigger=early, ) out = torch.empty( rows, self.hq, HEAD_DIM, dtype=torch.bfloat16, device=self.device @@ -1845,6 +2020,7 @@ def _ungated_main( and epilogue[2].stride(1) == 1 ): epilogue = None + plain = early_merge is False and epilogue is None self._merge_h( q, k_cache, @@ -1858,7 +2034,8 @@ def _ungated_main( out, {**gate, "thr": 0}, pdl=True, - early=self.early_merge, + early=early, + plain=plain, gate_epilogue=epilogue, ) if epilogue is not None: @@ -1886,6 +2063,7 @@ def _merge_h( gate, pdl=None, early=False, + plain=False, gate_epilogue=None, ): pdl = self.cp_pdl if pdl is None else pdl @@ -1902,7 +2080,9 @@ def _merge_h( "GATE_EPILOGUE": True, } ) - _cp_merge_h_ll[(rows, hq)]( + merge_kernel = _cp_merge_h_ll_plain if plain else _cp_merge_h_ll + merge_options = {} if plain else {"EARLY": early and pdl, **epilogue_args} + merge_kernel[(rows, hq)]( self.po_x.view(torch.int32), self.pl_x.view(torch.int32), hq * HEAD_DIM // 2, @@ -1937,10 +2117,9 @@ def _merge_h( NB=gate["NB"], GMAX=gate["GMAX"], WAIT=pdl, - EARLY=early and pdl, num_warps=4, launch_pdl=pdl, - **epilogue_args, + **merge_options, ) def _own_workspace(self): @@ -2161,10 +2340,12 @@ def _fold_merge_h( out, gate, fused=None, + plain=False, ): """Merge the fold layout with one program per row and query head.""" hq = self.hq - _cp_merge_h_ll[(rows, hq)]( + merge_kernel = _cp_merge_h_ll_plain if plain else _cp_merge_h_ll + merge_kernel[(rows, hq)]( self.po.view(torch.int32), self.plse.view(torch.int32), self.cp * hq * HEAD_DIM // 2, diff --git a/test/kernels/test_qwen3_5_attention.py b/test/kernels/test_qwen3_5_attention.py index 823c6b99..c54768c9 100644 --- a/test/kernels/test_qwen3_5_attention.py +++ b/test/kernels/test_qwen3_5_attention.py @@ -172,6 +172,236 @@ def test_skv2_as_a_programmatic_dependent_gives_the_same_bytes(length): assert torch.equal(outs[0].view(torch.int16), outs[1].view(torch.int16)) +@pytest.mark.parametrize("chunk", [False, True], ids=["exact", "fold-chunk"]) +def test_cp_r11_plain_merge_matches_r11_at_b32(chunk): + """The draft-multi B32 merge's causal math is independent of whether it runs + before or after publishing the prefix partial. The D2 plain kernel is the R11 + program (including its volatile-load wait); EARLY is retained here only as a + numerical comparison, not as the reference. The direct FP32 calculation is + diagnostic-only because the fold path merges its chunks in a different order; + the correctness gate is the plain kernel against the exact R11 program. + + CP=1 makes the low-latency exchange self-contained on one GPU: every word this + kernel polls is published by the same launch. PDL liveness and the production + CP=4 exchange remain covered by the no-simulation server gate. + """ + import hashlib + import importlib.util + import os + import sys + from pathlib import Path + + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import ( + _cp_merge_h_ll as feature_merge, + _cp_merge_h_ll_plain as plain_merge, + ) + + r11_path = os.environ.get("V3D_R11_Q35_CP_LITE") + r11_sha = os.environ.get("V3D_R11_Q35_CP_LITE_SHA256") + if r11_path is None and r11_sha is None: + pytest.skip("exact R11 q35_cp_lite source was not provided") + assert r11_path is not None and r11_sha is not None + r11_source = Path(r11_path) + assert r11_source.is_file() + actual_sha = hashlib.sha256(r11_source.read_bytes()).hexdigest() + assert actual_sha == r11_sha, (actual_sha, r11_sha) + spec = importlib.util.spec_from_file_location("_v3d_r11_q35_cp_lite", r11_source) + assert spec is not None and spec.loader is not None + r11_module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = r11_module + spec.loader.exec_module(r11_module) + r11_merge = r11_module._cp_merge_h_ll + + rows, q_len, heads, dim, page, words = 32 * 7, 7, 4, 256, 64, 520 + g = torch.Generator(device="cuda").manual_seed(20261029) + q = (torch.randn(rows, heads, dim, device="cuda", generator=g) * 2).to( + torch.float8_e4m3fn + ) + po = (torch.randn(1, rows, heads, dim, device="cuda", generator=g) * 0.1).to( + torch.bfloat16 + ) + plse = torch.randn(1, rows, heads, device="cuda", generator=g) + chunks_per_request = 2 + po_v = ( + torch.randn( + 32 * chunks_per_request, + q_len * heads, + dim, + device="cuda", + generator=g, + ) + * 0.1 + ).to(torch.bfloat16) + pl_v = torch.randn( + 32 * chunks_per_request, + q_len * heads, + device="cuda", + generator=g, + ) + vstart = ( + torch.arange(32, dtype=torch.int32, device="cuda") * chunks_per_request + ) + vcnt = torch.full( + (32,), chunks_per_request, dtype=torch.int32, device="cuda" + ) + k, v = ( + (torch.randn(32, 1, page, dim, device="cuda", generator=g) * 2).to( + torch.float8_e4m3fn + ) + for _ in range(2) + ) + table = torch.zeros(32, 2, dtype=torch.int32, device="cuda") + table[:, 1] = torch.arange(32, dtype=torch.int32, device="cuda") + lengths = torch.full((32,), page + q_len, dtype=torch.int32, device="cuda") + + def run(kernel, early=None): + out = torch.empty(rows, heads, dim, dtype=torch.bfloat16, device="cuda") + ll = torch.empty(rows * words, dtype=torch.int64, device="cuda") + peers = torch.tensor([ll.data_ptr()], dtype=torch.int64, device="cuda") + counter = torch.zeros(rows * heads, dtype=torch.int32, device="cuda") + options = {} if early is None else {"EARLY": early} + if chunk: + options.update( + po_v=po_v.view(torch.int32), + pl_v=pl_v, + vstart=vstart, + vcnt=vcnt, + CHUNK=True, + ) + kernel[(rows, heads)]( + po.view(torch.int32), + plse.view(torch.int32), + heads * dim // 2, + po.stride(0) // 2, + heads, + plse.stride(0), + q, + q.stride(0), + k, + v, + k.stride(0), + k.stride(2), + table, + table.stride(0), + lengths, + out, + out.stride(0), + peers, + 0, + 0, + counter, + 0.75 * dim**-0.5, + 0.5, + 32, + 0, + QL=q_len, + HQ=heads, + CP=1, + D=dim, + PG=page, + NWP=words, + NB=32, + GMAX=False, + WAIT=False, + num_warps=4, + launch_pdl=False, + **options, + ) + torch.cuda.synchronize() + assert torch.isfinite(out).all() + assert bool((counter == 1).all()) + return out + + # Direct FP32 evaluation of the same causal block followed by its prefix partial. + # Elementwise products and reductions avoid TF32 matmul in the reference. + requests = rows // q_len + qf = q.float().view(requests, q_len, heads, dim) + page_ids = table[:, 1].long() + kf = k[page_ids, 0, :q_len].float() + vf = v[page_ids, 0, :q_len].float() + scores = (qf.unsqueeze(3) * kf[:, None, None]).sum(-1) * (0.75 * dim**-0.5) + causal = torch.arange(q_len, device="cuda") + scores.masked_fill_( + causal[None, None, None] > causal[None, :, None, None], -math.inf + ) + peak = scores.max(-1).values + weight = torch.exp(scores - peak.unsqueeze(-1)) + causal_sum = weight.sum(-1) + causal_out = (weight.unsqueeze(-1) * vf[:, None, None]).sum(-2) * 0.5 + if chunk: + chunk_lse = pl_v.view( + requests, chunks_per_request, q_len, heads + ).float() + chunk_peak = chunk_lse.max(1).values + chunk_weight = torch.exp2(chunk_lse - chunk_peak.unsqueeze(1)) + chunk_sum = chunk_weight.sum(1) + prefix_out = ( + po_v.float().view( + requests, chunks_per_request, q_len, heads, dim + ) + * chunk_weight.unsqueeze(-1) + ).sum(1) / chunk_sum.unsqueeze(-1) + prefix_lse = (chunk_peak + torch.log2(chunk_sum)) * math.log(2.0) + else: + prefix_out = po[0].float().view(requests, q_len, heads, dim) + prefix_lse = plse[0].view(requests, q_len, heads) * math.log(2.0) + merged_peak = torch.maximum(peak, prefix_lse) + causal_weight = torch.exp(peak - merged_peak) + prefix_weight = torch.exp(prefix_lse - merged_peak) + reference = ( + causal_out * causal_weight.unsqueeze(-1) + + prefix_out * prefix_weight.unsqueeze(-1) + ) / (causal_sum * causal_weight + prefix_weight).unsqueeze(-1) + reference = reference.view(rows, heads, dim) + + rtol, atol = 2**-7, 1e-6 + relative_floor = atol / rtol + rounded_reference = reference.to(torch.bfloat16) + + def stats(got): + error = (got.float() - reference).abs() + relative = error / reference.abs().clamp_min(relative_floor) + violation = error > (atol + rtol * reference.abs()) + return { + "max_abs": error.max().item(), + "max_rel": relative.max().item(), + "tolerance_violations": violation.count_nonzero().item(), + "bf16_mismatch": ( + got.view(torch.int16) != rounded_reference.view(torch.int16) + ) + .count_nonzero() + .item(), + } + + def pair_stats(got, baseline): + error = (got.float() - baseline.float()).abs() + relative = error / baseline.float().abs().clamp_min(relative_floor) + return { + "max_abs": error.max().item(), + "max_rel": relative.max().item(), + "bf16_mismatch": ( + got.view(torch.int16) != baseline.view(torch.int16) + ) + .count_nonzero() + .item(), + } + + r11 = run(r11_merge) + plain, early = run(plain_merge), run(feature_merge, True) + r11_stats = stats(r11) + plain_stats, early_stats = stats(plain), stats(early) + plain_vs_early = pair_stats(plain, early) + plain_vs_r11 = pair_stats(plain, r11) + print( + "DRAFT_MERGE_NUMERICS " + f"chunk={chunk} rtol={rtol:.9g} atol={atol:.9g} " + f"relative_floor={relative_floor:.9g} " + f"r11={r11_stats} plain={plain_stats} early={early_stats} " + f"plain_vs_early={plain_vs_early} plain_vs_r11={plain_vs_r11}" + ) + assert torch.equal(plain.view(torch.int16), r11.view(torch.int16)), plain_vs_r11 + + HIDDEN, QKV_ROWS, POOL_SLOTS = 4096, (2 * HEADS + 2) * DIM, 1024 diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 9487e0fe..de36b77c 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -255,6 +255,207 @@ def test_cp_exact_selector_keeps_r11s_unbalanced_batch_cap(): assert [cp._exact_supported(1, batch) for batch in (32, 64)] == [True, False] +def test_cp_call_override_pairs_the_exact_partial_and_merge_order(monkeypatch): + from sgl_nvidia_mega.qwen3_5.sm103.attention import q35_cp_lite as cp_module + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + class Kernel: + def __getitem__(self, grid): + return lambda *args, **kwargs: None + + monkeypatch.setattr(cp_module, "_cp_gather_q_ll", Kernel()) + cp = CpLite.__new__(CpLite) + cp.hq, cp.device, cp.me, cp.llq_off, cp.cp = 4, torch.device("cpu"), 0, 0, 1 + cp.peers = torch.zeros(1, dtype=torch.int64) + cp.q_all = torch.zeros(14, 4, 256, dtype=torch.uint8) + cp.cnt = torch.zeros(2, 14, dtype=torch.int32) + cp.pre_lens = cp.stock_lens = torch.zeros(2, dtype=torch.int32) + cp.gather_pdl, cp.early_merge = True, True + cp._gate_epilogue, cp._gate_epilogue_done, cp._logged = None, False, set() + calls = [] + cp._exact_partial = lambda *args, **kwargs: calls.append( + ("partial", kwargs["early_trigger"]) + ) + cp._merge_h = lambda *args, **kwargs: calls.append( + ("merge", kwargs["early"], kwargs["plain"]) + ) + q = torch.zeros(14, 4, 256, dtype=torch.uint8) + cache = torch.zeros(2, 1, 64, 256, dtype=torch.uint8) + table = torch.zeros(2, 2, dtype=torch.int32) + lengths = torch.full((2,), 71, dtype=torch.int32) + gate = {"nb": 2, "thr": 0, "GATE": False, "NB": 2, "GMAX": False} + + for override, expected in ((False, False), (None, True)): + calls.clear() + cp._ungated_main( + q, + cache, + cache, + table, + lengths, + 7, + 2, + 14, + 1.0, + 1.0, + gate, + early_merge=override, + ) + assert calls == [ + ("partial", expected), + ("merge", expected, override is False), + ] + + +def test_cp_fold_call_override_selects_the_r11_plain_merge(monkeypatch): + from sgl_nvidia_mega.qwen3_5.sm103.attention import q35_cp_lite as cp_module + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + calls = [] + + class Kernel: + def __init__(self, name): + self.name = name + + def __getitem__(self, grid): + return lambda *args, **kwargs: calls.append((self.name, grid, kwargs)) + + monkeypatch.setattr(cp_module, "_cp_merge_h_ll", Kernel("feature-rich")) + monkeypatch.setattr(cp_module, "_cp_merge_h_ll_plain", Kernel("r11-plain")) + cp = CpLite.__new__(CpLite) + cp.cp, cp.hq, cp.me, cp.llh_off = 4, 4, 0, 0 + cp.po = torch.zeros(4, 7, 4, 256, dtype=torch.bfloat16) + cp.plse = torch.zeros(4, 7, 4, dtype=torch.float32) + cp.peers = torch.zeros(4, dtype=torch.int64) + cp.cnt_h = torch.zeros(7 * 4, dtype=torch.int32) + q = torch.zeros(7, 4, 256, dtype=torch.uint8) + cache = torch.zeros(2, 1, 64, 256, dtype=torch.uint8) + table = torch.zeros(1, 2, dtype=torch.int32) + lengths = torch.full((1,), 71, dtype=torch.int32) + out = torch.zeros(7, 4, 256, dtype=torch.bfloat16) + gate = {"nb": 1, "thr": 0, "NB": 1, "GMAX": False} + + for plain, expected in ((False, "feature-rich"), (True, "r11-plain")): + cp._fold_merge_h( + q, + cache, + cache, + table, + lengths, + 7, + 7, + 1.0, + 1.0, + out, + gate, + plain=plain, + ) + assert calls[-1][0] == expected + assert calls[-1][1] == (7, 4) + assert calls[-1][2]["WAIT"] is False + assert "EARLY" not in calls[-1][2] + assert "GATE_EPILOGUE" not in calls[-1][2] + + vstart = torch.zeros(1, dtype=torch.int32) + vcnt = torch.ones(1, dtype=torch.int32) + fused = ( + (None, None, None, vstart, vcnt), + None, + torch.zeros(1, 7 * 4 * 4, 256, dtype=torch.bfloat16), + torch.zeros(1, 7 * 4 * 4, dtype=torch.float32), + ) + cp._fold_merge_h( + q, + cache, + cache, + table, + lengths, + 7, + 7, + 1.0, + 1.0, + out, + gate, + fused=fused, + plain=True, + ) + assert calls[-1][0] == "r11-plain" + assert calls[-1][2]["CHUNK"] is True + assert "EARLY" not in calls[-1][2] + assert "GATE_EPILOGUE" not in calls[-1][2] + + +def test_cp_exact_plain_merge_does_not_pass_feature_constexprs(monkeypatch): + from sgl_nvidia_mega.qwen3_5.sm103.attention import q35_cp_lite as cp_module + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite + + calls = [] + + class Kernel: + def __init__(self, name): + self.name = name + + def __getitem__(self, grid): + return lambda *args, **kwargs: calls.append((self.name, grid, kwargs)) + + monkeypatch.setattr(cp_module, "_cp_merge_h_ll", Kernel("feature-rich")) + monkeypatch.setattr(cp_module, "_cp_merge_h_ll_plain", Kernel("r11-plain")) + cp = CpLite.__new__(CpLite) + cp.cp, cp.hq, cp.me, cp.llh_off = 4, 4, 0, 0 + cp.po_x = torch.zeros(4, 7, 4, 256, dtype=torch.bfloat16) + cp.pl_x = torch.zeros(4, 7, 4, dtype=torch.float32) + cp.peers = torch.zeros(4, dtype=torch.int64) + cp.cnt_h = torch.zeros(7 * 4, dtype=torch.int32) + q = torch.zeros(7, 4, 256, dtype=torch.uint8) + cache = torch.zeros(2, 1, 64, 256, dtype=torch.uint8) + table = torch.zeros(1, 2, dtype=torch.int32) + lengths = torch.full((1,), 71, dtype=torch.int32) + out = torch.zeros(7, 4, 256, dtype=torch.bfloat16) + gate = {"nb": 1, "thr": 0, "NB": 1, "GMAX": False} + + cp._merge_h( + q, + cache, + cache, + table, + lengths, + 7, + 7, + 1.0, + 1.0, + out, + gate, + pdl=True, + early=False, + plain=True, + ) + name, grid, kwargs = calls[-1] + assert (name, grid) == ("r11-plain", (7, 4)) + assert kwargs["WAIT"] is True + assert "EARLY" not in kwargs + assert "GATE_EPILOGUE" not in kwargs + + cp._merge_h( + q, + cache, + cache, + table, + lengths, + 7, + 7, + 1.0, + 1.0, + out, + gate, + pdl=True, + early=True, + plain=False, + ) + name, grid, kwargs = calls[-1] + assert (name, grid) == ("feature-rich", (7, 4)) + assert kwargs["EARLY"] is True + + def test_cp_gate_epilogue_offer_is_consumed_once(): from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite @@ -732,9 +933,21 @@ def test_the_draft_decode_keeps_cp_lite_for_cp_variants_and_multi_request( class CpLite: def attend( - self, q, kv, pages, lens, q_len, bmm1, bmm2, ws, layer, owner, min_bs=None + self, + q, + kv, + pages, + lens, + q_len, + bmm1, + bmm2, + ws, + layer, + owner, + min_bs=None, + early_merge=None, ): - calls.append((tuple(q.shape), q_len, bmm1, bmm2, min_bs)) + calls.append((tuple(q.shape), q_len, bmm1, bmm2, min_bs, early_merge)) return torch.ones(q.shape[0], heads, dim, dtype=torch.bfloat16) monkeypatch.setattr(split_kv, "cp_lite", CpLite) @@ -772,7 +985,16 @@ def attend( q8 = torch.zeros(requests * rows, heads * dim, dtype=torch.float8_e4m3fn) out = attn.decode(q8, meta, kv, rows_per_request=rows) assert out.shape == (requests * rows, heads * dim) and bool((out == 1).all()) - assert calls == [((requests * rows, heads, dim), rows, 1.0, 3.0, min_bs)] + assert calls == [ + ( + (requests * rows, heads, dim), + rows, + 1.0, + 3.0, + min_bs, + False if requests > 1 else None, + ) + ] @pytest.mark.parametrize( From 6b1f5df61b451e0f071fd113bbe843fd8cda00a7 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 2 Nov 2026 15:35:32 -0800 Subject: [PATCH 44/62] style(qwen3.5): apply the repository's isort and ruff formatting Format the two files the pre-commit lint job flags on the current head; no code change. --- sgl_nvidia_mega/qwen3_5/plan.py | 4 +++- test/kernels/test_qwen3_5_attention.py | 22 +++++++--------------- 2 files changed, 10 insertions(+), 16 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index 1ca1c3c0..d9c1232b 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -61,7 +61,9 @@ # The router front's 16-row tiles replace SGLang's router GEMM through two # waves; wider route_topk rows keep the stock GEMM. ROUTE_GEMM_MAX_ROWS = 32 -GDN_CORE_MAX_ROWS = 896 # and gdn_core's launch (TRT-L4: checked at 64 and 128 requests) +GDN_CORE_MAX_ROWS = ( + 896 # and gdn_core's launch (TRT-L4: checked at 64 and 128 requests) +) MKM_MAX_ROWS = 7 MKM_WIDE_MAX_ROWS = 16 MKM_M32_MAX_ROWS = 32 diff --git a/test/kernels/test_qwen3_5_attention.py b/test/kernels/test_qwen3_5_attention.py index c54768c9..6ad78fd8 100644 --- a/test/kernels/test_qwen3_5_attention.py +++ b/test/kernels/test_qwen3_5_attention.py @@ -193,6 +193,8 @@ def test_cp_r11_plain_merge_matches_r11_at_b32(chunk): from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import ( _cp_merge_h_ll as feature_merge, + ) + from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import ( _cp_merge_h_ll_plain as plain_merge, ) @@ -238,12 +240,8 @@ def test_cp_r11_plain_merge_matches_r11_at_b32(chunk): device="cuda", generator=g, ) - vstart = ( - torch.arange(32, dtype=torch.int32, device="cuda") * chunks_per_request - ) - vcnt = torch.full( - (32,), chunks_per_request, dtype=torch.int32, device="cuda" - ) + vstart = torch.arange(32, dtype=torch.int32, device="cuda") * chunks_per_request + vcnt = torch.full((32,), chunks_per_request, dtype=torch.int32, device="cuda") k, v = ( (torch.randn(32, 1, page, dim, device="cuda", generator=g) * 2).to( torch.float8_e4m3fn @@ -329,16 +327,12 @@ def run(kernel, early=None): causal_sum = weight.sum(-1) causal_out = (weight.unsqueeze(-1) * vf[:, None, None]).sum(-2) * 0.5 if chunk: - chunk_lse = pl_v.view( - requests, chunks_per_request, q_len, heads - ).float() + chunk_lse = pl_v.view(requests, chunks_per_request, q_len, heads).float() chunk_peak = chunk_lse.max(1).values chunk_weight = torch.exp2(chunk_lse - chunk_peak.unsqueeze(1)) chunk_sum = chunk_weight.sum(1) prefix_out = ( - po_v.float().view( - requests, chunks_per_request, q_len, heads, dim - ) + po_v.float().view(requests, chunks_per_request, q_len, heads, dim) * chunk_weight.unsqueeze(-1) ).sum(1) / chunk_sum.unsqueeze(-1) prefix_lse = (chunk_peak + torch.log2(chunk_sum)) * math.log(2.0) @@ -379,9 +373,7 @@ def pair_stats(got, baseline): return { "max_abs": error.max().item(), "max_rel": relative.max().item(), - "bf16_mismatch": ( - got.view(torch.int16) != baseline.view(torch.int16) - ) + "bf16_mismatch": (got.view(torch.int16) != baseline.view(torch.int16)) .count_nonzero() .item(), } From ec40427958d32be3022c3d00dd107c5f4ace2f57 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Sun, 1 Nov 2026 21:38:26 -0800 Subject: [PATCH 45/62] perf(qwen3.5): extend GDN verify layouts and add prefill split conv --- sgl_nvidia_mega/qwen3_5/gdn/__init__.py | 46 +++- sgl_nvidia_mega/qwen3_5/gdn/front.py | 3 +- sgl_nvidia_mega/qwen3_5/gdn/prefill.py | 191 +++++++++++++ sgl_nvidia_mega/qwen3_5/gdn/verify.py | 16 +- sgl_nvidia_mega/qwen3_5/runtime/prefill.py | 297 +++++++++++++++++++++ test/kernels/test_qwen3_5_gdn.py | 206 ++++++++++++-- test/unit/test_qwen3_5_runtime.py | 51 ++++ test/unit/test_qwen3_5_spec.py | 8 - test/unit/test_qwen3_5_steps.py | 93 ++++++- 9 files changed, 857 insertions(+), 54 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/gdn/prefill.py create mode 100644 sgl_nvidia_mega/qwen3_5/runtime/prefill.py diff --git a/sgl_nvidia_mega/qwen3_5/gdn/__init__.py b/sgl_nvidia_mega/qwen3_5/gdn/__init__.py index cef99e4f..5413f1c5 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/__init__.py @@ -99,6 +99,7 @@ def build(model, ctx) -> Bound: reason = verify.unsupported(gdn) if reason: return Bound(reason=f"layer {i}: {reason}") + reasons = [] packed = {} if "gdn_push" in SLOTS: import torch @@ -115,12 +116,18 @@ def build(model, ctx) -> Bound: try: packed[gdn.attn.layer_id] = pack_out_proj(gdn.out_proj.weight.t()) except ValueError as e: - return Bound(reason=f"layer {i}: {e}") + reasons.append(f"gdn_push: layer {i}: {e}") + packed = {} + break rlin = {} - if packed and "route_lin" in SLOTS: - rlin, reason = _route_lin(m, gdns) - if reason: - return Bound(reason=f"slot route_lin: {reason}") + if "route_lin" in SLOTS: + if packed: + rlin, reason = _route_lin(m, gdns) + if reason: + reasons.append(f"route_lin: {reason}") + rlin = {} + else: + reasons.append("route_lin: requires supported gdn_push") front_packed, error = {}, None if "gdn_front" in SLOTS: import torch @@ -128,13 +135,16 @@ def build(model, ctx) -> Bound: from sgl_nvidia_mega.qwen3_5.gdn import front if ctx.tokens_per_request != front.ROWS: - return Bound(reason=f"gdn_front: {ctx.tokens_per_request} rows a request") - for i, gdn in gdns: - reason = front.unsupported(gdn) - if reason: - return Bound(reason=f"gdn_front: layer {i}: {reason}") - front_packed = front.build(model, gdns) - error = torch.zeros(1, dtype=torch.int32, device="cuda") + reasons.append(f"gdn_front: {ctx.tokens_per_request} rows a request") + else: + for i, gdn in gdns: + reason = front.unsupported(gdn) + if reason: + reasons.append(f"gdn_front: layer {i}: {reason}") + break + else: + front_packed = front.build(model, gdns) + error = torch.zeros(1, dtype=torch.int32, device="cuda") core = "gdn_core" in SLOTS if core: import torch @@ -149,6 +159,9 @@ def build(model, ctx) -> Bound: attn.num_v_heads, gdns[0][1].norm.weight.dtype == torch.float32, ) + trt_front = "gdn_trt_front" in SLOTS and bool(front_packed) and core + if "gdn_trt_front" in SLOTS and not trt_front: + reasons.append("gdn_trt_front: requires supported gdn_front and gdn_core") return Bound( gdn_verify=True, gdn_pdl="gdn_pdl" in SLOTS, @@ -156,9 +169,10 @@ def build(model, ctx) -> Bound: gdn_conv="gdn_conv" in SLOTS, packed=packed, front=front_packed, - gdn_trt_front="gdn_trt_front" in SLOTS, + gdn_trt_front=trt_front, rlin=rlin, error=error, + reason="; ".join(reasons), ) @@ -217,6 +231,12 @@ def step(layer, x, st): z_shape = z.shape core = core.reshape(-1, core.shape[-1]) z = z.reshape(-1, z.shape[-1]) + # DP attention may return only the real rows while the projection and z + # retain the graph's padded rows. Match SGLang before the gated norm. + if core.shape != z.shape: + padded = z.new_zeros(z.shape) + padded[: core.shape[0], :] = core + core = padded core = gdn.norm(core, z) core = core.reshape(z_shape) core = core.reshape(*core.shape[:-2], core.shape[-2] * core.shape[-1]) diff --git a/sgl_nvidia_mega/qwen3_5/gdn/front.py b/sgl_nvidia_mega/qwen3_5/gdn/front.py index 99f57a9d..82f6744f 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/front.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/front.py @@ -4,8 +4,7 @@ # of SGLang branch nvidia/qwen35-agentx-gb300 and SGLang's # python/sglang/srt/layers/attention/linear/gdn_backend.py (GDNAttnBackend.forward_extend). """Slot ``gdn_front`` (numeric): slot gdn_verify's step for one request of 7 rows with -the fused GDN front of SGLang branch nvidia/qwen35-agentx-gb300 (SGLANG_Q35MF_MKG_FUSED + -MKG_BA_TC + MKM_GDN_FRONT): +the fused GDN front of SGLang branch nvidia/qwen35-agentx-gb300: 1. ``sm103/gdn/front.py`` in one launch: the static-FP8 ``in_proj_qkvz`` projection of the BF16 rows, the conv update of its q / k / v channels, the BF16 ``in_proj_ba`` diff --git a/sgl_nvidia_mega/qwen3_5/gdn/prefill.py b/sgl_nvidia_mega/qwen3_5/gdn/prefill.py new file mode 100644 index 00000000..379559c8 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/gdn/prefill.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Adapted from SGLang (Apache-2.0): GDNAttnBackend.forward_extend and +# MambaAttnBackendBase._track_mamba_state_extend. +"""Slot gdn_prefill_split_conv: channel-sliced conv into contiguous Q/K/V. + +The native Triton recurrence, state writes, gated norm and output projection +retain SGLang's order. Only the conv output layout changes, above the fused +Q/K/V split limit. The caller checks the view before any layer mutates state. +""" + +from __future__ import annotations + +SGLANG_READS = { + "kernels/ops/mamba/causal_conv1d_triton.py": ("causal_conv1d_fn",), + "kernels/ops/attention/triton_gdn_fused_proj.py": ( + "qwen3_5_gdn_prefill_projection_views", + ), + "kernels/ops/attention/linear/fused_gdn_gating.py": ("fused_gdn_gating",), + "kernels/ops/attention/fla/chunk.py": ("chunk_gated_delta_rule",), + "srt/models/qwen3_5.py": ( + "_forward_input_proj", + "num_k_heads", + "num_v_heads", + "attn_tp_size", + "head_k_dim", + "head_v_dim", + "attn", + "norm", + "out_proj", + ), + "srt/layers/radix_linear_attention.py": ( + "layer_id", + "q_dim", + "k_dim", + "v_dim", + "num_q_heads", + "num_k_heads", + "num_v_heads", + "head_q_dim", + "head_k_dim", + "head_v_dim", + "conv_weights", + "bias", + "activation", + "A_log", + "dt_bias", + ), + "srt/model_executor/forward_batch_info.py": ( + "extend_prefix_lens", + "extend_seq_lens_cpu", + ), + "srt/kernel_plugins/nvidia/eagle_binding.py": ( + "LinearPrefillView", + "LinearPrefillBuffers", + "state", + "layer_ids", + "conv", + "temporal", + "query_start_loc", + "state_indices", + "conv_block_table", + "kernel_id", + "is_graph", + "has_mis", + "has_track", + "track_conv_indices", + "track_conv_dst", + "track_ssm_h_src", + "track_ssm_h_dst", + "track_ssm_final_src", + "track_ssm_final_dst", + "track_ssm_recompute_dst", + ), +} + +SGLANG_MIRRORS = { + # v3 routes committed state via STATE_PART. LinearPrefillView already carries + # those exact committed conv/temporal buffers; no recurrence/order change. + "srt/layers/attention/linear/gdn_backend.py:GDNAttnBackend.forward_extend": ( + "12a835fc4f78af33", + "6ed4eb6a818f7e20", + ), + "srt/layers/attention/hybrid_linear_attn_backend.py:MambaAttnBackendBase._track_mamba_state_extend": "84c447ca707eba64", +} + +FUSED_SPLIT_LIMIT = 8192 + + +def unsupported(view) -> str: + """The first implementation follows the eager native Triton prefill chain.""" + if view.kernel_id != 1: + return "not the native Triton prefill kernel" + if view.is_graph or view.has_mis: + return "prefill graph or multi-item layout" + if not view.state.temporal.is_contiguous() or any( + not x.is_contiguous() for x in view.state.conv + ): + return "strided committed state pool" + if ( + view.track_ssm_recompute_dst is not None + and view.track_ssm_recompute_dst.numel() + ): + return "recomputed checkpoint layout" + return "" + + +def split_conv(attn, mixed_qkv, view, forward_batch, conv_states): + """Three native conv launches, writing dense [1, T, heads, dim] tensors.""" + from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_fn + + extra = ( + {"block_table": view.conv_block_table} + if view.conv_block_table is not None + else {} + ) + has_initial_state = forward_batch.extend_prefix_lens > 0 + tokens = mixed_qkv.shape[0] + outputs = [] + start = 0 + for heads, dim in ( + (attn.num_q_heads, attn.head_q_dim), + (attn.num_k_heads, attn.head_k_dim), + (attn.num_v_heads, attn.head_v_dim), + ): + channels = slice(start, start + heads * dim) + start = channels.stop + out = causal_conv1d_fn( + mixed_qkv[:, channels].transpose(0, 1), + attn.conv_weights[channels], + None if attn.bias is None else attn.bias[channels], + activation=attn.activation, + conv_states=conv_states[:, channels], + has_initial_state=has_initial_state, + cache_indices=view.state_indices, + query_start_loc=view.query_start_loc, + seq_lens_cpu=forward_batch.extend_seq_lens_cpu, + **extra, + ) + outputs.append(out.transpose(0, 1)[:tokens].view(1, tokens, heads, dim)) + return outputs + + +def step(gdn, hidden, view, forward_batch): + from sglang.kernels.ops.attention.fla.chunk import chunk_gated_delta_rule + from sglang.kernels.ops.attention.linear.fused_gdn_gating import fused_gdn_gating + from sglang.kernels.ops.attention.triton_gdn_fused_proj import ( + qwen3_5_gdn_prefill_projection_views, + ) + + attn = gdn.attn + qkvz, ba = gdn._forward_input_proj(hidden) + mixed, z, b, a = qwen3_5_gdn_prefill_projection_views( + qkvz, + ba, + gdn.num_k_heads // gdn.attn_tp_size, + gdn.num_v_heads // gdn.attn_tp_size, + gdn.head_k_dim, + gdn.head_v_dim, + ) + li = view.state.layer_ids.index(attn.layer_id) + conv_states = view.state.conv[0][li] + ssm = view.state.temporal[li] + if view.has_track: + conv_states[view.track_conv_dst] = mixed[view.track_conv_indices].transpose( + 1, 2 + ) + query, key, value = split_conv(attn, mixed, view, forward_batch, conv_states) + g, beta = fused_gdn_gating(attn.A_log, a, b, attn.dt_bias) + core, _, h = chunk_gated_delta_rule( + q=query, + k=key, + v=value, + g=g, + beta=beta, + initial_state=ssm, + initial_state_indices=view.state_indices, + cu_seqlens=view.query_start_loc, + head_first=False, + use_qk_l2norm_in_kernel=True, + inplace_update=True, + ) + if view.has_track: + if view.track_ssm_h_src.numel(): + ssm[view.track_ssm_h_dst] = h.squeeze(0)[view.track_ssm_h_src].to( + ssm.dtype, copy=False + ) + if view.track_ssm_final_src.numel(): + ssm[view.track_ssm_final_dst] = ssm[view.track_ssm_final_src] + core = gdn.norm(core.reshape(-1, core.shape[-1]), z) + return gdn.out_proj(core.reshape(qkvz.shape[0], -1))[0] diff --git a/sgl_nvidia_mega/qwen3_5/gdn/verify.py b/sgl_nvidia_mega/qwen3_5/gdn/verify.py index 57cd453b..44ae3d31 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/verify.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/verify.py @@ -11,15 +11,17 @@ target-verify arm for the circular ReplaySSM ring (the causal conv update, the strided q / k / v split, the ReplaySSM verify), reading the step's GDN metadata and the mamba pools from the forward's view (``Step.view.gdn`` and ``.mamba``) and its request slots -from the forward batch, then the gated norm and ``out_proj``. The choices, all bitwise -to SGLang's own: +from the forward batch, then the gated norm and ``out_proj``. This explicitly +requested numeric slot changes these choices: - token-major projection views, with a token-major conv output (no split copy); - a 128-channel conv tile; - the ReplaySSM verify with 32-wide value tiles on 2 warps, on SGLang's kernel (``replayssm.py``), and with slot ``gdn_pdl`` launched dependent on the conv; - the gated norm writing the static FP8 ``out_proj`` takes (``norm_fp8.py``), which - ``out_proj``'s FlashInfer BMM then reads without quantizing it again. + ``out_proj``'s FlashInfer BMM then reads without quantizing it again. Its output + layout can change the FP32 reduction and BF16 midpoint rounding; the other + choices above retain native arithmetic. Slot ``gdn_push`` (numeric, not bitwise) runs ``out_proj`` on the decode hop's GEMM instead, whose epilogue pushes the partial into the hop's mailbox; the MoE input's @@ -161,7 +163,7 @@ ), } -RING_LEN, MAX_DRAFT = 16, 8 # the verify tiles' bounds, as measured +RING_LENGTHS, MAX_DRAFT = (16, 32), 8 # aggregate and DEP8 ReplaySSM pools # the step widths that have run the slot (each logged once, as engagement evidence) _ENGAGED: set = set() @@ -234,7 +236,7 @@ def pool_reasons(pools, tokens_per_request: int) -> list[str]: reasons = [] if pools.ring_d is None or pools.ring_cache_base is None: reasons.append("no circular ReplaySSM verify ring") - elif pools.ring_d.shape[-2] != RING_LEN: + elif pools.ring_d.shape[-2] not in RING_LENGTHS: reasons.append(f"ring length {pools.ring_d.shape[-2]}") if pools.temporal.dtype != torch.bfloat16: reasons.append(f"{pools.temporal.dtype} checkpoints") @@ -263,6 +265,10 @@ def core_pool_reasons(pools) -> list[str]: compact ring (update / key low parts, no exact-replay buffers) of BF16 checkpoints, int32 cursors and int8 flush flags.""" reasons = [] + # The fused CUDA core has two eight-entry history tiles. The generic + # gdn_verify path handles DEP8's 32-entry ring instead. + if pools.ring_d.shape[-2] != 16: + reasons.append(f"core ring length {pools.ring_d.shape[-2]} (serves 16)") if ( pools.ring_beta is not None or pools.ring_rawk is None diff --git a/sgl_nvidia_mega/qwen3_5/runtime/prefill.py b/sgl_nvidia_mega/qwen3_5/runtime/prefill.py new file mode 100644 index 00000000..9fb74ae2 --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/runtime/prefill.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Adapted from SGLang (Apache-2.0): Qwen3VLForConditionalGeneration.forward, +# Qwen3_5ForCausalLM.forward and Qwen3_5LinearDecoderLayer.forward. +"""The text prefill of one PP stage; native layers except the GDN conv layout.""" + +from __future__ import annotations + +import logging + +import torch +from sglang.kernels.plugin_kernels.nvidia.spec import Step + +from sgl_nvidia_mega.qwen3_5.gdn import prefill as gdn_prefill +from sgl_nvidia_mega.qwen3_5.runtime import parity + +logger = logging.getLogger(__name__) + +SGLANG_READS = { + "kernels/plugin_kernels/nvidia/spec.py": ("Step",), + "srt/models/qwen3_5.py": ( + "Qwen3_5LinearDecoderLayer", + "residual_batch", + "layers", + "start_layer", + "end_layer", + "pp_group", + "embed_tokens", + "attn_boundary", + "ffn_boundary", + "linear_attn", + "mlp", + "norm", + "flashinfer_mnnvl_cutedsl_fusion", + '"_is_layer_to_capture"', + ), + "srt/models/qwen3_vl.py": ( + "language_model_only", + "is_mrope_enabled", + "capture_aux_hidden_states", + "logits_processor", + "lm_head", + ), + "srt/layers/layer_boundary/residual/batch.py": ("start", "to_pp", "final_norm"), + "srt/layers/layer_boundary/stage.py": ("from_pp", "prepare", "finish", "exit"), + "srt/layers/layer_boundary/exit.py": ("defer_moe_finalize", "finish"), + "srt/layers/aux_hidden_states.py": ("AuxHiddenStateList",), + "srt/eplb/expert_distribution.py": ( + "get_global_expert_distribution_recorder", + "with_current_layer", + ), + "srt/model_executor/forward_batch_info.py": ( + "mrope_positions", + "mm_inputs", + "tensors", + ), + "srt/model_executor/forward_mode.py|srt/model_executor/forward_batch_info.py": ( + "is_extend_without_speculative", + ), + "srt/layers/logits_processor.py": ( + "next_token_logits", + "hidden_states", + "last_hidden_states", + ), +} + +SGLANG_MIRRORS = { + "srt/models/qwen3_vl.py:Qwen3VLForConditionalGeneration.forward": ( + "d394d79ba7fa7b31", + "dbdcf4b88a2c0d40", + ), + "srt/models/qwen3_5.py:Qwen3_5ForCausalLM.forward": "9e530b264ce30ad6", + "srt/models/qwen3_5.py:Qwen3_5LinearDecoderLayer.forward": "2298a57111f9430e", +} + + +class Qwen35Prefill: + draft_sampler = draft_loop = draft_extend = begin_draft_round = None + graph_variants = None + + def __init__(self, model, ctx, enabled): + from sglang.srt.models.qwen3_5 import Qwen3_5LinearDecoderLayer + + self.lm = model + self.enabled = enabled + self.parity = PrefillParity() if parity.ENABLED else None + # Widths describe decode-graph installation; P uses the ordinary eager + # forward wrapper and plans its actual chunk rows independently. + self.widths = (ctx.tokens_per_request,) + self._linear = { + i + for i in range(model.model.start_layer, model.model.end_layer) + if isinstance(model.model.layers[i], Qwen3_5LinearDecoderLayer) + } + self._engaged = False + for i in self._linear: + gdn = model.model.layers[i].linear_attn + if ( + gdn.attn.q_dim + gdn.attn.k_dim + gdn.attn.v_dim + <= gdn_prefill.FUSED_SPLIT_LIMIT + or gdn.attn_tp_size != 1 + ): + self.enabled = False + + def summary(self): + return "prefill TP1/PP4; gdn_prefill_split_conv " + ( + "built" if self.enabled else "stock" + ) + + def plan( + self, + input_ids, + forward_batch, + get_embedding=False, + pp_proxy_tensors=None, + **kwargs, + ): + if ( + (not self.enabled and self.parity is None) + or kwargs + or get_embedding + or any(item is not None for item in (forward_batch.mm_inputs or ())) + or not forward_batch.forward_mode.is_extend_without_speculative() + ): + return None + return Step(mode="prefill", rows=input_ids.shape[0]) + + def forward( + self, + step, + original, + input_ids, + positions, + forward_batch, + get_embedding=False, + pp_proxy_tensors=None, + **kwargs, + ): + if gdn_prefill.unsupported(step.view): + return original( + input_ids, + positions, + forward_batch, + get_embedding=get_embedding, + pp_proxy_tensors=pp_proxy_tensors, + **kwargs, + ) + if self.parity is not None: + stock = lambda: original( + input_ids, + positions, + forward_batch, + get_embedding=get_embedding, + pp_proxy_tensors=pp_proxy_tensors, + **kwargs, + ) + candidate = ( + ( + lambda: self._forward( + step.view, input_ids, positions, forward_batch, pp_proxy_tensors + ) + ) + if self.enabled + else stock + ) + return self.parity.run(step.rows, step.view, stock, candidate) + if not self._engaged: + logger.info( + "NVIDIA megakernel: Qwen3.5 slot gdn_prefill_split_conv runs at %d rows", + step.rows, + ) + self._engaged = True + return self._forward( + step.view, input_ids, positions, forward_batch, pp_proxy_tensors + ) + + def _forward(self, view, input_ids, positions, fb, pp_proxy_tensors): + from sglang.srt.eplb.expert_distribution import ( + get_global_expert_distribution_recorder, + ) + from sglang.srt.layers.aux_hidden_states import AuxHiddenStateList + from sglang.srt.layers.layer_boundary.residual import batch as residual_batch + + lm, model = self.lm, self.lm.model + if lm.is_mrope_enabled: + positions = fb.mrope_positions + if model.pp_group.is_first_rank: + hidden = model.embed_tokens(input_ids) + residual_batch.start(fb) + else: + hidden = model.layers[model.start_layer].attn_boundary.from_pp( + pp_proxy_tensors, fb + ) + aux = AuxHiddenStateList() + for i in range(model.start_layer, model.end_layer): + layer = model.layers[i] + capture = aux if getattr(layer, "_is_layer_to_capture", False) else None + with get_global_expert_distribution_recorder().with_current_layer(i): + if i not in self._linear: + hidden = layer( + positions=positions, + hidden_states=hidden, + forward_batch=fb, + captured_last_layer_outputs=capture, + ) + continue + x = layer.attn_boundary.prepare(hidden, fb, capture_gathered=capture) + part = gdn_prefill.step(layer.linear_attn, x, view, fb) + x = layer.attn_boundary.finish(part, fb) + x = layer.ffn_boundary.prepare(x, fb) + with layer.ffn_boundary.exit(fb) as ffn_exit: + out = layer.mlp(x, fb, defer_finalize=ffn_exit.defer_moe_finalize) + hidden = ffn_exit.finish(out) + if not model.pp_group.is_last_rank: + return residual_batch.to_pp(hidden, fb) + hidden = residual_batch.final_norm( + hidden, + fb, + model.norm, + finalize_norm=model.flashinfer_mnnvl_cutedsl_fusion, + skip_empty=True, + ) + return lm.logits_processor( + input_ids, + hidden, + lm.lm_head, + fb, + aux if lm.capture_aux_hidden_states and len(aux) else None, + ) + + +class PrefillParity: + """Test only: restore every state destination before the second forward. + + Unlike verify, prefill updates committed SSM state as well as conv history. + PP inputs stay fixed for the local comparison; each PP rank checks its output. + """ + + @staticmethod + def _outputs(out): + if isinstance(out, torch.Tensor): + return {"tensor": out.clone()} + tensors = getattr(out, "tensors", None) + if tensors is None: + tensors = { + name: getattr(out, name, None) + for name in ("next_token_logits", "hidden_states", "last_hidden_states") + } + return { + name: value.clone() + for name, value in tensors.items() + if isinstance(value, torch.Tensor) + } + + @staticmethod + def _indices(view, names): + rows = [view.state_indices.to(torch.long)] + if view.has_track: + rows += [ + value.to(torch.long) + for name in names + if (value := getattr(view, name)) is not None + ] + indices = torch.cat(rows) + return indices[indices >= 0].unique() + + def run(self, rows, view, original, candidate): + ci = self._indices(view, ("track_conv_dst",)) + si = self._indices( + view, ("track_ssm_h_dst", "track_ssm_final_dst", "track_ssm_recompute_dst") + ) + pools = [*((value, ci) for value in view.state.conv), (view.state.temporal, si)] + before = [value[:, indices].clone() for value, indices in pools] + want = self._outputs(original()) + state_want = [value[:, indices].clone() for value, indices in pools] + for (value, indices), saved in zip(pools, before): + value[:, indices] = saved + got = candidate() + outputs = self._outputs(got) + assert outputs.keys() == want.keys() + output_differ = sum(int(parity.differing(outputs[n], want[n])) for n in want) + state_differ = sum( + int(parity.differing(value[:, indices], saved)) + for (value, indices), saved in zip(pools, state_want) + ) + logger.info( + "NVIDIA megakernel prefill parity rows=%d output_differ=%d " + "state_differ=%d outputs_compared=%d state_compared=%d", + rows, + output_differ, + state_differ, + sum(value.numel() for value in want.values()), + sum(value.numel() for value in state_want), + ) + if output_differ or state_differ: + raise RuntimeError("GDN prefill output/state differs from stock") + return got diff --git a/test/kernels/test_qwen3_5_gdn.py b/test/kernels/test_qwen3_5_gdn.py index bd39fece..b0b66c6a 100644 --- a/test/kernels/test_qwen3_5_gdn.py +++ b/test/kernels/test_qwen3_5_gdn.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """The kernels slot gdn_verify launches its own way, on one SM100-family GPU (GB300), -against SGLang's, bit for bit, at Qwen3.5's TP8 GatedDeltaNet shapes, for steps of 1, 2 -and 4 requests: +against SGLang's, bit for bit, at Qwen3.5's TP8, TP4, TP2 and DP-attention +GatedDeltaNet shapes, including TP2 decode B10 and DEP8 per-rank B32: - the causal conv update with the 128-channel tile and the token-major output, against ``causal_conv1d_update`` (output, conv states and intermediate windows), on the @@ -55,32 +55,43 @@ def _gpu_ready() -> bool: pytestmark = pytest.mark.skipif(not _gpu_ready(), reason="needs an SM100-family GPU") -def _projection(bs, views, g): +def _projection(bs, views, g, h=H, hv=HV): """A verify step's mixed q / k / v rows [bs * 7, D]: a contiguous copy, or the token-major view of the q / k / v / z projection.""" - qkvz = torch.randn(bs * DRAFT, D + HV * V, device="cuda", generator=g).to( + dim = 2 * h * K + hv * V + qkvz = torch.randn(bs * DRAFT, dim + hv * V, device="cuda", generator=g).to( torch.bfloat16 ) - return qkvz[:, :D] if views else qkvz[:, :D].contiguous() + return qkvz[:, :dim] if views else qkvz[:, :dim].contiguous() # a covered step's requests (slot gdn_verify covers B1 to B128) and their slots REQUESTS = (1, 2, 4, 8, 16, 32, 64, 128) +# TP8 plus the TP4, TP2 and DEP8 (attention TP1 per DP rank) decode shapes +TOPOLOGY_BATCHES = [(8, bs) for bs in REQUESTS] + [(4, 1), (2, 10), (1, 32)] + + +def _heads(tp): + h, hv = 16 // tp, 64 // tp + return h, hv, 2 * h * K + hv * V + + STATE_SLOTS = (4, 7, 0, 8, *range(39, 11, -1), *range(135, 39, -1)) RING_SLOTS = (3, 5, 0, 8, *range(12, 136)) REPLAY_SLOTS = (1, 2, 4, 6, *range(10, 134)) -@pytest.mark.parametrize("bs", REQUESTS) +@pytest.mark.parametrize("tp, bs", TOPOLOGY_BATCHES) @pytest.mark.parametrize("views", [False, True]) @pytest.mark.parametrize("block_n", [256, 128]) -def test_conv_update_is_sglangs(views, block_n, bs): +def test_conv_update_is_sglangs(views, block_n, tp, bs): + H, HV, D = _heads(tp) from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_update from sgl_nvidia_mega.qwen3_5.gdn.conv import conv_update g = torch.Generator(device="cuda").manual_seed(block_n + views + 10 * bs) - rows = _projection(bs, views, g) + rows = _projection(bs, views, g, H, HV) weight = torch.randn(D, WIDTH, device="cuda", generator=g).to(torch.bfloat16) bias = torch.randn(D, device="cuda", generator=g).to(torch.bfloat16) state = torch.randn(SLOTS, D, WIDTH - 1, device="cuda", generator=g).to( @@ -121,18 +132,19 @@ def test_conv_update_is_sglangs(views, block_n, bs): assert got.transpose(1, 2).is_contiguous() -@pytest.mark.parametrize("bs", REQUESTS) +@pytest.mark.parametrize("tp, bs", TOPOLOGY_BATCHES) @pytest.mark.parametrize("tile", [(128, 4), (64, 2)]) @pytest.mark.parametrize("window_mode", [0, 1]) -def test_hoisted_conv_is_sglangs(window_mode, tile, bs): +def test_hoisted_conv_is_sglangs(window_mode, tile, tp, bs): """Slot gdn_conv: x rows loaded first (and with window_mode 1 one window row a channel) writes the output, the conv states and the deduplicated windows bit for bit as SGLang's kernel.""" + H, HV, D = _heads(tp) from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_update from sgl_nvidia_mega.qwen3_5.gdn.conv import conv_update g = torch.Generator(device="cuda").manual_seed(window_mode + tile[0] + 10 * bs) - rows = _projection(bs, True, g) + rows = _projection(bs, True, g, H, HV) weight = torch.randn(D, WIDTH, device="cuda", generator=g).to(torch.bfloat16) state = torch.randn(SLOTS, D, WIDTH - 1, device="cuda", generator=g).to( torch.bfloat16 @@ -183,9 +195,10 @@ def windows(t): assert torch.equal(state, want_state) and torch.equal(shared, want_shared) -@pytest.mark.parametrize("bs", REQUESTS) +@pytest.mark.parametrize("tp, bs", TOPOLOGY_BATCHES) @pytest.mark.parametrize("views", [False, True]) -def test_norm_fp8_is_sglangs_norm_then_static_quant(views, bs): +def test_norm_fp8_is_within_v2_numeric_bound(views, tp, bs): + _, HV, D = _heads(tp) from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 @@ -211,17 +224,25 @@ def test_norm_fp8_is_sglangs_norm_then_static_quant(views, bs): got = rms_norm_gated_fp8( core, norm.weight, norm.eps, z, scale, activation=norm.activation ) - assert torch.equal( - got.view(torch.uint8).reshape(tokens, -1), want.view(torch.uint8) + from sgl_nvidia_mega.qwen3_5.spec import NUMERIC_TOLERANCES + + ref = want.float().reshape(-1) + error = torch.linalg.vector_norm(got.float().reshape(-1) - ref) + relative = error / torch.linalg.vector_norm(ref) + assert ( + relative.item() + <= NUMERIC_TOLERANCES["gdn_verify"]["gdn_verify.norm_fp8.rel_l2_vs_stock"] ) -@pytest.mark.parametrize("bs", REQUESTS) +@pytest.mark.parametrize("tp, bs", TOPOLOGY_BATCHES) @pytest.mark.parametrize("pdl, gate", [(False, False), (True, False), (True, True)]) -def test_replayssm_verify_is_sglangs(pdl, gate, bs): +@pytest.mark.parametrize("ring_len", [16, 32]) +def test_replayssm_verify_is_sglangs(pdl, gate, ring_len, tp, bs): """The slot's verify launch (32-wide value tiles on 2 warps, and with or without the dependent launch, waiting before the gates with ``gate``) against gdn_replayssm_spec_decode's verify launch: the output and this step's ring entries.""" + H, HV, D = _heads(tp) try: from sglang.kernels.ops.attention.linear.gdn_replayssm_spec_decode import ( gdn_replayssm_spec_decode, @@ -236,7 +257,7 @@ def test_replayssm_verify_is_sglangs(pdl, gate, bs): g = torch.Generator(device="cuda").manual_seed( 3 + int(pdl) + 2 * int(gate) + 10 * bs ) - ring_len, slots = 16, SLOTS + slots = SLOTS tokens = bs * DRAFT def rand(*shape, dtype=torch.bfloat16): @@ -332,15 +353,16 @@ def rand(*shape, dtype=torch.bfloat16): # -------------------------------------------------------------------- slot gdn_core -@pytest.mark.parametrize("bs", REQUESTS) +@pytest.mark.parametrize("tp, bs", TOPOLOGY_BATCHES) @pytest.mark.parametrize("history", [False, True]) @pytest.mark.parametrize("pad", [False, True]) -def test_core_is_gdn_verify_numerically(pad, history, bs): +def test_core_is_gdn_verify_numerically(pad, history, tp, bs): """Slot gdn_core's one launch against slot gdn_verify's verify (32-wide value tiles, 2 warps) and FP8 gated norm on the same inputs: the replay output and the FP8 bytes up to rounding, this step's ring entries (update and key high + low parts, log-decay) up to FP32 summation order; a padded request writes zeros and leaves its ring alone. History: an empty ring (BF16 checkpoints fold every commit) or 1-9 committed steps.""" + H, HV, D = _heads(tp) from sgl_nvidia_mega.qwen3_5.gdn import replayssm from sgl_nvidia_mega.qwen3_5.gdn.norm_fp8 import rms_norm_gated_fp8 from sgl_nvidia_mega.qwen3_5.sm103.gdn.core import core @@ -1151,3 +1173,147 @@ def test_trt_front_replays_from_a_cuda_graph_and_never_compiles_in_one(): trt_lowm.prepare(unprepared) with pytest.raises(RuntimeError, match="not prepared"): trt_lowm.run(unprepared, out=outs["ba"], x16=o["x"], w_bf16=o["ba_w"]) + + +@pytest.mark.parametrize("on", [False, True]) +def test_prefill_split_conv_matches_native_state_chain(on): + """P TP1's conv, recurrence and tracked writes on the same native operands; + ``on`` adds the bias, the compact block table, a prefix and a padded request.""" + bias = compact = prefix = pad = on + from types import SimpleNamespace as NS + + from sglang.kernels.ops.attention.fla.chunk import chunk_gated_delta_rule + from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm + from sglang.kernels.ops.attention.linear.fused_gdn_gating import fused_gdn_gating + from sglang.kernels.ops.attention.triton_gdn_fused_proj import ( + qwen3_5_gdn_prefill_projection_views, + ) + from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_fn + + from sgl_nvidia_mega.qwen3_5.gdn import prefill + + g = torch.Generator(device="cuda").manual_seed(880 + bias + 10 * prefix) + h, hv, dim = _heads(1) + lengths = [129, 17] + tokens = sum(lengths) + + def rand(*shape): + return torch.randn(*shape, generator=g, device="cuda").bfloat16() + + projected, ba = rand(tokens, dim + hv * V), rand(tokens, 2 * hv) + mixed, z, b, a = qwen3_5_gdn_prefill_projection_views(projected, ba, h, hv, K, V) + attn = NS( + layer_id=12, + q_dim=h * K, + k_dim=h * K, + v_dim=hv * V, + num_q_heads=h, + num_k_heads=h, + num_v_heads=hv, + head_q_dim=K, + head_k_dim=K, + head_v_dim=V, + conv_weights=rand(dim, WIDTH), + bias=rand(dim) if bias else None, + activation="silu", + A_log=torch.rand(hv, device="cuda", generator=g), + dt_bias=rand(hv), + ) + norm = RMSNorm( + V, eps=1e-6, norm_before_gate=True, device="cuda", dtype=torch.bfloat16 + ) + norm.weight.data = rand(V) + module = NS( + attn=attn, + num_k_heads=h, + num_v_heads=hv, + attn_tp_size=1, + head_k_dim=K, + head_v_dim=V, + norm=norm, + _forward_input_proj=lambda _: (projected, ba), + out_proj=lambda x: (x, None), + ) + conv = rand(9, dim, WIDTH - 1) + ssm = rand(9, hv, V, K) * 0.02 + indices = torch.tensor([2, -1 if pad else 4], device="cuda", dtype=torch.int32) + qsl = torch.tensor([0, 129, tokens], device="cuda", dtype=torch.int32) + table = ( + torch.tensor( + [(i, c) for i, n in enumerate(lengths) for c in range((n + 7) // 8)], + device="cuda", + dtype=torch.int32, + ) + if compact + else None + ) + fb = NS( + extend_prefix_lens=torch.tensor([8192 if prefix else 0, 0], device="cuda"), + extend_seq_lens_cpu=lengths, + ) + if not prefix: + ssm[2].zero_() + ssm[4].zero_() + conv_ref, ssm_ref = conv.clone(), ssm.clone() + ix = lambda vals: torch.tensor(vals, device="cuda", dtype=torch.long) + view = NS( + state=NS(layer_ids=(12,), conv=(conv.unsqueeze(0),), temporal=ssm.unsqueeze(0)), + state_indices=indices, + query_start_loc=qsl, + conv_block_table=table, + kernel_id=1, + is_graph=False, + has_mis=False, + has_track=True, + track_conv_indices=ix([[126, 127, 128]]), + track_conv_dst=ix([6]), + track_ssm_h_src=ix([0]), + track_ssm_h_dst=ix([7]), + track_ssm_final_src=ix([2]), + track_ssm_final_dst=ix([8]), + track_ssm_recompute_dst=None, + ) + assert not prefill.unsupported(view) + # Native tracking runs before convolution overwrites its committed state. + conv_ref[view.track_conv_dst] = mixed[view.track_conv_indices].transpose(1, 2) + convolved = causal_conv1d_fn( + mixed.T, + attn.conv_weights, + attn.bias, + conv_states=conv_ref, + has_initial_state=fb.extend_prefix_lens > 0, + cache_indices=indices, + query_start_loc=qsl, + seq_lens_cpu=lengths, + **({"block_table": table} if compact else {}), + ).T + q, k, v = [ + x.view(1, tokens, heads, d) + for x, heads, d in zip( + convolved.split([h * K, h * K, hv * V], dim=-1), (h, h, hv), (K, K, V) + ) + ] + gate, beta = fused_gdn_gating(attn.A_log, a, b, attn.dt_bias) + want, _, history = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + initial_state=ssm_ref, + initial_state_indices=indices, + cu_seqlens=qsl, + head_first=False, + use_qk_l2norm_in_kernel=True, + inplace_update=True, + ) + ssm_ref[view.track_ssm_h_dst] = history.squeeze(0)[view.track_ssm_h_src].to( + ssm_ref.dtype + ) + ssm_ref[view.track_ssm_final_dst] = ssm_ref[view.track_ssm_final_src] + want = norm(want.reshape(-1, V), z).reshape(tokens, -1) + got = prefill.step(module, projected, view, fb) + live = slice(0, 129) if pad else slice(None) + assert torch.equal(got[live], want[live]) + assert torch.equal(conv, conv_ref) + assert torch.equal(ssm, ssm_ref) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index de36b77c..8648e9b1 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -1694,3 +1694,54 @@ def test_the_step_state_carries_the_binding(): "EagleBinding", "EagleForwardView", ) + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("mm_inputs", [None, [None]]) +def test_prefill_plan_uses_actual_chunk_rows_and_stock_when_slot_off( + enabled, mm_inputs +): + from sglang.srt.model_executor.forward_batch_info import ForwardMode + + from sgl_nvidia_mega.qwen3_5.runtime.prefill import Qwen35Prefill + + model = SimpleNamespace( + model=SimpleNamespace(start_layer=0, end_layer=0, layers=[]) + ) + built = Qwen35Prefill(model, SimpleNamespace(tokens_per_request=7), enabled) + fb = SimpleNamespace(mm_inputs=mm_inputs, forward_mode=ForwardMode.EXTEND) + step = built.plan(torch.zeros(33, dtype=torch.int64), fb) + if enabled: + assert step.mode == "prefill" and step.rows == 33 + else: + assert step is None + fb.forward_mode = ForwardMode.TARGET_VERIFY + assert built.plan(torch.zeros(7, dtype=torch.int64), fb) is None + + +def test_prefill_plan_keeps_real_multimodal_batches_on_stock(): + from sglang.srt.model_executor.forward_batch_info import ForwardMode + + from sgl_nvidia_mega.qwen3_5.runtime.prefill import Qwen35Prefill + + model = SimpleNamespace( + model=SimpleNamespace(start_layer=0, end_layer=0, layers=[]) + ) + built = Qwen35Prefill(model, SimpleNamespace(tokens_per_request=7), True) + fb = SimpleNamespace(mm_inputs=[None, object()], forward_mode=ForwardMode.EXTEND) + assert built.plan(torch.zeros(33, dtype=torch.int64), fb) is None + + +def test_prefill_unsupported_view_returns_the_original_before_state_writes(): + from sglang.kernels.plugin_kernels.nvidia.spec import Step + + from sgl_nvidia_mega.qwen3_5.runtime.prefill import Qwen35Prefill + + model = SimpleNamespace( + model=SimpleNamespace(start_layer=0, end_layer=0, layers=[]) + ) + built = Qwen35Prefill(model, SimpleNamespace(tokens_per_request=7), True) + marker = object() + original = lambda *a, **k: marker + step = Step(mode="prefill", rows=33, view=SimpleNamespace(kernel_id=0)) + assert built.forward(step, original, None, None, None) is marker diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index f61f9071..2a13bcb8 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -89,14 +89,6 @@ def test_the_decode_side_of_pd_is_accepted(): assert _refusal(disaggregation_mode="decode") == "" -def test_the_prefill_side_of_pd_is_refused(): - # SGLang refuses ReplaySSM on a PD prefill server itself - message = _refusal( - disaggregation_mode="prefill", enable_linear_replayssm_spec=False - ) - assert "a PD prefill server runs no covered step" in message, message - - BREAKABLE = with_phase( default_cuda_graph_config(), Phase.DECODE, backend=Backend.BREAKABLE ) diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index cda1e943..5c89fa87 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -962,13 +962,74 @@ def test_the_route_slot_is_not_built_for_a_layer_it_cannot_serve(monkeypatch): runtime.build(model, ctx) -def _gdn_module(log, x): +def _gdn_build_model(): + layer = qwen3_5.Qwen3_5LinearDecoderLayer.__new__(qwen3_5.Qwen3_5LinearDecoderLayer) + torch.nn.Module.__init__(layer) + layer.linear_attn = SimpleNamespace( + attn=SimpleNamespace(layer_id=0, num_k_heads=2, num_v_heads=8), + norm=SimpleNamespace(weight=torch.ones(128)), + out_proj=SimpleNamespace(weight=torch.ones(4, 4)), + ) + return SimpleNamespace( + model=SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) + ) + + +@pytest.mark.parametrize("tokens,reason", [(6, "6 rows a request"), (7, "TP 4")]) +def test_optional_gdn_front_falls_back_without_disabling_verify( + monkeypatch, tokens, reason +): + from sgl_nvidia_mega.qwen3_5.gdn import front, verify + from sgl_nvidia_mega.qwen3_5.sm103.gdn import core + + slots = { + "gdn_verify", + "gdn_pdl", + "gdn_conv", + "gdn_core", + "gdn_front", + "gdn_trt_front", + } + monkeypatch.setattr(runtime, "SLOTS", frozenset(slots)) + monkeypatch.setattr(verify, "unsupported", lambda _: "") + monkeypatch.setattr(front, "unsupported", lambda _: "TP 4") + monkeypatch.setattr(front, "build", lambda *a: pytest.fail("front built")) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(core, "build", lambda *a: None) + bound = gdn.build(_gdn_build_model(), SimpleNamespace(tokens_per_request=tokens)) + assert bound.gdn_verify and bound.gdn_pdl and bound.gdn_conv and bound.gdn_core + assert bound.front == {} and not bound.gdn_trt_front and reason in bound.reason + + +@pytest.mark.parametrize("bad_slot", ["gdn_push", "route_lin"]) +def test_optional_gdn_push_or_route_lin_falls_back_independently(monkeypatch, bad_slot): + from sgl_nvidia_mega.qwen3_5.gdn import verify + from sgl_nvidia_mega.qwen3_5.sm103 import hop + + slots = {"gdn_verify", "gdn_pdl", "gdn_push", "route_lin"} + monkeypatch.setattr(runtime, "SLOTS", frozenset(slots)) + monkeypatch.setattr(verify, "unsupported", lambda _: "") + + def pack(weight): + if bad_slot == "gdn_push": + raise ValueError("unsupported out_proj shape") + return "packed" + + monkeypatch.setattr(hop, "pack_out_proj", pack) + monkeypatch.setattr(gdn, "_route_lin", lambda *a: ({}, "TP 4")) + bound = gdn.build(_gdn_build_model(), SimpleNamespace(tokens_per_request=7)) + assert bound.gdn_verify and bound.gdn_pdl and bound.rlin == {} + assert bad_slot in bound.reason + assert bound.packed == ({} if bad_slot == "gdn_push" else {0: "packed"}) + + +def _gdn_module(log, x, core_rows=7): """A Qwen3_5GatedDeltaNet whose projections, backend call, norm and out_proj log.""" mod = qwen3_5.Qwen3_5GatedDeltaNet.__new__(qwen3_5.Qwen3_5GatedDeltaNet) torch.nn.Module.__init__(mod) qkvz, ba = torch.zeros(7, 10), torch.zeros(7, 2) z = torch.zeros(7, 2, 3) - core = torch.zeros(7, 2, 3) + core = torch.ones(core_rows, 2, 3) log.known.update(qkvz=qkvz, ba=ba, z=z, core=core) mod.__dict__.update( num_k_heads=16, @@ -986,12 +1047,21 @@ def _gdn_module(log, x): @pytest.mark.parametrize("packed_verify", [False, True]) -def test_the_gdn_step_issues_what_sglang_issues(packed_verify): +@pytest.mark.parametrize("core_rows", [3, 7]) +def test_the_gdn_step_issues_what_sglang_issues(packed_verify, core_rows): runs = [] for which in ("sglang", "plugin"): x = torch.zeros(7, 4) log = _log(x=x) - mod, split = _gdn_module(log, x) + mod, split = _gdn_module(log, x, core_rows=core_rows) + norm_inputs = [] + norm = mod.norm + + def checked_norm(core, z): + norm_inputs.append(core.clone()) + return norm(core, z) + + mod.norm = checked_norm backend = SimpleNamespace( linear_attn_backend=SimpleNamespace(supports_packed_verify=packed_verify) ) @@ -1027,13 +1097,14 @@ def test_the_gdn_step_issues_what_sglang_issues(packed_verify): out = gdn.step(SimpleNamespace(linear_attn=mod), x, st) except (NotImplementedError, RuntimeError, AttributeError) as e: out = type(e).__name__ - runs.append((log, out)) - (want_log, want), (got_log, got) = runs + runs.append((log, out, norm_inputs)) + (want_log, want, want_norm), (got_log, got, got_norm) = runs if packed_verify: # SGLang folds the split and conv into such a backend; the step refuses it assert got == "NotImplementedError" and want != got return assert got_log == want_log and want == got == "out" + assert torch.equal(got_norm[0], want_norm[0]) def _gdn_verify_setup(log, *, ring_len=16, temporal_dtype=torch.bfloat16, draft=7): @@ -2112,6 +2183,16 @@ def test_the_gdn_verify_slot_names_the_pools_it_cannot_serve(kw, want): built.forward(step, None, torch.zeros(7, dtype=torch.int64), None, None) +@pytest.mark.parametrize("ring_len", [16, 32]) +def test_gdn_verify_accepts_aggregate_and_dep8_rings(ring_len): + verify, _, st = _gdn_verify_setup(_log(), ring_len=ring_len) + assert verify.pool_reasons(st.view.mamba, st.binding.tokens_per_request) == [] + if ring_len == 32: + assert "core ring length 32 (serves 16)" in verify.core_pool_reasons( + st.view.mamba + ) + + def test_the_pools_are_checked_once_at_the_first_covered_step(monkeypatch): verify, _, st = _gdn_verify_setup(_log()) built = _verify(binding=st.binding) From 67cbe3095a57fae719534a7d659bca18a69a8833 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 2 Nov 2026 00:46:52 -0800 Subject: [PATCH 46/62] perf(qwen3.5): extend attention fronts to TP2 and wide TP4 --- sgl_nvidia_mega/qwen3_5/attention/__init__.py | 37 ++++++---- sgl_nvidia_mega/qwen3_5/attention/mka.py | 36 ++++++---- sgl_nvidia_mega/qwen3_5/attention/route.py | 2 +- sgl_nvidia_mega/qwen3_5/plan.py | 2 +- test/kernels/test_qwen3_5_attention.py | 35 ++++++---- test/unit/test_qwen3_5_runtime.py | 6 +- test/unit/test_qwen3_5_spec.py | 2 +- test/unit/test_qwen3_5_steps.py | 69 +++++++++++++++---- 8 files changed, 127 insertions(+), 62 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/attention/__init__.py b/sgl_nvidia_mega/qwen3_5/attention/__init__.py index 3c3a7c68..f4712bbf 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/attention/__init__.py @@ -55,7 +55,7 @@ class Bound(msgspec.Struct, frozen=True): skv2: bool = False # slot skv2 serves every full-attention layer mka: bool = False - mka_wide: bool = False # mka's wide entries serve the steps of several requests + mka_wide: bool = False # mka's front path also serves steps of several requests mka_push: bool = False route_lin: bool = False # router-before-all-reduce packed for every layer reason: str = "" # why a requested slot is not built @@ -77,34 +77,41 @@ def build(model, ctx) -> Bound: for i in range(m.start_layer, m.end_layer) if isinstance(m.layers[i], Qwen3_5AttentionDecoderLayer) ] + reasons = [] + skv2_on, mka_on = "skv2" in SLOTS, "mka" in SLOTS for i, layer in layers: - reason = "" - if "skv2" in SLOTS: + if skv2_on: reason = split_kv.unsupported(layer.attn, ctx.tokens_per_request) - if not reason and "mka" in SLOTS: + if reason: + reasons.append(f"layer {i}: {reason}") + skv2_on = False + if mka_on: reason = ( mka.unsupported(layer) if model.is_mrope_enabled else "no M-RoPE positions (MKA reads them)" ) - if reason: - return Bound(reason=f"layer {i}: {reason}") + if reason: + reasons.append(f"layer {i}: {reason}") + mka_on = False if "skv2_cp" in SLOTS: reason = split_kv.cp_reason() if reason: - return Bound(reason=f"slot skv2_cp: {reason}") - if "mka" in SLOTS: + reasons.append(f"slot skv2_cp: {reason}") + route_lin = mka_on and "route_lin" in SLOTS + if mka_on: for i, layer in layers: - reason = mka.pack(layer, route_lin="route_lin" in SLOTS) + reason = mka.pack(layer, route_lin=route_lin) if reason: - return Bound(reason=f"slot route_lin: layer {i}: {reason}") - mka_on = "mka" in SLOTS + reasons.append(f"slot route_lin: layer {i}: {reason}") + route_lin = False return Bound( - skv2="skv2" in SLOTS, + skv2=skv2_on, mka=mka_on, - mka_wide=mka_on and all(layer.num_heads == 4 for _, layer in layers), - mka_push="mka_push" in SLOTS, - route_lin="route_lin" in SLOTS, + mka_wide=mka_on, + mka_push=mka_on and "mka_push" in SLOTS, + route_lin=route_lin, + reason="; ".join(reasons), ) diff --git a/sgl_nvidia_mega/qwen3_5/attention/mka.py b/sgl_nvidia_mega/qwen3_5/attention/mka.py index 4b943e90..c170c15d 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/mka.py +++ b/sgl_nvidia_mega/qwen3_5/attention/mka.py @@ -16,8 +16,8 @@ the MoE input's all-reduce will run the hop for them, and that all-reduce only consumes (the branch's MKA_PUSH). -Steps of several requests (TP8, with ``multi_request``): up to ``FUSED_MAX_ROWS`` rows -the wide entries of the same kernels; past it the front path, which keeps SGLang's qkv +Steps of several requests (with ``multi_request``): TP8 uses the fused entries through +14 rows, TP4 through 8, and TP2 and DEP8 always use the front path, which keeps SGLang's qkv and o_proj GEMMs and fuses the rest bit for bit: the q / k norm, RoPE and gate split with the FP8 q and K / V write in one launch (the attention front of slot ``draft_front``), and the output gate with o_proj's static FP8 input quantization in one launch. @@ -106,10 +106,10 @@ def unsupported(layer) -> str: rope = layer.rotary_emb if not layer.attn_output_gate: return "no attention output gate" - if layer.num_heads not in (4, 8) or layer.num_kv_heads != 1: + if (layer.num_heads, layer.num_kv_heads) not in ((4, 1), (8, 1), (16, 1), (32, 2)): return ( f"{layer.num_heads} query / {layer.num_kv_heads} KV heads a rank " - "(MKA runs 4 or 8 / 1)" + "(MKA runs 4/1, 8/1, 16/1 or 32/2)" ) if layer.head_dim != 256 or layer.hidden_size != 4096: return f"head_dim {layer.head_dim}, hidden {layer.hidden_size} (MKA runs 256, 4096)" @@ -123,10 +123,14 @@ def unsupported(layer) -> str: def pool_reasons(kv) -> list[str]: - """Why the prologue cannot write these KV pools ([] when it can): one FP8 KV head - of 256 per slot.""" + """Why the prologue cannot write these KV pools ([] when it can): contiguous + NHD FP8 rows of one (TP8/TP4) or two (DEP8) heads of 256.""" for buf in (*kv.k, *kv.v): - if buf.dtype != torch.float8_e4m3fn or buf.stride(0) != 256: + if ( + buf.dtype != torch.float8_e4m3fn + or buf.shape[1:] not in ((1, 256), (2, 256)) + or not buf.is_contiguous() + ): return [f"a {buf.dtype} KV pool buffer of strides {buf.stride()}"] return [] @@ -136,6 +140,9 @@ def pack(layer, route_lin: bool = False) -> str: Returns why its router-before-all-reduce weights cannot be packed, when requested. """ + # TP2 and DEP8 keep the stock GEMMs; its generic front needs no packed weights. + if layer.num_heads > 8: + return _pack_route_lin(layer) if route_lin else "" from sgl_nvidia_mega.qwen3_5.sm103.attention import mka as kernels kernels.build(layer.qkv_proj.weight.device) @@ -204,14 +211,19 @@ def step(layer, x, st): from sgl_nvidia_mega.qwen3_5.sm103.attention import mka as kernels m = x.shape[0] - if m not in _ENGAGED: - _ENGAGED.add(m) + fused_max = ( + FUSED_MAX_ROWS if layer.num_heads == 4 else 8 if layer.num_heads == 8 else 0 + ) + key = (layer.num_heads, m) + if key not in _ENGAGED: + _ENGAGED.add(key) logger.info( - "NVIDIA megakernel: Qwen3.5 slot mka runs at %d rows (%s)", + "NVIDIA megakernel: Qwen3.5 slot mka runs at %d rows, %d query heads (%s)", m, - "fused" if m <= FUSED_MAX_ROWS else "front path", + layer.num_heads, + "fused" if m <= fused_max else "front path", ) - if m > FUSED_MAX_ROWS: + if m > fused_max: return _front_step(layer, x, st) attn = layer.attn width = layer.num_heads * layer.head_dim diff --git a/sgl_nvidia_mega/qwen3_5/attention/route.py b/sgl_nvidia_mega/qwen3_5/attention/route.py index 3339bcfb..b302fb81 100644 --- a/sgl_nvidia_mega/qwen3_5/attention/route.py +++ b/sgl_nvidia_mega/qwen3_5/attention/route.py @@ -8,7 +8,7 @@ picked without reading the device. ``VERIFY`` (8192 tokens, as the TP8 serving configuration on SGLang branch -nvidia/qwen35-agentx-gb300 sets ``SGLANG_Q35MF_MKA_SKV2_VERIFY_MIN_LENGTH``) routes the +nvidia/qwen35-agentx-gb300) routes the target verify and ``EXTEND`` (32768, that branch's ``SPLIT_KV_MIN_LENGTH``) the draft extend. A built megakernel offers its variants per SGLang runner (``graph_variants``: "target_verify", "draft_loop", "draft_extend"); a runner captures each of diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index d9c1232b..f6addea8 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -46,7 +46,7 @@ REQUESTS = tuple(range(1, 129)) # the TP sizes whose steps of more than one request are covered when requested (checked at # TP8); at the others a covered step is one request -MULTI_REQUEST_TP = (8,) +MULTI_REQUEST_TP = (2, 4, 8) # the name that requests them in the slot list MULTI_REQUEST = "multi_request" GDN_FRONT_ROWS = 7 # the one-request custom front's fixed width diff --git a/test/kernels/test_qwen3_5_attention.py b/test/kernels/test_qwen3_5_attention.py index 6ad78fd8..ad616e1f 100644 --- a/test/kernels/test_qwen3_5_attention.py +++ b/test/kernels/test_qwen3_5_attention.py @@ -591,14 +591,11 @@ def _bits(a, b): ) -def _front_operands(rows, g, mrope=True, slots=64 * 16): - q_gate = (torch.randn(rows, HEADS * 2 * DIM, device="cuda", generator=g) * 3).to( - torch.bfloat16 - ) - k, v = ( - (torch.randn(rows, DIM, device="cuda", generator=g) * 3).to(torch.bfloat16) - for _ in range(2) - ) +def _front_operands(rows, g, mrope=True, slots=64 * 16, heads=HEADS, kv_heads=1): + # As qkv_proj's split: all three inputs keep the packed projection's row stride. + widths = (heads * 2 * DIM, kv_heads * DIM, kv_heads * DIM) + qkv = (torch.randn(rows, sum(widths), device="cuda", generator=g) * 3).bfloat16() + q_gate, k, v = qkv.split(widths, dim=-1) norms = [ (torch.randn(DIM, device="cuda", generator=g) * 0.1).to(torch.bfloat16) for _ in range(2) @@ -611,11 +608,13 @@ def _front_operands(rows, g, mrope=True, slots=64 * 16): else: pos, axis = torch.randint(0, 4096, (rows,), device="cuda", generator=g), None pools = tuple( - torch.randn(slots, 1, DIM, device="cuda", generator=g).to(torch.float8_e4m3fn) + torch.randn(slots, kv_heads, DIM, device="cuda", generator=g).to( + torch.float8_e4m3fn + ) for _ in range(2) ) loc = torch.randperm(slots, device="cuda", generator=g)[:rows] - args = (q_gate, k, *norms, cos_sin, pos, 1e-6, HEADS, 1, DIM, 64) + args = (q_gate, k, *norms, cos_sin, pos, 1e-6, heads, kv_heads, DIM, 64) return args, axis, v, pools, loc @@ -669,7 +668,10 @@ def run(): @pytest.mark.parametrize("rows", [7, 14, 28, 56, 112, 224, 448, 896]) -def test_front_with_the_cache_write_is_sglangs_front_then_fp8_kv_write(rows): +@pytest.mark.parametrize("heads,kv_heads", [(4, 1), (16, 1), (32, 2)]) +def test_front_with_the_cache_write_is_sglangs_front_then_fp8_kv_write( + rows, heads, kv_heads +): """The front with its FP8 KV write (one launch) gives the bytes of SGLang's front followed by its fused FP8 q / k / v + KV write, at every covered width.""" from sglang.kernels.ops.attention.fused_qk_rmsnorm_rope_gate import ( @@ -684,7 +686,9 @@ def test_front_with_the_cache_write_is_sglangs_front_then_fp8_kv_write(rows): ) g = torch.Generator(device="cuda").manual_seed(300 + rows) - args, axis, v, pools, loc = _front_operands(rows, g, slots=64 * 32) + args, axis, v, pools, loc = _front_operands( + rows, g, slots=64 * 32, heads=heads, kv_heads=kv_heads + ) scales = [torch.full((1,), s, device="cuda") for s in (0.8, 1.3)] want_pools = tuple(p.clone() for p in pools) q, k, gate = sglang_front(*args, has_gate=True, mrope_axis_map=axis) @@ -710,15 +714,16 @@ def test_front_with_the_cache_write_is_sglangs_front_then_fp8_kv_write(rows): @pytest.mark.parametrize("rows", [7, 14, 28, 56, 112, 224, 448, 896]) -def test_gate_quant_is_sglangs_sigmoid_mul_then_static_quant(rows): +@pytest.mark.parametrize("heads", [4, 16, 32]) +def test_gate_quant_is_sglangs_sigmoid_mul_then_static_quant(rows, heads): from sglang.kernels.ops.elementwise.elementwise import fused_sigmoid_mul from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 from sgl_nvidia_mega.qwen3_5.sm103.attention.gate_quant import gate_quant g = torch.Generator(device="cuda").manual_seed(400 + rows) - attn = (torch.randn(rows, HEADS * DIM, device="cuda", generator=g) * 4).bfloat16() - gate = (torch.randn(rows, HEADS * DIM, device="cuda", generator=g) * 3).bfloat16() + attn = (torch.randn(rows, heads * DIM, device="cuda", generator=g) * 4).bfloat16() + gate = (torch.randn(rows, heads * DIM, device="cuda", generator=g) * 3).bfloat16() scale = torch.full((1,), 0.013, device="cuda") want, _ = static_quant_fp8(fused_sigmoid_mul(attn, gate), scale) assert _bits(want, gate_quant(attn, gate, scale)) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 8648e9b1..9fa311cc 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -1310,9 +1310,9 @@ def test_the_draft_loop_writes_each_steps_token_and_restores_the_positions( [ (8, {"multi_request"}, (1, 2)), (8, set(), (1,)), - (4, {"multi_request"}, (1,)), + (4, {"multi_request"}, (1, 2)), (8, {"multi_request", "draft_multi"}, tuple(range(1, 129))), - (4, {"multi_request", "draft_multi"}, (1,)), + (4, {"multi_request", "draft_multi"}, tuple(range(1, 129))), ], ) def test_the_draft_loop_serves_two_requests_at_tp8_on_request( @@ -1331,7 +1331,7 @@ def test_the_draft_loop_serves_two_requests_at_tp8_on_request( [ (8, {"multi_request", "draft_multi"}, True, tuple(range(1, 129))), (8, {"multi_request", "draft_multi"}, False, (1,)), - (4, {"multi_request", "draft_multi"}, True, (1,)), + (4, {"multi_request", "draft_multi"}, True, tuple(range(1, 129))), (8, {"draft_multi"}, True, (1,)), ], ) diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 2a13bcb8..133841b1 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -440,7 +440,7 @@ def test_the_slot_plan_gates_each_slot_by_the_rows_its_kernels_serve(): # one-request step it is checked at assert plan.multi_request(8, frozenset({"multi_request"})) assert not plan.multi_request(8, frozenset()) - assert not plan.multi_request(4, frozenset({"multi_request"})) + assert plan.multi_request(4, frozenset({"multi_request"})) assert plan.covered((7, 14, 28), every, 7) == (7, 14, 28) assert plan.covered((7, 448, 896), every, 7) == (7, 448, 896) # one request is always covered, more only where a slot runs diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 5c89fa87..72cf5e81 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -2890,7 +2890,13 @@ def symmetric(group, disabled=False): @pytest.mark.parametrize("consumed", [False, True]) -def test_the_mka_wide_path_offers_gate_quantization_to_cp_lite(consumed): +@pytest.mark.parametrize( + "heads,kv_heads,rows", + [(4, 1, 21), (16, 1, 7), (32, 2, 224)], +) +def test_the_mka_wide_path_offers_gate_quantization_to_cp_lite( + consumed, heads, kv_heads, rows +): from sglang.srt.layers.quantization import fp8_utils from sgl_nvidia_mega.qwen3_5 import lowm, runtime @@ -2899,16 +2905,16 @@ def test_the_mka_wide_path_offers_gate_quantization_to_cp_lite(consumed): from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite from sgl_nvidia_mega.qwen3_5.sm103.draft import front - rows = 21 log = _log() - attn, st = _skv2_setup(log, batch_size=3, rows=rows) - layer = _mka_layer(log, attn) + attn, st = _skv2_setup(log, batch_size=rows // 7, rows=rows) + layer = _mka_layer(log, attn, num_heads=heads, num_kv_heads=kv_heads) + width = heads * 256 x = torch.zeros(rows, 4096, dtype=torch.bfloat16) - qkv = torch.zeros(rows, 2560, dtype=torch.bfloat16) - q = torch.zeros(rows, 1024, dtype=torch.float8_e4m3fn) - gate = torch.zeros(rows, 1024, dtype=torch.bfloat16) - attention_out = torch.zeros(rows, 1024, dtype=torch.bfloat16) - fallback = torch.zeros(rows, 1024, dtype=torch.float8_e4m3fn) + qkv = torch.zeros(rows, 2 * width + 2 * kv_heads * 256, dtype=torch.bfloat16) + q = torch.zeros(rows, width, dtype=torch.float8_e4m3fn) + gate = torch.zeros(rows, width, dtype=torch.bfloat16) + attention_out = torch.zeros(rows, width, dtype=torch.bfloat16) + fallback = torch.zeros(rows, width, dtype=torch.float8_e4m3fn) result = torch.zeros(rows, 4096, dtype=torch.bfloat16) qkv_linear = layer.qkv_proj layer.qkv_proj = _Callable(lambda _: (qkv, None), **vars(qkv_linear)) @@ -2952,13 +2958,18 @@ def row_parallel(linear, o_proj_input, dtype, what, hop): out = mka.step(layer, x, st) assert out is result + if rows == 7: + # CP multi needs more than one request: B1 quantizes its own gate + assert [call[0] for call in calls] == ["gate_quant", "o_proj"] + assert calls[1][2] is fallback + return assert calls[0][0] == "offer" offered_gate = calls[0][1] assert offered_gate.shape == gate.shape assert offered_gate.data_ptr() == gate.data_ptr() assert calls[0][2] is layer.o_proj.input_scale offered = calls[0][3] - assert tuple(offered.shape) == (rows, 1024) + assert tuple(offered.shape) == (rows, width) assert offered.dtype == torch.float8_e4m3fn assert calls[1] == ("take",) if consumed: @@ -3043,6 +3054,8 @@ def symmetric(group, disabled=False): "kw, want", [ (dict(), ""), + (dict(num_heads=16), ""), + (dict(num_heads=32, num_kv_heads=2), ""), (dict(attn_output_gate=False), "no attention output gate"), (dict(num_heads=2), "2 query / 1 KV heads a rank"), (dict(num_kv_heads=2), "4 query / 2 KV heads a rank"), @@ -3108,8 +3121,8 @@ def layer(kind): ctx = SimpleNamespace(tokens_per_request=7) for made in layers[1:]: made.num_heads = 8 - # TP4's eight query heads a rank: the narrow kernels only - assert attention.build(model, ctx) == attention.Bound(mka=True) + # TP4's narrow kernels and the generic front at wider rows. + assert attention.build(model, ctx) == attention.Bound(mka=True, mka_wide=True) packed.clear() for made in layers[1:]: made.num_heads = 4 @@ -3136,12 +3149,40 @@ def kv(*bufs): assert ( mka.pool_reasons(kv(torch.zeros(64, 1, 256, dtype=torch.float8_e4m3fn))) == [] ) - assert mka.pool_reasons(kv(torch.zeros(64, 2, 256, dtype=torch.float8_e4m3fn))) == [ - "a torch.float8_e4m3fn KV pool buffer of strides (512, 256, 1)" + assert ( + mka.pool_reasons(kv(torch.zeros(64, 2, 256, dtype=torch.float8_e4m3fn))) == [] + ) + assert mka.pool_reasons(kv(torch.zeros(64, 4, 256, dtype=torch.float8_e4m3fn))) == [ + "a torch.float8_e4m3fn KV pool buffer of strides (1024, 256, 1)" ] + assert mka.pool_reasons( + kv(torch.zeros(64, 1, 512, dtype=torch.float8_e4m3fn)[..., ::2]) + ) assert "torch.float32" in mka.pool_reasons(kv(torch.zeros(64, 1, 256)))[0] +@pytest.mark.parametrize("missing", ["mka", "skv2"]) +def test_an_unavailable_attention_slot_keeps_the_independent_slot(monkeypatch, missing): + from sgl_nvidia_mega.qwen3_5.attention import mka, split_kv + + layer = qwen3_5.Qwen3_5AttentionDecoderLayer.__new__( + qwen3_5.Qwen3_5AttentionDecoderLayer + ) + torch.nn.Module.__init__(layer) + layer.attn = SimpleNamespace() + model = SimpleNamespace( + is_mrope_enabled=True, + model=SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]), + ) + monkeypatch.setattr(runtime, "SLOTS", frozenset({"mka", "skv2"})) + reasons = {missing: "unsupported"} + monkeypatch.setattr(mka, "unsupported", lambda _: reasons.get("mka", "")) + monkeypatch.setattr(split_kv, "unsupported", lambda *_: reasons.get("skv2", "")) + monkeypatch.setattr(mka, "pack", lambda *a, **k: "") + bound = attention.build(model, SimpleNamespace(tokens_per_request=7)) + assert (bound.mka, bound.skv2) == (missing != "mka", missing != "skv2") + + class _Exit: def __init__(self, log, defer): self.log, self.defer_moe_finalize = log, defer From 92950d6f8d4d4b81daabc2d33d719bdb02f29bfd Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 2 Nov 2026 00:46:52 -0800 Subject: [PATCH 47/62] perf(qwen3.5): support router-only slots independently of expert engines --- sgl_nvidia_mega/qwen3_5/moe/__init__.py | 11 ++++++++++- sgl_nvidia_mega/qwen3_5/moe/route.py | 22 +++++++++++++--------- test/kernels/test_qwen3_5_moe.py | 19 ++++++++++++++----- test/unit/test_qwen3_5_steps.py | 9 +++++---- 4 files changed, 42 insertions(+), 19 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/moe/__init__.py b/sgl_nvidia_mega/qwen3_5/moe/__init__.py index ef1b8df1..5cb94e69 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/moe/__init__.py @@ -50,6 +50,7 @@ "TopKOutputChecker", "format_is_bypassed", "format_is_standard", + "empty_topk_output", ), "srt/layers/moe/fused_moe_triton/layer.py": ("forward_deferred_finalize",), "srt/layers/moe/utils.py": ("sum_post_experts_output",), @@ -167,9 +168,17 @@ def step(layer, x: torch.Tensor, st, defer: bool, reduce=None): mlp = layer.mlp rows, hidden = x.shape x = x.view(-1, hidden) + if defer and rows == 0: + raise RuntimeError("Qwen deferred finalize does not support M=0") # the gate is a plain linear on CUDA (no AMX, no NPU) use_fused_gate = mlp.shared_expert_gate is not None - if ( + if rows == 0: + # Empty DP ranks skip the shared expert and gate, but must participate + # in the experts' collective in the same way as SGLang's forward. + shared = None + topk_output = mlp.topk.empty_topk_output(x.device) + routed = mlp.experts(x, topk_output) + elif ( mlp.alt_stream is not None and get_is_capture_mode() and not torch.compiler.is_compiling() diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index b7d488ee..76dafc68 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -12,8 +12,7 @@ or slot ``mkm`` (``mkm.py``) takes the router front's outputs in its place. Slot ``route_shared`` (numeric) also computes the shared expert's static-FP8 gate/up -projection in the router front's launch (SGLANG_Q35MF_ROUTE_SHARED on SGLang branch -nvidia/qwen35-agentx-gb300), so the shared expert on the alternate stream starts at its +projection in the router front's launch, so the shared expert on the alternate stream starts at its activation, after that launch only; the gate/up sums in another order than the FP8 linear's GEMM. @@ -74,6 +73,9 @@ "TopKOutputChecker", "format_is_bypassed", "StandardTopKOutput", + "TopKOutputFormat", + "BYPASSED", + "output_format", ), "srt/layers/moe/fused_moe_triton/layer.py": ( "dispatcher", @@ -289,6 +291,7 @@ def build_m32( def unsupported(mlp) -> str: """Why the router front cannot stand in for ``mlp``'s gate and the runner's input quantizer ("" when it can).""" + from sglang.srt.layers.moe.topk import TopKOutputFormat from sglang.srt.models.qwen2_moe import Qwen2MoeSparseMoeBlock from sgl_nvidia_mega.qwen3_5.sm103.moe import router @@ -308,23 +311,24 @@ def unsupported(mlp) -> str: return "per-token activation scales" if not mlp.supports_deferred_finalize: return "no deferred finalize" + # TopK binds None to BYPASSED for the TRTLLM runner during model setup. + # Check the bound contract, before compiling or claiming engagement. + if mlp.topk.topk_config.output_format != TopKOutputFormat.BYPASSED: + return "the router front requires logits-based bypassed TopK" return "" def topk_unsupported(model) -> str: """Why slot route_topk cannot serve every MoE layer of ``model`` ("" when it can): - the router front's gate and input quantizer, and the trtllm-gen weights with a - softmax top-10 renormalized over 512 experts. Under expert parallelism only the - materialized-routing contract applies; the plugin's full-expert engines do not.""" + the router front's gate and input quantizer, and a softmax top-10 renormalized + over 512 experts. The stock runner owns the expert GEMMs: TP2's intermediate + of 512 does not need a matching mkm engine.""" from sgl_nvidia_mega.qwen3_5.moe import mkm m = model.model for i in range(m.start_layer, m.end_layer): mlp = m.layers[i].mlp - expert_parallel = mlp.experts.num_local_experts != mlp.experts.num_experts - reason = unsupported(mlp) or ( - mkm.routing_unsupported(mlp) if expert_parallel else mkm.unsupported(mlp) - ) + reason = unsupported(mlp) or mkm.routing_unsupported(mlp) if reason: return f"layer {i}: {reason}" return "" diff --git a/test/kernels/test_qwen3_5_moe.py b/test/kernels/test_qwen3_5_moe.py index 12d0d961..af32c262 100644 --- a/test/kernels/test_qwen3_5_moe.py +++ b/test/kernels/test_qwen3_5_moe.py @@ -80,11 +80,20 @@ def _experts(g, rows, cols, std, n=E): def layer(request): """One TP8 (intermediate 128 a rank) or TP4 (256) Qwen3.5 MoE layer on the trtllm-gen weight layout, as SGLang loads it.""" + return _layer(request.param, maps=True) + + +@pytest.fixture(scope="module", params=[128, 256, 512], ids=lambda i: f"I{i}") +def routing_layer(request): + """TP8/TP4/TP2 weights for the stock routed runner, without an mkm engine.""" + return _layer(request.param, maps=False) + + +def _layer(inter, maps): from sglang.srt.layers.quantization.utils import ( prepare_static_weights_for_trtllm_fp4_moe, ) - inter = request.param # intermediate per rank g = torch.Generator(device="cuda").manual_seed(11) w13, sf13, s13 = _experts(g, 2 * inter, H, 0.02) w2, sf2, s2 = _experts(g, H, inter, 0.02) @@ -100,7 +109,7 @@ def layer(request): "g1": g1, "g1c": (g1 / H_SCALE).float(), "g2": (H_SCALE / s2).float(), - "maps": mkm.weight_maps(w13, w2), + "maps": mkm.weight_maps(w13, w2) if maps else None, "inter": inter, } @@ -605,7 +614,7 @@ def _topk_inputs(m, seed): @pytest.mark.parametrize("m", [1, 7, 17, 21, 28, 56, 112, 224, 448, 896]) -def test_route_topk_is_the_trtllm_routing_and_the_runners_quantizer(layer, m): +def test_route_topk_is_the_trtllm_routing_and_the_runners_quantizer(routing_layer, m): from flashinfer import fp4_quantize for seed in range(2): @@ -616,8 +625,8 @@ def test_route_topk_is_the_trtllm_routing_and_the_runners_quantizer(layer, m): assert torch.equal(x_sf, q_sf.view(torch.uint8).reshape(m, -1)) # the routed experts on this routing equal the runner's own routing from the # logits, row for row, weights included - want_rows, want_w = _stock(layer, logits, x_fp4, x_sf) - got_rows, got_w = _routed_stock(layer, (ids, weights), x_fp4, x_sf) + want_rows, want_w = _stock(routing_layer, logits, x_fp4, x_sf) + got_rows, got_w = _routed_stock(routing_layer, (ids, weights), x_fp4, x_sf) assert torch.equal(got_w.view(torch.int16), want_w.view(torch.int16)) assert torch.equal(got_rows.view(torch.int16), want_rows.view(torch.int16)) diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 72cf5e81..1a5bc6e6 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -897,12 +897,13 @@ def meta(dtype, *shape): assert mkm.unsupported(mlp).startswith("an expert activation") -def test_route_topk_checks_only_materialized_routing_under_expert_parallelism( - monkeypatch, +@pytest.mark.parametrize("local_experts", [128, 512]) +def test_route_topk_checks_materialized_routing_without_a_fused_expert_engine( + monkeypatch, local_experts ): from sgl_nvidia_mega.qwen3_5.moe import mkm, route - experts = SimpleNamespace(num_experts=512, num_local_experts=128) + experts = SimpleNamespace(num_experts=512, num_local_experts=local_experts) mlp = SimpleNamespace(experts=experts) model = SimpleNamespace( model=SimpleNamespace( @@ -915,7 +916,7 @@ def test_route_topk_checks_only_materialized_routing_under_expert_parallelism( mkm, "routing_unsupported", lambda _: checked.append("routing") or "" ) monkeypatch.setattr( - mkm, "unsupported", lambda _: pytest.fail("full-expert engine checked under EP") + mkm, "unsupported", lambda _: pytest.fail("full-expert engine checked") ) assert route.topk_unsupported(model) == "" From a39fab2c9ffd5b33411fc75c5c6c69a6c1b0c57d Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 2 Nov 2026 00:46:52 -0800 Subject: [PATCH 48/62] perf(qwen3.5): enable BF16 draft on measured TP4 layouts --- sgl_nvidia_mega/qwen3_5/draft/attention.py | 2 +- sgl_nvidia_mega/qwen3_5/draft/extend.py | 6 ++- sgl_nvidia_mega/qwen3_5/draft/hop.py | 4 +- sgl_nvidia_mega/qwen3_5/draft/xchg.py | 4 ++ sgl_nvidia_mega/qwen3_5/sm103/draft/front.py | 2 +- sgl_nvidia_mega/qwen3_5/sm103/draft/q1.py | 2 +- .../qwen3_5/sm103/draft/src/mkd_xchg.cu | 9 ++-- sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py | 22 +++++---- test/kernels/test_qwen3_5_draft_front.py | 4 +- test/kernels/test_qwen3_5_draft_moe.py | 32 ++++++++----- test/kernels/test_qwen3_5_draft_xchg.py | 48 ++++++++++++++++++- test/unit/test_qwen3_5_runtime.py | 32 +++++++++++++ 12 files changed, 132 insertions(+), 35 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/draft/attention.py b/sgl_nvidia_mega/qwen3_5/draft/attention.py index 552ed214..8535bc75 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/attention.py +++ b/sgl_nvidia_mega/qwen3_5/draft/attention.py @@ -26,7 +26,7 @@ # the KV-length graph variants of the draft loop (2.2's graph_variants.Q1_STOCK / # Q1_LONG): trtllm-gen below Q1_MIN_LENGTH, SKV2-Q1 from it (2.2's -# SGLANG_Q35MF_MKA_Q1_MIN_LENGTH default, which its TP8 serving configuration keeps) +# the source branch's default Q1 threshold, which its TP8 serving configuration keeps) Q1_STOCK, Q1_LONG = "q1_stock", "q1_long" Q1_MIN_LENGTH = 16384 # slot draft_q1_cp: one request past attention.route.CP_MIN_KV["draft"] runs its steps' diff --git a/sgl_nvidia_mega/qwen3_5/draft/extend.py b/sgl_nvidia_mega/qwen3_5/draft/extend.py index 33265b7b..a2c6f21d 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/extend.py +++ b/sgl_nvidia_mega/qwen3_5/draft/extend.py @@ -108,10 +108,14 @@ def _multi_step(self, io) -> None: """More requests (slot draft_multi): SGLang's model step on every verify row, then the selected rows' sharded top-1 (torch.argmax's choice, a NaN first) straight into SGLang's outputs.""" + from sgl_nvidia_mega.qwen3_5.draft.xchg import FC_SPLIT_MAX_TOKENS from sgl_nvidia_mega.qwen3_5.sm103.draft.top1 import NAN_KEY_TORCH_ARGMAX parts = self.parts - x = layer.fc_input(parts, io.input_ids, io.hidden_states) + if parts.mkd is not None and io.input_ids.shape[0] <= FC_SPLIT_MAX_TOKENS: + x = parts.mkd.split_fc(parts.embed(io.input_ids), io.hidden_states) + else: + x = layer.fc_input(parts, io.input_ids, io.hidden_states) hidden = layer.layer( parts, x, diff --git a/sgl_nvidia_mega/qwen3_5/draft/hop.py b/sgl_nvidia_mega/qwen3_5/draft/hop.py index a0b0dd25..3d03e57b 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/hop.py +++ b/sgl_nvidia_mega/qwen3_5/draft/hop.py @@ -38,9 +38,9 @@ MAX_M = 1 # Whether the attention output's o_proj runs as the one-row GEMV that pushes into the FFN # input's hop, per TP size, as SGLang branch nvidia/qwen35-agentx-gb300 runs them: TP8 does -# (SGLANG_Q35MF_MKD_GEMV); TP4 runs SGLang's o_proj and the hop publishes its output (that +# (its GEMV path); TP4 runs SGLang's o_proj and the hop publishes its output (that # branch's TP4 setup also pushes the single-launch MoE into the final norm's hop, -# SGLANG_Q35MF_MKD_HOP_PUSH, which changes no bits and is not ported). +# its hop-push path, which changes no bits and is not ported). GEMV_O_PROJ = {8: True, 4: False} diff --git a/sgl_nvidia_mega/qwen3_5/draft/xchg.py b/sgl_nvidia_mega/qwen3_5/draft/xchg.py index 69460f62..7a9383a0 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/xchg.py +++ b/sgl_nvidia_mega/qwen3_5/draft/xchg.py @@ -390,6 +390,10 @@ def head(self, hidden_states: torch.Tensor) -> torch.Tensor: """The one-row draft logits [1, N] bf16 on this rank's shard (the exchange then reads the per-CTA maxima). For an armed step with the head installed.""" x = hidden_states.to(torch.bfloat16).contiguous() + if self.head_part is None: + # TP2's BF16 head shard is wider than the fused head's 16-bit row key. + # The exchange selects over its full row ids, behind SGLang's linear op. + return torch.nn.functional.linear(x, self.lm_head_weight) logits = torch.empty( 1, self.lm_head_weight.shape[0], dtype=torch.bfloat16, device=x.device ) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/front.py b/sgl_nvidia_mega/qwen3_5/sm103/draft/front.py index b2cc6551..e726a277 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/front.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/front.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Adapted from SGLang (Apache-2.0, SGLang Team): # python/sglang/kernels/ops/attention/fused_qk_rmsnorm_rope_gate.py, as SGLang branch -# nvidia/qwen35-agentx-gb300 extends it with the FP8 KV-cache write (SGLANG_Q35MF_ATTN_FRONT; +# nvidia/qwen35-agentx-gb300 extends it with the FP8 KV-cache write (attention-front path; # 98573b3e69, 0d86b61136, 4115dd1469): copied verbatim. """Fused Q/K GemmaRMSNorm + NeoX RoPE + gate deinterleave (Triton), and with ``kv_cache`` 2.2's attention front (ATTN_FRONT): FP8 q, and K / V written straight diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/q1.py b/sgl_nvidia_mega/qwen3_5/sm103/draft/q1.py index d2c4042b..6ba899a2 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/q1.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/q1.py @@ -54,7 +54,7 @@ KV_HEADS = 1 # the draft's KV heads on a rank: 2 KV heads over TP8, replicated QUERY_HEADS = (4, 8) # the KV length from which the branch serves the draft's q1 decode with SKV2-Q1 -# (SGLANG_Q35MF_MKA_Q1_MIN_LENGTH, unset in its TP8 serving configuration) +# (the default threshold, unchanged in its TP8 serving configuration) MIN_LENGTH = 16384 # the merge's head-dim block and warps (Q1SplitPlan.run) MERGE_BLOCK_D, MERGE_WARPS = 32, 4 diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_xchg.cu b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_xchg.cu index d94ea7c3..b4952b7b 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_xchg.cu +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/src/mkd_xchg.cu @@ -175,9 +175,12 @@ SGL_DEVICE void mkd_xchg_kernel(const MkdXchgParams& p) { constexpr int kScanUnroll = 4; const int nvec = p.vocab / 8; const bool from_head = p.head_part != nullptr; - const bool split = !from_head && p.scan_part != nullptr && nvec >= mkd_xchg::kSplitScanMinVecs; + // TP2's vocabulary shard exceeds the packed key's 16-bit row. Reuse the exact + // full-index scan below; narrow shards retain their packed/head fast path. + const bool wide = p.vocab > 65536; + const bool split = !wide && !from_head && p.scan_part != nullptr && nvec >= mkd_xchg::kSplitScanMinVecs; const int lo = split ? static_cast((static_cast(nvec) * c) / kChunks) : 0; - const int hi = from_head ? 0 : split ? static_cast((static_cast(nvec) * (c + 1)) / kChunks) : nvec; + const int hi = (wide || from_head) ? 0 : split ? static_cast((static_cast(nvec) * (c + 1)) / kChunks) : nvec; for (int v0 = lo + static_cast(threadIdx.x); v0 < hi; v0 += kThreads * kScanUnroll) { uint4 x[kScanUnroll]; #pragma unroll @@ -250,7 +253,7 @@ SGL_DEVICE void mkd_xchg_kernel(const MkdXchgParams& p) { int32_t bidx = static_cast(0xffffu - (best & 0xffffu)); const uint32_t k16 = best >> 16; best = (k16 << 16) | (k16 < 0x8000u ? 0xffffu : 0u); - if (nan_any != 0u) { + if (wide || nan_any != 0u) { __syncthreads(); best = 0u; bidx = 0x7fffffff; diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py b/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py index fb089f59..e284229f 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py @@ -37,7 +37,7 @@ SRC = Path(__file__).parent / "src" HIDDEN = 4096 -WORLDS = (4, 8) +WORLDS = (2, 4, 8) # CTAs per exchange; each owns HIDDEN / CHUNKS columns of the candidate and partial rows CHUNKS = 16 XCHG_THREADS = 1024 @@ -379,8 +379,10 @@ def mkd_xchg( fc_out = None _check(topk_index, "topk_index", (1, 1), torch.int64) _check(topk_p, "topk_p", (1, 1), torch.float32) - if v % 8 != 0 or v > 65536: - raise ValueError("mkd_xchg: local vocab must be a multiple of 8 and <= 65536") + if v % 8 != 0: + raise ValueError("mkd_xchg: local vocab must be a multiple of 8") + if v > 65536 and head_part is not None: + raise ValueError("mkd_xchg: packed head maxima require a <=65536-row shard") head_parts = head_part.numel() if head_part is not None else 0 if head_parts > XCHG_THREADS: raise ValueError("mkd_xchg: at most one head partial per exchange thread") @@ -626,11 +628,15 @@ def mkd_draft_head( def draft_fc_partial_rows(out_rows: int, width: int, num_sms: int) -> Optional[int]: - """Weight rows per CTA for the fused kernel, or None if they do not fit in shared - memory.""" + """Weight rows per CTA, capped by shared memory for TP2's full-half slice. + + The dot and norm reductions are independent of the output-row partition. TP4/8 + retain one CTA per SM; TP2 may need more CTAs to fit the same BF16 operands. + """ rows = -(-out_rows // num_sms) static = MAX_TOKENS * width * 2 - return rows if rows * width * 2 + static <= _MAX_SMEM_BYTES else None + fit = (_MAX_SMEM_BYTES - static) // (width * 2) + return min(rows, fit) if fit > 0 else None def draft_fc_partial( @@ -663,9 +669,7 @@ def draft_fc_partial( _check(out, "out", (m, HIDDEN), torch.float32) if not (col_start >= 0 and col_start + width <= HIDDEN and col_start % 8 == 0): raise ValueError("draft_fc_partial: slice out of range") - grid = _sm_count(device.index) - if (HIDDEN + grid - 1) // grid > max_rows: - raise ValueError("draft_fc_partial: more weight rows per CTA than kMaxRows") + grid = max(_sm_count(device.index), -(-HIDDEN // max_rows)) p = DraftFcPartialParams( x=x.data_ptr(), norm_w=norm_weight.data_ptr(), diff --git a/test/kernels/test_qwen3_5_draft_front.py b/test/kernels/test_qwen3_5_draft_front.py index 05354a98..b527d0a2 100644 --- a/test/kernels/test_qwen3_5_draft_front.py +++ b/test/kernels/test_qwen3_5_draft_front.py @@ -109,7 +109,7 @@ def _same(a, b): @cuda -@pytest.mark.parametrize("heads", [4, 8]) +@pytest.mark.parametrize("heads", [4, 8, 16]) @pytest.mark.parametrize("rows", [1, 2, 7]) @pytest.mark.parametrize("seed", [0, 1, 2]) def test_the_front_writes_what_sglangs_two_launches_write(rows, seed, heads): @@ -123,7 +123,7 @@ def test_the_front_writes_what_sglangs_two_launches_write(rows, seed, heads): @cuda -@pytest.mark.parametrize("heads", [4, 8]) +@pytest.mark.parametrize("heads", [4, 8, 16]) def test_the_front_replays_in_a_graph_with_new_positions_and_slots(heads): rows = 1 inputs = _inputs(rows, 5, heads) diff --git a/test/kernels/test_qwen3_5_draft_moe.py b/test/kernels/test_qwen3_5_draft_moe.py index 5e69f5b0..c29a5335 100644 --- a/test/kernels/test_qwen3_5_draft_moe.py +++ b/test/kernels/test_qwen3_5_draft_moe.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """The draft package's single-launch MoE (sm103/draft/moe.py) on one SM100-family GPU -(GB300), at the TP8 draft shape (hidden 4096, 512 experts of intermediate 128, top-10): +(GB300), at the TP8 / TP4 draft shapes (hidden 4096, 512 experts, intermediate 128 / 256, top-10): - the top-10 is exactly the kernel's own selection rule on the logits it published (bf16 order keys, ties to the lowest id), and the weights its renormalized softmax; - logits, weights and output match a float64 reference with the kernel's rounding @@ -32,9 +32,10 @@ def _gpu_ready() -> bool: pytestmark = pytest.mark.skipif(not _gpu_ready(), reason="needs an SM100-family GPU") -@pytest.fixture(scope="module") -def weights(): - """One TP8 draft MoE block's bf16 weights in checkpoint layout.""" +@pytest.fixture(scope="module", params=[128, 256], ids=["tp8", "tp4"]) +def weights(request): + """One BF16 draft MoE block's bf16 weights in checkpoint layout.""" + inter = request.param g = torch.Generator(device="cuda").manual_seed(5) def randn(*shape, std): @@ -44,10 +45,10 @@ def randn(*shape, std): return { "router_w": randn(E, H, std=0.02), - "w13": randn(E, 2 * I, H, std=0.02), - "w2": randn(E, H, I, std=0.05), - "w13_shared": randn(2 * I, H, std=0.02), - "w2_shared": randn(H, I, std=0.05), + "w13": randn(E, 2 * inter, H, std=0.02), + "w2": randn(E, H, inter, std=0.05), + "w13_shared": randn(2 * inter, H, std=0.02), + "w2_shared": randn(H, inter, std=0.05), "gate_w": randn(H, std=0.02), } @@ -105,13 +106,18 @@ def _r(t: torch.Tensor) -> torch.Tensor: return t.to(torch.bfloat16).double() +def _inter(w): + return w["w2"].shape[-1] + + def _reference(w, x, ids, weights): """float64 dots with the kernel's bf16 rounding points.""" + inter = _inter(w) xd = x.double().view(-1) acts = [] for s in range(K + 1): rows = w["w13_shared"] if s == K else w["w13"][int(ids[s])] - gate, up = _r(rows[:I].double() @ xd), _r(rows[I:].double() @ xd) + gate, up = _r(rows[:inter].double() @ xd), _r(rows[inter:].double() @ xd) acts.append(_r(gate / (1.0 + torch.exp(-gate)) * up)) routed = torch.zeros(H, dtype=torch.float64, device="cuda") for s in range(K): @@ -126,7 +132,7 @@ def _reference(w, x, ids, weights): def test_selection_weights_and_output(weights): from sgl_nvidia_mega.qwen3_5.sm103.draft import moe - scratch = moe.draft_moe_fused_scratch(E, K, I, torch.device("cuda")) + scratch = moe.draft_moe_fused_scratch(E, K, _inter(weights), torch.device("cuda")) for seed in range(4): x = _x(seed) topk_weights, topk_ids, out = _run(weights, x, scratch, _outputs()) @@ -161,7 +167,7 @@ def test_selection_weights_and_output(weights): def test_repeat_and_graph_replay(weights): from sgl_nvidia_mega.qwen3_5.sm103.draft import moe - scratch = moe.draft_moe_fused_scratch(E, K, I, torch.device("cuda")) + scratch = moe.draft_moe_fused_scratch(E, K, _inter(weights), torch.device("cuda")) eager = {} for seed in range(3): a = [t.clone() for t in _run(weights, _x(seed), scratch, _outputs())] @@ -204,7 +210,7 @@ def test_draft_moe_two_tokens(weights): got = block(x) from sgl_nvidia_mega.qwen3_5.sm103.draft import moe - scratch = moe.draft_moe_fused_scratch(E, K, I, torch.device("cuda")) + scratch = moe.draft_moe_fused_scratch(E, K, _inter(weights), torch.device("cuda")) for m in range(2): *_, out = _run(weights, x[m : m + 1], scratch, _outputs()) assert torch.equal(got[m : m + 1], out) @@ -213,7 +219,7 @@ def test_draft_moe_two_tokens(weights): def test_rejects_bad_operands(weights): from sgl_nvidia_mega.qwen3_5.sm103.draft import moe - scratch = moe.draft_moe_fused_scratch(E, K, I, torch.device("cuda")) + scratch = moe.draft_moe_fused_scratch(E, K, _inter(weights), torch.device("cuda")) with pytest.raises(ValueError): _run(weights, _x(0).float(), scratch, _outputs()) with pytest.raises(ValueError): diff --git a/test/kernels/test_qwen3_5_draft_xchg.py b/test/kernels/test_qwen3_5_draft_xchg.py index 400432c7..c1a07091 100644 --- a/test/kernels/test_qwen3_5_draft_xchg.py +++ b/test/kernels/test_qwen3_5_draft_xchg.py @@ -23,7 +23,7 @@ The bitwise comparison with the kernels of SGLang branch nvidia/qwen35-agentx-gb300 needs that branch installed as ``sglang`` and is not part of this test. -Run with ``python test/kernels/test_qwen3_5_draft_xchg.py --num-gpu 4``. +Run with ``python test/kernels/test_qwen3_5_draft_xchg.py --num-gpu 2`` (or 4 / 8). """ import os @@ -279,6 +279,50 @@ def same(got, want): dist.barrier(group=cpu_group) +def test_exchange_full_row_ids_on_tp2_sized_shards(ranks, boundary): + """A local winner above 65535 and the tie rule survive the exchange.""" + from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels + + rank, world, group, _ = ranks + mkd, _ = boundary + vocab = 124544 + # Token-only exchange never reads the candidate payload. + candidates = torch.empty(vocab, H, dtype=torch.float32, device="cuda") + hot_ids = torch.arange(world * vocab, dtype=torch.int64, device="cuda") + index = torch.empty(1, 1, dtype=torch.int64, device="cuda") + probability = torch.empty(1, 1, dtype=torch.float32, device="cuda") + key = kernels.NAN_KEY_SPLIT_ARGMAX + for ties in (False, True): + logits = torch.zeros(1, vocab, dtype=torch.bfloat16, device="cuda") + logits[0, vocab - 1] = 5.0 + if ties: + logits[0, 65536] = 5.0 + kernels.mkd_xchg( + logits, + candidates, + hot_ids, + mkd.own_partial, + mkd.epoch_done, + mkd.ext_cand, + mkd.ext_partials, + index, + probability, + None, + None, + None, + mkd.mc, + mkd.uc, + rank, + world, + 0, + key, + token_only=True, + scan_part=mkd.scan_part, + ) + want = _want_token(logits, group, -1e30) + assert int(index) == want and float(probability) == 1.0 + + @pytest.mark.parametrize("m", [1, 2, 7]) def test_split_fc_gathers_and_sums_in_rank_order(ranks, boundary, m): from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels @@ -705,4 +749,4 @@ def run(): if __name__ == "__main__": - multigpu_pytest_main(__name__, __file__, num_gpus=(4,)) + multigpu_pytest_main(__name__, __file__, num_gpus=(2, 4, 8)) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 9fa311cc..ace74d05 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -1511,6 +1511,38 @@ def test_the_draft_extend_writes_the_selected_rows_outputs(monkeypatch): # ------------------------------------------------------------------- the parity check +@pytest.mark.parametrize("requests", [2, 6]) +def test_multi_draft_extend_uses_split_fc_only_within_its_bound(monkeypatch, requests): + from sgl_nvidia_mega.qwen3_5.draft import layer + from sgl_nvidia_mega.qwen3_5.draft.extend import DraftExtend + + parts = _DraftParts() + calls = [] + parts.embed = lambda ids: calls.append("embed") or ids + parts.mkd = SimpleNamespace( + split_fc=lambda embeds, hidden: calls.append("split") or hidden + ) + monkeypatch.setattr(layer, "fc_input", lambda *a: calls.append("stock") or a[2]) + monkeypatch.setattr(layer, "layer", lambda p, x, *a, **kw: x) + monkeypatch.setattr(layer, "top1_into", lambda *a, **kw: None) + rows = requests * 7 + io = SimpleNamespace( + input_ids=torch.arange(rows), + hidden_states=torch.ones(rows, 1), + positions=None, + out_cache_loc=None, + attn=None, + kv=None, + tokens_per_request=7, + select_index=torch.arange(requests), + hidden_out=torch.empty(requests, 1), + topk_index_out=None, + topk_p_out=None, + ) + DraftExtend(parts, tokens_per_request=7)._multi_step(io) + assert calls == (["embed", "split"] if rows <= 16 else ["stock"]) + + @pytest.fixture def eager(monkeypatch): """Outside CUDA graph capture (a CPU build of torch has no capture to query), with From 1046f06b0dacdae653f8602b6a0d88780689eb8a Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 2 Nov 2026 00:46:52 -0800 Subject: [PATCH 49/62] perf(qwen3.5): select aggregate and PD slots with explicit topology plans --- README.md | 19 +- sgl_nvidia_mega/qwen3_5/draft/__init__.py | 1 + sgl_nvidia_mega/qwen3_5/plan.py | 18 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 122 +++++++---- sgl_nvidia_mega/qwen3_5/runtime/executor.py | 17 +- sgl_nvidia_mega/qwen3_5/runtime/parity.py | 29 ++- sgl_nvidia_mega/qwen3_5/spec.py | 73 +++++-- sgl_nvidia_mega/qwen3_5/topology.py | 216 ++++++++++++++++++++ test/unit/test_qwen3_5_runtime.py | 80 +++++++- test/unit/test_qwen3_5_spec.py | 61 +++++- test/unit/test_qwen3_5_steps.py | 56 +++-- 11 files changed, 586 insertions(+), 106 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/topology.py diff --git a/README.md b/README.md index 8ce32131..4a9b0093 100644 --- a/README.md +++ b/README.md @@ -18,8 +18,8 @@ The package loads as an SGLang [general plugin](https://github.com/sgl-project/s | Qwen3.5-397B-A17B (NVFP4) | GB300 (SM103) | 8 (two trays, one NVLink domain) | MTP target verify steps of 7 rows (batch 1) | | Kimi K3 (packed MXFP4) | GB300 (SM103) | 8 | One-token decode steps of 1–8, 16, 32 and 64 rows | -For DeepSeek-V4.1-Flash and Qwen3.5, larger verify steps, prefill and extend run -SGLang's standard path. Their draft specs provide the DSpark proposal and the +For DeepSeek-V4.1-Flash, larger verify steps, prefill and extend run +SGLang's standard path. Qwen3.5 also offers the explicit GDN prefill slot below. Their draft specs provide the DSpark proposal and the Qwen3.5 EAGLE entries, respectively. For Kimi K3, the plugin covers the 93-layer language model (24 MLA and 69 KDA layers) at the listed widths; every other call runs the standard path. @@ -55,7 +55,6 @@ Qwen3.5 (`qwen3_5/runtime`) runs the same way: per layer the stage boundary (`bo | Slot | Rows | Kernel | | --- | --- | --- | | `route_topk`, `mkm_m32` … `mkm_m512` | 1–896 | Materialized stock routing and the shape-selected fused wide routed-expert engines | -| `gdn_verify`, `gdn_pdl`, `gdn_conv` | 7–896 | The opened GatedDeltaNet verify chain, dependent launch, and bitwise conv update | | `hop`, `hop_fp8` | 1–15 | The layer boundary's fused all-reduce + residual add + RMSNorm, also writing the next static-FP8 input | | `multi_request` | 2–128 requests | The target graphs keep every offered request count instead of padding to a coarser width | | `draft_loop`, `draft_extend`, `draft_multi`, `draft_embed` | 1–128 requests | EAGLE draft entries, exact sharded top-1 and embedding exchange; the pre-FC norms write directly into the FC input | @@ -65,6 +64,7 @@ The remaining kernel families are opt-in through that one slot policy. These cha | Slot | Rows | Kernel | | --- | --- | --- | | `route`, `route_gemm`, `mkm`, `mkm_wide` | 1–32 | Fused router fronts and the small routed-expert engines | +| `gdn_verify`, `gdn_pdl`, `gdn_conv` | 7–896 | Explicit numeric verify chain and dependent variants; FP8 norm uses the v2 relative-L2 bound of 2e-2 | | `gdn_core` | 7–224 | ReplaySSM verify, gated norm and FP8 output in one launch | | `hop_mid` (with `hop`) | 16–28 | Extends the one-shot boundary mailbox through four requests instead of returning those widths to the stock boundary | | `gdn_push` (with `gdn_verify`, `hop`) | 7 | The GatedDeltaNet `out_proj` as the hop's GEMM, whose epilogue pushes the partial into the hop's mailbox | @@ -79,7 +79,18 @@ includes `hop`. It remains opt-in rather than part of the bitwise default polic an explicit `SGLANG_NVIDIA_MEGAKERNEL_SLOTS` value replaces, rather than extends, that default list. -The sampling and acceptance after the verify, and the GatedDeltaNet extend of prefill, stay SGLang's own code. In an SGLang that has the switches, `SGLANG_ENABLE_EAGLE_POST_VERIFY_CUDA_GRAPH=1` runs the sampling and acceptance in CUDA graphs, and `SGLANG_ENABLE_GDN_PREFILL_GRAPH_EXTEND=1` captures the GatedDeltaNet extend inside the breakable prefill CUDA graphs. +The router-only `route`, `route_topk` and `route_gemm` slots support TP2, TP4, TP8 and TEP8 with the TRTLLM runner and no all-to-all. PD decode TP2 uses the same local row counts: batch 6/10 gives 42/70 verify rows. These fronts leave the expert GEMMs with the stock runner, so TP2 needs no fused expert engine. `route_shared` keeps its TP8/TEP8 projection bounds; DEP8 keeps stock MoE dispatch and routing. + +`mka` uses the fused projection kernels through 14 verify rows on TP8 and 8 on +TP4. At wider rows, and at every TP2 width, it keeps SGLang's projection GEMMs +and fuses Q/K normalization, RoPE and the FP8 KV write, then the output gate +and static FP8 quantization. TP2 has 16 query heads and one KV head per rank; +TP4 has 8/1 and TP8 has 4/1. The same front path serves PD decode TP2 and DEP8 (32/2 per rank). +`multi_request` coverage is required for more than one request. `skv2` keeps +its one-request, one-KV-head scope; `skv2_cp` additionally requires replicated +KV heads across ranks. + +The sampling and acceptance after the verify stay SGLang's own code. `gdn_prefill_split_conv` is an explicit bitwise slot for ordinary text prefill on TP1/PP4 with the data-only prefill binding: it splits the native conv by Q/K/V channels above 8192 channels, retaining native recurrence, norm and state writes. Unsupported layouts use SGLang; an empty slot policy calls the original prefill forward directly. In an SGLang that has the switches, `SGLANG_ENABLE_EAGLE_POST_VERIFY_CUDA_GRAPH=1` runs the sampling and acceptance in CUDA graphs, and `SGLANG_ENABLE_GDN_PREFILL_GRAPH_EXTEND=1` captures the GatedDeltaNet extend inside the breakable prefill CUDA graphs. ### Design diff --git a/sgl_nvidia_mega/qwen3_5/draft/__init__.py b/sgl_nvidia_mega/qwen3_5/draft/__init__.py index 13473c54..f807735a 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/draft/__init__.py @@ -221,6 +221,7 @@ def _build(model, hf_config, ctx) -> MtpDraft: parallel = get_parallel() runtime.apply_expert_parallel(parallel.moe_ep_size) runtime.apply_cp_lite_topology(parallel.attn_tp_size) + runtime.apply_topology(parallel) runtime.build_cp_lite(ctx.binding) return MtpDraft( model, diff --git a/sgl_nvidia_mega/qwen3_5/plan.py b/sgl_nvidia_mega/qwen3_5/plan.py index f6addea8..373a86ea 100644 --- a/sgl_nvidia_mega/qwen3_5/plan.py +++ b/sgl_nvidia_mega/qwen3_5/plan.py @@ -27,8 +27,8 @@ expert's gate/up in route's launch) up to 8 rows. ``skv2`` runs the full attention's core of one request through the split-KV attention of SGLang branch nvidia/qwen35-agentx-gb300, ``mka`` the fused attention prologue and epilogue around it -(up to 8 rows; with four query heads a rank, its wide entries up to 14 rows and past -them its front path, bit for bit SGLang's, at every width), and ``mka_push`` (with +(TP8 fused through 14 rows, TP4 through 8, then its front path at every width; +TP2 always uses the front path, bit for bit SGLang's), and ``mka_push`` (with ``hop``) has the epilogue push its rows into the hop. ``lowm_in`` / ``lowm_out`` (numeric) run the static FP8 input / output projections on a low-M warp-MMA launch with the weights loaded before the dependency wait, up to each one's measured rows (``LOWM_*_MAX_ROWS``), and with ``lowm_push`` (with ``hop``) push their rows into the hop; @@ -78,7 +78,7 @@ # the rows slot route_shared's shared stage fits shared memory for ROUTE_SHARED_ROWS = 8 MKA_MAX_ROWS = 8 -# with four query heads a rank: the wide entries, then the front path (attention/mka.py) +# The fused entries where supported, then the front path (attention/mka.py). MKA_WIDE_MAX_ROWS = MAX_ROWS # slots lowm / lowm_se: the rows up to which each low-M launch measured faster than SGLang's quantize + # cuBLAS path on GB300 (results/trt-1022/l6): the input projections (in_proj_qkvz + in_proj_ba, qkv_proj), @@ -197,7 +197,7 @@ class Facts(msgspec.Struct, frozen=True): ) skv2: bool = False # SKV2 for every full-attention layer mka: bool = False # the fused prologue and epilogue for every full-attention layer - mka_wide: bool = False # with its wide entries (four query heads a rank) + mka_wide: bool = False # with its front path at every width mka_push: bool = False # with the epilogue's push into the hop route_lin: bool = False # the router-before-all-reduce path built for every layer lowm_in: bool = False # the low-M input projections built for every layer @@ -344,12 +344,18 @@ def widest(widths: tuple[int, ...], max_rows: int) -> int: def covered( - widths: tuple[int, ...], facts: Facts, tokens_per_request: int + widths: tuple[int, ...], + facts: Facts, + tokens_per_request: int, + *, + stock: bool = False, ) -> tuple[int, ...]: """The widths a build covers: one request always (with no slot it runs the stock step), more requests only where a slot runs.""" return tuple( w for w in widths - if w == tokens_per_request or slots(w, facts, tokens_per_request) != Slots() + if stock + or w == tokens_per_request + or slots(w, facts, tokens_per_request) != Slots() ) diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index c02f0d5d..82d427b6 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -23,9 +23,6 @@ "mkm_m128", "mkm_m256", "mkm_m512", - "gdn_verify", - "gdn_pdl", - "gdn_conv", "hop", "hop_fp8", "multi_request", @@ -35,6 +32,12 @@ "draft_embed", ) OPT_IN_SLOTS: tuple[str, ...] = ( + # The fused FP8 norm can select a different reduction layout from BF16. + # Keep its verify chain and dependent variants explicitly requested. + "gdn_verify", + "gdn_pdl", + "gdn_conv", + "gdn_prefill_split_conv", "route", "route_gemm", "mkm", @@ -191,8 +194,20 @@ def apply_expert_parallel(ep_size: int) -> frozenset: return dropped -def slot_reasons() -> list[str]: +def apply_topology(parallel) -> frozenset: + """Keep only each slot's declared capabilities before package builds run.""" + from sgl_nvidia_mega.qwen3_5 import topology + + global SLOTS + selected = topology.select(SLOTS, topology.from_parallel(parallel)) + dropped = SLOTS - selected + SLOTS = selected + return dropped + + +def slot_reasons(requested=None) -> list[str]: """Why the requested slots cannot be served ([] when they can).""" + slots = SLOTS if requested is None else requested known = tuple( dict.fromkeys( ( @@ -216,113 +231,113 @@ def slot_reasons() -> list[str]: "with post-verify graphs turns them on with " "SGLANG_ENABLE_EAGLE_POST_VERIFY_CUDA_GRAPH=1" ) - if "mkm" in SLOTS and "route" not in SLOTS: + if "mkm" in slots and "route" not in slots: reasons.append(f"{SLOTS_ENV}: slot 'mkm' runs on slot 'route''s outputs") - if "mkm_wide" in SLOTS and "mkm" not in SLOTS: + if "mkm_wide" in slots and "mkm" not in slots: reasons.append(f"{SLOTS_ENV}: slot 'mkm_wide' is slot 'mkm''s 16-token engine") - if "mkm_wide" in SLOTS and "multi_request" not in SLOTS: + if "mkm_wide" in slots and "multi_request" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'mkm_wide' serves the steps of more than one request " "('multi_request')" ) - if "mkm_m32" in SLOTS and "route_topk" not in SLOTS: + if "mkm_m32" in slots and "route_topk" not in slots: reasons.append(f"{SLOTS_ENV}: slot 'mkm_m32' takes slot 'route_topk''s rows") - if "mkm_m32" in SLOTS and "multi_request" not in SLOTS: + if "mkm_m32" in slots and "multi_request" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'mkm_m32' serves the steps of more than two requests " "('multi_request')" ) - if "mkm_m64" in SLOTS and "route_topk" not in SLOTS: + if "mkm_m64" in slots and "route_topk" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'mkm_m64' runs on slot 'route_topk''s routing" ) - if "mkm_m64" in SLOTS and "multi_request" not in SLOTS: + if "mkm_m64" in slots and "multi_request" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'mkm_m64' serves the steps of more than four requests " "('multi_request')" ) - if "mkm_m128" in SLOTS and "route_topk" not in SLOTS: + if "mkm_m128" in slots and "route_topk" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'mkm_m128' runs on slot 'route_topk''s routing" ) - if "mkm_m128" in SLOTS and "multi_request" not in SLOTS: + if "mkm_m128" in slots and "multi_request" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'mkm_m128' serves the steps of more than nine requests " "('multi_request')" ) for name, more in (("mkm_m256", "18"), ("mkm_m512", "36")): - if name in SLOTS and "route_topk" not in SLOTS: + if name in slots and "route_topk" not in slots: reasons.append( f"{SLOTS_ENV}: slot '{name}' runs on slot 'route_topk''s routing" ) - if name in SLOTS and "multi_request" not in SLOTS: + if name in slots and "multi_request" not in slots: reasons.append( f"{SLOTS_ENV}: slot '{name}' serves the steps of more than {more} requests " "('multi_request')" ) if ( - "mkm_wide_bal" in SLOTS + "mkm_wide_bal" in slots and not { "mkm_m256", "mkm_m512", } - & SLOTS + & slots ): reasons.append( f"{SLOTS_ENV}: slot 'mkm_wide_bal' changes the work split of slot " "'mkm_m256' or 'mkm_m512'" ) - if "route_shared" in SLOTS and "route" not in SLOTS: + if "route_shared" in slots and "route" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'route_shared' runs in slot 'route''s launch" ) - if "route_gemm" in SLOTS and "route_topk" not in SLOTS: + if "route_gemm" in slots and "route_topk" not in slots: reasons.append(f"{SLOTS_ENV}: slot 'route_gemm' runs inside slot 'route_topk'") - if "hop_mid" in SLOTS and "hop" not in SLOTS: + if "hop_mid" in slots and "hop" not in slots: reasons.append(f"{SLOTS_ENV}: slot 'hop_mid' widens slot 'hop'") - if "gdn_push" in SLOTS and not {"gdn_verify", "hop"} <= SLOTS: + if "gdn_push" in slots and not {"gdn_verify", "hop"} <= slots: reasons.append( f"{SLOTS_ENV}: slot 'gdn_push' pushes slot 'gdn_verify''s out_proj into " "slot 'hop'" ) - if "gdn_front" in SLOTS and not {"gdn_verify", "gdn_pdl"} <= SLOTS: + if "gdn_front" in slots and not {"gdn_verify", "gdn_pdl"} <= slots: reasons.append( f"{SLOTS_ENV}: slot 'gdn_front' runs slot 'gdn_pdl''s dependent verify" ) - if "gdn_trt_front" in SLOTS and not {"gdn_front", "gdn_core"} <= SLOTS: + if "gdn_trt_front" in slots and not {"gdn_front", "gdn_core"} <= slots: reasons.append( f"{SLOTS_ENV}: slot 'gdn_trt_front' replaces slot 'gdn_front' in front " "of slot 'gdn_core'" ) - if "gdn_core" in SLOTS and "gdn_verify" not in SLOTS: + if "gdn_core" in slots and "gdn_verify" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'gdn_core' runs slot 'gdn_verify''s step in one launch" ) - if "gdn_conv" in SLOTS and "gdn_verify" not in SLOTS: + if "gdn_conv" in slots and "gdn_verify" not in slots: reasons.append(f"{SLOTS_ENV}: slot 'gdn_conv' runs slot 'gdn_verify''s conv") - if "gdn_pdl" in SLOTS and "gdn_verify" not in SLOTS: + if "gdn_pdl" in slots and "gdn_verify" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'gdn_pdl' launches slot 'gdn_verify''s verify" ) - for name in sorted(SLOTS & set(DRAFT_SLOTS) - {"draft_skv2", "draft_skv2_cp"}): - if "draft_loop" not in SLOTS: + for name in sorted(slots & set(DRAFT_SLOTS) - {"draft_skv2", "draft_skv2_cp"}): + if "draft_loop" not in slots: reasons.append(f"{SLOTS_ENV}: slot {name!r} runs inside entry 'draft_loop'") - if "draft_hop" in SLOTS and "hop" not in SLOTS: + if "draft_hop" in slots and "hop" not in slots: reasons.append(f"{SLOTS_ENV}: slot 'draft_hop' runs on slot 'hop''s mailbox") - if "draft_fin" in SLOTS and "draft_hop" not in SLOTS: + if "draft_fin" in slots and "draft_hop" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'draft_fin' extends slot 'draft_hop''s fused final norm" ) - if "draft_multi" in SLOTS and "multi_request" not in SLOTS: + if "draft_multi" in slots and "multi_request" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'draft_multi' serves the draft steps of two or more " "requests ('multi_request')" ) - if "draft_embed" in SLOTS and "draft_multi" not in SLOTS: + if "draft_embed" in slots and "draft_multi" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'draft_embed' runs in slot 'draft_multi''s steps" ) - if "draft_skv2" in SLOTS and "draft_extend" not in SLOTS: + if "draft_skv2" in slots and "draft_extend" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'draft_skv2' runs inside entry 'draft_extend'" ) @@ -331,25 +346,25 @@ def slot_reasons() -> list[str]: ("draft_q1_cp", "draft_q1"), ("draft_skv2_cp", "draft_skv2"), ): - if name in SLOTS and base not in SLOTS: + if name in slots and base not in slots: reasons.append( f"{SLOTS_ENV}: slot {name!r} takes slot {base!r}'s long-KV graph " "variant's place from its length up" ) - if "hop_fp8" in SLOTS and "hop" not in SLOTS: + if "hop_fp8" in slots and "hop" not in slots: reasons.append(f"{SLOTS_ENV}: slot 'hop_fp8' is an output of slot 'hop'") - if "lowm_push" in SLOTS and not {"lowm_out", "hop"} <= SLOTS: + if "lowm_push" in slots and not {"lowm_out", "hop"} <= slots: reasons.append( f"{SLOTS_ENV}: slot 'lowm_push' pushes slot 'lowm_out''s rows into slot " "'hop''s mailbox" ) - if "mka_push" in SLOTS and not {"mka", "hop"} <= SLOTS: + if "mka_push" in slots and not {"mka", "hop"} <= slots: reasons.append( f"{SLOTS_ENV}: slot 'mka_push' pushes slot 'mka''s epilogue into slot " "'hop''s mailbox" ) if ( - "route_lin" in SLOTS + "route_lin" in slots and not { "route", "mkm", @@ -358,13 +373,13 @@ def slot_reasons() -> list[str]: "mka", "mka_push", } - <= SLOTS + <= slots ): reasons.append( f"{SLOTS_ENV}: slot 'route_lin' runs the router before the all-reduce " "through slots 'route', 'mkm', 'hop', 'gdn_push', 'mka' and 'mka_push'" ) - if "route_lin" in SLOTS and "route_shared" in SLOTS and "lowm_se" not in SLOTS: + if "route_lin" in slots and "route_shared" in slots and "lowm_se" not in slots: reasons.append( f"{SLOTS_ENV}: slot 'route_lin' needs 'lowm_se' to replace " "'route_shared' at one request" @@ -418,6 +433,11 @@ def build(model, ctx): from sgl_nvidia_mega.qwen3_5.runtime.executor import Qwen35Verify parallel = get_parallel() + if getattr(parallel, "pp_size", 1) == 4: + from sgl_nvidia_mega.qwen3_5.runtime.prefill import Qwen35Prefill + + return Qwen35Prefill(model, ctx, "gdn_prefill_split_conv" in SLOTS) + dropped = apply_expert_parallel(parallel.moe_ep_size) if dropped: logger.info( @@ -428,8 +448,18 @@ def build(model, ctx): ) apply_cp_lite_topology(parallel.attn_tp_size) build_cp_lite(ctx.binding) - multi = slot_plan.MULTI_REQUEST in SLOTS - if multi: + from sgl_nvidia_mega.qwen3_5 import topology + + current = topology.from_parallel(parallel) + unsupported = apply_topology(parallel) + if unsupported: + logger.info( + "NVIDIA megakernel: Qwen3.5 %s: slots %s use stock (topology)", + current.name, + sorted(unsupported), + ) + multi = current.stock_widths or slot_plan.MULTI_REQUEST in SLOTS + if multi and not current.stock_widths: multi = slot_plan.multi_request(parallel.tp_size, SLOTS) widths = slot_plan.candidates(ctx.widths, ctx.tokens_per_request, multi) if not widths: @@ -488,7 +518,9 @@ def build(model, ctx): lowm_push=bound["lowm"].push, hop_fp8=bound["lowm"].hop_fp8, ) - covered = slot_plan.covered(pctx.widths, facts, pctx.tokens_per_request) + covered = slot_plan.covered( + pctx.widths, facts, pctx.tokens_per_request, stock=current.stock_widths + ) slots = slot_plan.plan(covered, facts, pctx.tokens_per_request) # the slots of every width wider than one request, for the startup line wider = "".join( @@ -546,8 +578,8 @@ def build(model, ctx): vacuous.update({"route_topk", "route_gemm"}) missing = sorted(name for name in target_slots - vacuous if not built[name]) if missing: - raise RuntimeError( - f"NVIDIA megakernel: requested slots {missing} built no width (MoE: " + logger.info( + f"NVIDIA megakernel: requested slots {missing} use stock (MoE: " f"{bound['moe'].reason or 'built'}; GDN: {bound['gdn'].reason or 'built'}; " f"boundary: {bound['boundary'].reason or 'built'}; " f"attention: {bound['attention'].reason or 'built'}; " diff --git a/sgl_nvidia_mega/qwen3_5/runtime/executor.py b/sgl_nvidia_mega/qwen3_5/runtime/executor.py index 758edd2c..fa8b4ed0 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/executor.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/executor.py @@ -64,6 +64,9 @@ "model", ), "srt/model_executor/forward_batch_info.py": ("mrope_positions", "forward_mode"), + "srt/model_executor/forward_mode.py|srt/model_executor/forward_batch_info.py": ( + "is_idle", + ), "srt/models/qwen3_5.py": ( "Qwen3_5LinearDecoderLayer", "layers", @@ -85,7 +88,10 @@ # The SGLang functions whose body this module mirrors, by the digest of their code # (test/unit/test_contract.py recomputes them on the installed SGLang) SGLANG_MIRRORS = { - "srt/models/qwen3_vl.py:Qwen3VLForConditionalGeneration.forward": "d394d79ba7fa7b31", + "srt/models/qwen3_vl.py:Qwen3VLForConditionalGeneration.forward": ( + "d394d79ba7fa7b31", + "dbdcf4b88a2c0d40", + ), "srt/models/qwen3_5.py:Qwen3_5ForCausalLM.forward": "9e530b264ce30ad6", "srt/models/qwen3_5.py:Qwen3_5LinearDecoderLayer.forward": "2298a57111f9430e", "srt/models/qwen3_5.py:Qwen3_5AttentionDecoderLayer.forward": "6473f673a1a0796b", @@ -269,7 +275,14 @@ def _language_model(self, st, hidden): def _layer(self, i, layer, hidden, st, capture): """A decoder layer's forward: the mixer stage, then the MoE stage.""" x = boundary.pre(layer, hidden, st, capture) - part = self.mixers[i](layer, x, st) + # Idle DP ranks still cross the boundaries and run the MoE's collective, + # but the decoder layers skip their mixer (including empty FP4 GEMMs). + rows = (x[0] if isinstance(x, tuple) else x).shape[0] + part = ( + self.mixers[i](layer, x, st) + if not st.fb.forward_mode.is_idle() and rows > 0 + else x + ) x = boundary.post_attn(layer, part, st) with boundary.ffn_exit(layer, st) as ffn_exit: out = moe.step(layer, x, st, defer=ffn_exit.defer_moe_finalize) diff --git a/sgl_nvidia_mega/qwen3_5/runtime/parity.py b/sgl_nvidia_mega/qwen3_5/runtime/parity.py index 78958700..19993880 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/parity.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/parity.py @@ -105,11 +105,11 @@ def run(self, step, original, megakernel): if torch.cuda.is_current_stream_capturing(): counts = self.counters[step.rows] counts[0:1].add_(1) - valid = ( - None - if slots is None or step.rows % slots.numel() - else valid_rows(slots, step.rows) - ) + valid = None + if slots is not None and ( + slots.numel() == 0 or step.rows % slots.numel() == 0 + ): + valid = valid_rows(slots, step.rows) for i, (a, b) in enumerate(pairs, start=1): counts[i : i + 1].add_(differing(a, b)) counts[i + 2 : i + 3].add_(0 if a is None else a.numel()) @@ -232,16 +232,19 @@ def distance(logits_a, logits_b, hidden_a, hidden_b): logits.""" if logits_a is None or logits_b is None: return None + hidden = ( + torch.zeros((), device=logits_a.device) + if hidden_a is None or hidden_a.numel() == 0 + else (hidden_a.float() - hidden_b.float()).abs().max() + ) + if logits_a.numel() == 0: + zero = torch.zeros((), device=logits_a.device) + return torch.stack([zero, zero, zero, zero, hidden]).double() top = logits_a.argmax(-1, keepdim=True) dlp = ( torch.log_softmax(logits_b.float(), -1).gather(-1, top) - torch.log_softmax(logits_a.float(), -1).gather(-1, top) ).abs() - hidden = ( - torch.zeros((), device=logits_a.device) - if hidden_a is None - else (hidden_a.float() - hidden_b.float()).abs().max() - ) # no host-to-device copy: this also runs inside graph capture return torch.stack( [ @@ -268,6 +271,8 @@ def _distance_text(apart) -> str: def valid_rows(slots: torch.Tensor, rows: int) -> torch.Tensor: """[rows] bool: the rows of requests with a mamba slot (not the pad slot), for a step of ``slots.numel()`` requests of equal rows.""" + if slots.numel() == 0: + return torch.zeros(rows, dtype=torch.bool, device=slots.device) return (slots >= 0).repeat_interleave(rows // slots.numel()) @@ -284,7 +289,7 @@ def differing(a, b, valid=None): ) bits = _BITS[a.element_size()] differ = a.contiguous().view(bits) != b.contiguous().view(bits) - if valid is not None: + if valid is not None and differ.numel() > 0: differ = differ & valid.view(-1, *([1] * (differ.dim() - 1))) return differ.sum() @@ -296,4 +301,6 @@ def _clone(t): def _max_abs(a, b) -> str: if a is None: return "n/a" + if a.numel() == 0: + return "0" return f"{(a.float() - b.float()).abs().max().item():.3g}" diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index 580fe4b3..a84c3d6f 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -27,7 +27,7 @@ Requirements, ) -from sgl_nvidia_mega.qwen3_5 import TEXT_CONFIG, runtime +from sgl_nvidia_mega.qwen3_5 import TEXT_CONFIG, runtime, topology from sgl_nvidia_mega.qwen3_5.plan import MAX_ROWS # The SGLang names this module reads (test/unit/test_contract.py pins them) @@ -36,6 +36,7 @@ "Feature", "SPECULATIVE_DECODING", "MULTI_NODE_TP", + "DP_ATTENTION", "KV_OFFLOAD", "PD_DISAGGREGATION", "MegakernelSpec", @@ -66,11 +67,14 @@ "attn_dp_size", "startup_weight_load_mode", "disaggregation_mode", + "pp_size", + "ep_size", ), "srt/runtime_context.py": ("derive_attn_tp_size",), "srt/configs/qwen3_5.py": ("text_config",), "srt/kernel_plugins/nvidia/eagle_binding.py|srt/megakernels/eagle_binding.py": ( "EagleForwardView", + "LinearPrefillView", ), } @@ -81,7 +85,7 @@ REQUIREMENTS = Requirements( isas=frozenset({"sm_103"}), # TP8 over the two 4-GPU trays of one NVL72 rack, TP4 on one tray - tp_sizes=frozenset({4, 8}), + tp_sizes=frozenset({1, 2, 4, 8}), dtypes=frozenset({"bfloat16"}), # the decode CUDA graphs the covered step is tested with (eager steps too) decode_graph_backends=frozenset({"full", "disabled"}), @@ -92,17 +96,17 @@ "speculative_eagle_topk": frozenset({1}), # SGLang's EAGLE entries serve a chain without rejection sampling "speculative_use_rejection_sampling": frozenset({False}), - "moe_a2a_backend": frozenset({"none"}), + "moe_a2a_backend": frozenset({"none", "flashinfer"}), # the deferred MoE finalize the boundary takes (the measured TP8 setup) - "moe_runner_backend": frozenset({"flashinfer_trtllm"}), - "flashinfer_allreduce_fusion_backend": frozenset({"cutedsl"}), + "moe_runner_backend": frozenset({"flashinfer_trtllm", "flashinfer_cutedsl"}), + "flashinfer_allreduce_fusion_backend": frozenset({None, "cutedsl"}), "attention_backend": frozenset({"trtllm_mha"}), "decode_attention_backend": frozenset({None, "trtllm_mha"}), # the Triton GDN verify on the ReplaySSM ring with BF16 checkpoints "linear_attn_backend": frozenset({"triton"}), "linear_attn_decode_backend": frozenset({None, "triton"}), "linear_attn_verify_backend": frozenset({None, "triton"}), - "enable_linear_replayssm_spec": frozenset({True}), + "enable_linear_replayssm_spec": frozenset({False, True}), "mamba_ssm_dtype": frozenset({"bfloat16"}), # the kernels read the weights SGLang loaded, after its last conversion "startup_weight_load_mode": frozenset({"serial"}), @@ -122,6 +126,9 @@ "sharded_gumbel": { "sharded_gumbel.categorical.chi_square_p_value_min": 1e-4, }, + # Same FP8 norm bound as v2 gdn_core; output-store layout can change the + # FP32 reduction and the final BF16 rounding at a midpoint. + "gdn_verify": {"gdn_verify.norm_fp8.rel_l2_vs_stock": 2e-2}, # Four router GEMV chains sum in another order than SGLang's cuBLAS GEMM. "route": {"route.router_logits.bf16_ulp_vs_stock": 1.0}, # The tiled router front stays within one BF16 ULP of SGLang's whole-M GEMM. @@ -235,12 +242,30 @@ def binding_reasons() -> list[str]: def pd_reasons(cfg) -> list[str]: - """Why this PD role has no covered step ([] on the decode side).""" - if (cfg.disaggregation_mode or "null") == "prefill": - return [ - "a PD prefill server runs no covered step; enable it on the decode side" - ] - return [] + """The P role is TP1/PP4; decode retains its established TP/PP contract.""" + if (cfg.disaggregation_mode or "null") != "prefill": + reasons = [] + if cfg.tp_size not in (2, 4, 8): + reasons.append(f"tp_size {cfg.tp_size} is not a supported decode topology") + if cfg.pp_size != 1: + reasons.append( + "pipeline parallelism is in use outside the supported prefill role" + ) + if not cfg.enable_linear_replayssm_spec: + reasons.append("enable_linear_replayssm_spec must be True for decode") + return reasons + reasons = [] + if (cfg.tp_size, cfg.pp_size, cfg.ep_size or 1) != (1, 4, 1): + reasons.append("PD prefill requires TP1/PP4 without expert parallelism") + if cfg.enable_linear_replayssm_spec: + reasons.append("PD prefill does not use the speculative ReplaySSM ring") + try: + from sglang.srt.kernel_plugins.nvidia.eagle_binding import ( # noqa: F401 + LinearPrefillView, + ) + except ImportError: + reasons.append("PD prefill requires the data-only LinearPrefillView binding") + return reasons def _config_disable_reasons(cfg, hf_config) -> list[str]: @@ -254,10 +279,28 @@ def _config_disable_reasons(cfg, hf_config) -> list[str]: attn_dp_size=cfg.attn_dp_size, ) ) + configured = topology.from_config(cfg) + topology_reasons = ( + ["Qwen3.5 DP attention requires DEP8 (TP8, attention DP8, MoE EP8)"] + if configured.attn_dp_size > 1 and configured.name != "dep8" + else [] + ) + if configured.name != "dep8" and cfg.disaggregation_mode != "prefill": + for name, want in ( + ("moe_a2a_backend", "none"), + ("moe_runner_backend", "flashinfer_trtllm"), + ("flashinfer_allreduce_fusion_backend", "cutedsl"), + ): + value = getattr(cfg, name) + if value != want: + topology_reasons.append(f"{name} {value!r} is not {want!r}") return ( - text_config_reasons(hf_config) + topology_reasons + + text_config_reasons(hf_config) + binding_reasons() - + runtime.slot_reasons() + + runtime.slot_reasons( + topology.select(runtime.SLOTS, topology.from_config(cfg)) + ) + pd_reasons(cfg) ) @@ -274,12 +317,14 @@ def _build(model, hf_config, ctx): features=frozenset( { Feature.SPECULATIVE_DECODING, + Feature.PIPELINE_PARALLEL, # TP8 spans the two 4-GPU trays of one NVL72 rack Feature.MULTI_NODE_TP, # HiCache: covered steps read only running requests' pages Feature.KV_OFFLOAD, # TEP (MoE TP x EP): the slots in runtime.EP_STOCK_SLOTS stay off Feature.EXPERT_PARALLEL, + Feature.DP_ATTENTION, # The decode side of PD; pd_reasons refuses the prefill side. Feature.PD_DISAGGREGATION, } diff --git a/sgl_nvidia_mega/qwen3_5/topology.py b/sgl_nvidia_mega/qwen3_5/topology.py new file mode 100644 index 00000000..f4d96e9a --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/topology.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Slot capabilities, resolved once before any package allocates or compiles. + +Widths are local target-verify rows (requests times tokens_per_request). DEP8's +attention is DP8/TP1 and its MoE is EP8; it must not inherit TP8/TEP8 kernels just +because its global TP group also contains eight ranks. +""" + +from __future__ import annotations + +import msgspec + +SGLANG_READS = { + "srt/runtime_context.py": ('"tp_size"', '"attn_dp_size"', '"moe_ep_size"'), + "srt/arg_groups/fields/*.py": ( + "tp_size", + "attn_dp_size", + "ep_size", + "dp_size", + "enable_dp_attention", + ), +} + +TOPOLOGIES = frozenset({"tp2", "tp4", "tp8", "tep8", "dep8"}) +LEGACY = frozenset({"tp8", "tep8"}) +PREFILL_TOPOLOGIES = frozenset({"p_tp1_pp4"}) + +# Extend a slot's entry only alongside its topology-specific build/shape checks +# and correctness evidence. A missing entry supports no topology. +SLOT_TOPOLOGIES = { + name: LEGACY + for name in ( + "route", + "route_topk", + "route_gemm", + "mkm", + "mkm_wide", + "mkm_m32", + "mkm_m64", + "mkm_m128", + "mkm_m256", + "mkm_m512", + "mkm_wide_bal", + "gdn_verify", + "gdn_pdl", + "gdn_conv", + "gdn_core", + "gdn_push", + "gdn_front", + "gdn_trt_front", + "hop", + "hop_mid", + "hop_fp8", + "route_shared", + "route_lin", + "skv2", + "skv2_cp", + "mka", + "mka_push", + "lowm_in", + "lowm_out", + "lowm_se", + "lowm_push", + "draft_loop", + "draft_extend", + "draft_moe", + "draft_q1", + "draft_xchg", + "draft_hop", + "draft_fin", + "draft_front", + "draft_skv2", + "draft_multi", + "draft_embed", + "draft_q1_cp", + "draft_skv2_cp", + ) +} +# GDN's opened chain uses the layer's local head counts; the Triton verify +# supports the aggregate 16-entry ring and DEP8's 32-entry ring. The numeric +# CUDA core retains its measured two-tile (16-entry) history implementation. +SLOT_TOPOLOGIES.update( + gdn_verify=TOPOLOGIES, + gdn_pdl=TOPOLOGIES, + gdn_conv=TOPOLOGIES, + gdn_core=frozenset({"tp2", "tp4", "tp8", "tep8"}), +) +SLOT_TOPOLOGIES["gdn_prefill_split_conv"] = PREFILL_TOPOLOGIES +# a9212216c6 already serves these TP4 shapes. Preserve that baseline support: +# MoE's small engine accepts I=256, GDN uses the layer's local head counts, +# the hop serves world 4, attention serves 8 Q / 1 KV heads, and the draft's +# exchange/embedding kernels serve world 4. Each package keeps its shape checks. +# The I=128-only wide MoE engines, TP8 GDN fronts and TP8 router-lin stay excluded. +TP4_V2_SLOTS = frozenset( + { + "route", + "route_topk", + "route_gemm", + "mkm", + "gdn_verify", + "gdn_pdl", + "gdn_conv", + "gdn_core", + "gdn_push", + "hop", + "hop_mid", + "hop_fp8", + "route_shared", + "skv2", + "skv2_cp", + "mka", + "mka_push", + "lowm_in", + "lowm_out", + "lowm_se", + "lowm_push", + "draft_loop", + "draft_extend", + "draft_moe", + "draft_q1", + "draft_xchg", + "draft_hop", + "draft_fin", + "draft_front", + "draft_skv2", + "draft_multi", + "draft_embed", + "draft_q1_cp", + "draft_skv2_cp", + } +) +for name in TP4_V2_SLOTS: + SLOT_TOPOLOGIES[name] = SLOT_TOPOLOGIES[name] | {"tp4"} + +# These names select coverage or SGLang's own runtime policy, not a plugin kernel. +SLOT_TOPOLOGIES.update( + multi_request=TOPOLOGIES, + sharded_gumbel=TOPOLOGIES, +) +# Full-step AB+BA prefers v2's stock draft on TP2. Keep every draft entry +# filtered there, including PD decode TP2; TP4/TP8 retain the BF16 entries. +for _name in ( + "draft_loop", + "draft_extend", + "draft_multi", + "draft_embed", + "draft_front", +): + SLOT_TOPOLOGIES[_name] = frozenset({"tp4", "tp8", "tep8"}) +# TP4/TP8 retain the fused vocabulary exchange. +SLOT_TOPOLOGIES["draft_xchg"] = frozenset({"tp4", "tp8", "tep8"}) +SLOT_TOPOLOGIES["draft_moe"] = frozenset({"tp4", "tp8"}) +SLOT_TOPOLOGIES["draft_q1"] = frozenset({"tp4", "tp8", "tep8"}) +# Router-only fronts leave expert GEMMs with the stock TRTLLM runner. Their +# 4096-hidden / 512-expert routing does not depend on TP's intermediate width. +SLOT_TOPOLOGIES.update( + route=frozenset({"tp2", "tp4", "tp8", "tep8"}), + route_topk=frozenset({"tp2", "tp4", "tp8", "tep8"}), + route_gemm=frozenset({"tp2", "tp4", "tp8", "tep8"}), +) +# ATTN: the generic front uses stock projection GEMMs on TP2 and DEP8. +SLOT_TOPOLOGIES.update( + mka=TOPOLOGIES, + skv2=frozenset({"tp4", "tp8", "tep8"}), + skv2_cp=frozenset({"tp4", "tp8", "tep8"}), +) + + +class Topology(msgspec.Struct, frozen=True): + tp_size: int + attn_dp_size: int = 1 + moe_ep_size: int = 1 + + @property + def name(self) -> str: + if self.attn_dp_size > 1: + return ( + "dep8" + if (self.tp_size, self.attn_dp_size, self.moe_ep_size) == (8, 8, 8) + else "unsupported" + ) + if self.tp_size == 8 and self.moe_ep_size > 1: + return "tep8" + return f"tp{self.tp_size}" + + @property + def stock_widths(self) -> bool: + """New topologies cover all offered widths even with every slot off.""" + return self.name in {"tp2", "tp4", "dep8"} + + +def from_parallel(parallel) -> Topology: + return Topology( + parallel.tp_size, + getattr(parallel, "attn_dp_size", 1), + parallel.moe_ep_size, + ) + + +def from_config(cfg) -> Topology: + # Startup guards may run before legacy --dp-size/--enable-dp-attention is + # normalized to attn_dp_size; the published ParallelContext is already resolved. + attention_dp = cfg.attn_dp_size or 1 + if cfg.enable_dp_attention: + attention_dp = max(attention_dp, cfg.dp_size or 1) + return Topology(cfg.tp_size, attention_dp, cfg.ep_size or 1) + + +def select(requested: frozenset, topology: Topology) -> frozenset: + """Only explicitly declared support enables a kernel on this topology.""" + return frozenset( + name + for name in requested + if topology.name in SLOT_TOPOLOGIES.get(name, frozenset()) + ) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index ace74d05..609de1a5 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -61,9 +61,6 @@ def test_the_default_policy_ships_only_bitwise_shape_bounded_paths(monkeypatch): "mkm_m128", "mkm_m256", "mkm_m512", - "gdn_verify", - "gdn_pdl", - "gdn_conv", "hop", "hop_fp8", "multi_request", @@ -173,9 +170,17 @@ def test_config_filter_derives_attention_tp_before_publication(monkeypatch): ) monkeypatch.setattr(target, "text_config_reasons", lambda hf_config: []) monkeypatch.setattr(target, "binding_reasons", lambda: []) - monkeypatch.setattr(runtime, "slot_reasons", lambda: []) + monkeypatch.setattr(runtime, "slot_reasons", lambda requested=None: []) monkeypatch.setattr(target, "pd_reasons", lambda cfg: []) - cfg = SimpleNamespace(tp_size=8, attn_cp_size=1, attn_dp_size=8, ep_size=8) + cfg = SimpleNamespace( + tp_size=8, + attn_cp_size=1, + attn_dp_size=8, + ep_size=8, + enable_dp_attention=True, + dp_size=8, + disaggregation_mode="decode", + ) assert not hasattr(cfg, "attn_tp_size") assert target._config_disable_reasons(cfg, SimpleNamespace()) == [] @@ -598,7 +603,9 @@ def test_the_covered_widths_are_the_ones_the_build_returns(): assert built.plan(torch.zeros(21, dtype=torch.int64), _batch()) is None -@pytest.mark.parametrize("tp_size, want", [(8, (7,)), (4, (7,))]) +@pytest.mark.parametrize( + "tp_size, want", [(8, (7,)), (4, (7, 14, 21, 28)), (2, (7, 14, 21, 28))] +) def test_the_build_covers_more_requests_only_where_a_slot_runs( monkeypatch, tp_size, want ): @@ -618,7 +625,7 @@ def test_the_build_covers_more_requests_only_where_a_slot_runs( binding="EagleBinding", ) built = runtime.build(model, ctx) - assert built.widths == want and built.summary() == "slots none (stock)" + assert built.widths == want and built.summary().startswith("slots none (stock)") def test_the_build_hands_over_the_binding_and_ships_the_bitwise_draft_entries(): @@ -1728,6 +1735,65 @@ def test_the_step_state_carries_the_binding(): ) +@pytest.mark.parametrize( + "shape, name, on, off", + [ + ((2,), "tp2", {"gdn_core", "mka", "route_topk"}, {"gdn_push", "draft_loop"}), + ((4,), "tp4", {"gdn_push", "skv2", "draft_loop"}, {"gdn_front", "mkm_m32"}), + ((8,), "tp8", {"gdn_front", "route_lin", "mkm_m32", "draft_moe"}, set()), + ((8, 1, 8), "tep8", {"gdn_front", "draft_loop"}, {"draft_moe"}), + ((8, 8, 8), "dep8", {"gdn_verify", "mka"}, {"gdn_core", "route", "draft_loop"}), + ], +) +def test_each_topology_selects_only_its_declared_slots(shape, name, on, off): + from sgl_nvidia_mega.qwen3_5 import topology + + topo = topology.Topology(*shape) + got = topology.select(frozenset(topology.SLOT_TOPOLOGIES), topo) + assert topo.name == name and on <= got and not off & got + assert not topology.select(frozenset({"undeclared"}), topo) + + +def test_every_slot_declares_its_topologies(): + from sgl_nvidia_mega.qwen3_5 import topology + + known = set( + ( + *runtime.SHIPPED_SLOTS, + *runtime.OPT_IN_SLOTS, + *runtime.OPT_IN_VARIANTS, + *runtime.DRAFT_ENTRIES, + *runtime.DRAFT_SLOTS, + *runtime.OPT_IN_COVERAGE, + *runtime.HOOK_SLOTS, + ) + ) + assert set(topology.SLOT_TOPOLOGIES) == known + + +def test_unsupported_specialization_does_not_disable_a_supported_slot(monkeypatch): + from sgl_nvidia_mega.qwen3_5 import topology + + monkeypatch.setitem(topology.SLOT_TOPOLOGIES, "gdn_verify", topology.TOPOLOGIES) + requested = frozenset({"gdn_verify", "gdn_front", "gdn_pdl", "multi_request"}) + selected = topology.select(requested, topology.Topology(2)) + assert selected == frozenset({"gdn_verify", "gdn_pdl", "multi_request"}) + assert runtime.slot_reasons(selected) == [] + + +def test_tp2_keeps_the_draft_stock(monkeypatch): + from sgl_nvidia_mega.qwen3_5 import draft, topology + + monkeypatch.setattr(runtime, "build_cp_lite", lambda binding: None) + monkeypatch.setattr(runtime, "SLOTS", frozenset(topology.SLOT_TOPOLOGIES)) + monkeypatch.setattr( + "sglang.srt.runtime_context.get_parallel", + lambda: SimpleNamespace(tp_size=2, attn_tp_size=2, moe_ep_size=1), + ) + built = draft.SPEC.build("lm", SimpleNamespace(context_len=8192), _draft_ctx()) + assert not built.slots and built.draft_loop is None and built.draft_extend is None + + @pytest.mark.parametrize("enabled", [False, True]) @pytest.mark.parametrize("mm_inputs", [None, [None]]) def test_prefill_plan_uses_actual_chunk_rows_and_stock_when_slot_off( diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 133841b1..db4cb043 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -30,6 +30,7 @@ @pytest.fixture(autouse=True) def registered(monkeypatch): monkeypatch.setattr(registry, "_REGISTRY", {}) + monkeypatch.setattr(runtime, "SLOTS", runtime.SLOTS) monkeypatch.delenv(plugin.PARITY_SORT_ENV, raising=False) plugin.register() @@ -89,6 +90,28 @@ def test_the_decode_side_of_pd_is_accepted(): assert _refusal(disaggregation_mode="decode") == "" +def test_an_unsupported_prefill_topology_is_refused(): + message = _refusal( + disaggregation_mode="prefill", enable_linear_replayssm_spec=False + ) + assert "PD prefill requires TP1/PP4" in message, message + + +def test_prefill_tp1_pp4_is_admitted_with_the_data_only_binding(): + assert ( + _refusal( + tp_size=1, + nnodes=1, + pp_size=4, + disaggregation_mode="prefill", + enable_linear_replayssm_spec=False, + moe_runner_backend="flashinfer_cutedsl", + flashinfer_allreduce_fusion_backend=None, + ) + == "" + ) + + BREAKABLE = with_phase( default_cuda_graph_config(), Phase.DECODE, backend=Backend.BREAKABLE ) @@ -97,9 +120,9 @@ def test_the_decode_side_of_pd_is_accepted(): @pytest.mark.parametrize( "kw, want", [ - (dict(tp_size=2, nnodes=1), "tp_size 2"), + (dict(tp_size=1, nnodes=1), "tp_size 1"), (dict(pp_size=2), "pipeline parallelism is in use"), - (dict(enable_dp_attention=True, dp_size=8), "DP attention is in use"), + (dict(enable_dp_attention=True, dp_size=4), "DP attention requires DEP8"), (dict(kv_cache_dtype="auto"), "kv_cache_dtype 'auto'"), (dict(speculative_algorithm=None), "speculative_algorithm None"), (dict(speculative_eagle_topk=4), "speculative_eagle_topk 4"), @@ -220,7 +243,7 @@ def test_a_slot_set_that_cannot_run_is_refused_before_the_load( def test_every_reason_is_listed_at_once(): - message = _refusal(tp_size=2, nnodes=1, kv_cache_dtype="auto", pp_size=2) + message = _refusal(tp_size=1, nnodes=1, kv_cache_dtype="auto", pp_size=2) assert message.count("\n - ") == 3, message @@ -577,7 +600,10 @@ def test_a_numeric_slot_turns_the_claim_into_its_tolerances(): assert dict(tiled_route.tolerances) == spec.NUMERIC_TOLERANCES["route_gemm"] claim = spec.numerics(frozenset({"gdn_verify", "hop", "gdn_push"})) assert claim.kind is NumericsKind.TOLERANCE - assert dict(claim.tolerances) == spec.NUMERIC_TOLERANCES["gdn_push"] + assert dict(claim.tolerances) == { + **spec.NUMERIC_TOLERANCES["gdn_push"], + **spec.NUMERIC_TOLERANCES["gdn_verify"], + } shared = spec.numerics(frozenset({"route", "route_shared"})) assert dict(shared.tolerances) == { **spec.NUMERIC_TOLERANCES["route"], @@ -585,11 +611,15 @@ def test_a_numeric_slot_turns_the_claim_into_its_tolerances(): } both = spec.numerics(frozenset({"gdn_verify", "gdn_pdl", "gdn_front", "gdn_push"})) assert dict(both.tolerances) == { + **spec.NUMERIC_TOLERANCES["gdn_verify"], **spec.NUMERIC_TOLERANCES["gdn_push"], **spec.NUMERIC_TOLERANCES["gdn_front"], } core = spec.numerics(frozenset({"gdn_verify", "gdn_core"})) - assert dict(core.tolerances) == spec.NUMERIC_TOLERANCES["gdn_core"] + assert dict(core.tolerances) == { + **spec.NUMERIC_TOLERANCES["gdn_core"], + **spec.NUMERIC_TOLERANCES["gdn_verify"], + } attention = spec.numerics(frozenset({"skv2", "mka", "mka_push", "hop"})) assert attention.kind is NumericsKind.TOLERANCE assert dict(attention.tolerances) == { @@ -628,3 +658,24 @@ def test_the_shipped_policy_claims_bitwise_numerics(): shipped = frozenset(runtime.SHIPPED_SLOTS) assert spec.numerics(shipped).kind is NumericsKind.BITWISE assert draft.numerics(shipped).kind is NumericsKind.BITWISE + + +@pytest.mark.parametrize("tp, pd", [(2, "decode"), (4, "null")]) +def test_topology_admission_is_driven_by_plugin_capabilities(tp, pd): + assert _refusal(tp_size=tp, nnodes=1, disaggregation_mode=pd) == "" + + +def test_dep8_accepts_the_recipe_stock_moe_and_dp_lm_head(): + assert ( + _refusal( + disaggregation_mode="decode", + enable_dp_attention=True, + dp_size=8, + ep_size=8, + enable_dp_lm_head=True, + moe_runner_backend="flashinfer_cutedsl", + moe_a2a_backend="flashinfer", + flashinfer_allreduce_fusion_backend=None, + ) + == "" + ) diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 1a5bc6e6..b55cce5a 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -158,7 +158,7 @@ def _moe_block(log, x, *, alt_stream=True, shared=True, deferred_ok=True): mlp = qwen2_moe.Qwen2MoeSparseMoeBlock.__new__(qwen2_moe.Qwen2MoeSparseMoeBlock) torch.nn.Module.__init__(mlp) bypassed = SimpleNamespace(format="bypassed") - routed = torch.zeros(7, 4) + routed = torch.zeros(x.shape[0], 4) log.known["routed"] = routed experts = _Callable( log.op("experts", routed), @@ -172,7 +172,10 @@ def _moe_block(log, x, *, alt_stream=True, shared=True, deferred_ok=True): _use_mega_moe=False, layer_id=3, gate=_Callable(log.op("gate", ("logits", None))), - topk=_Callable(log.op("topk", bypassed)), + topk=_Callable( + log.op("topk", bypassed), + empty_topk_output=log.op("empty_topk", bypassed), + ), experts=experts, shared_expert=_Callable(log.op("shared", "shared_out")) if shared else None, shared_expert_gate=_Callable( @@ -246,6 +249,25 @@ def test_the_moe_step_issues_what_sglang_issues(capture, defer, alt_stream): assert want == got if not isinstance(want, torch.Tensor) else torch.equal(want, got) +@pytest.mark.parametrize("capture", [False, True]) +def test_empty_dp_rank_skips_gate_and_shared_but_runs_experts(capture): + runs = [] + for which in ("sglang", "plugin"): + x = torch.zeros(0, 4) + log = _log(x=x) + mlp = _moe_block(log, x) + with _streams(log, capture), _moe_kernels(log): + if which == "sglang": + out = mlp.forward(x, None) + else: + out = moe.step(SimpleNamespace(mlp=mlp), x, _st(), defer=False) + runs.append((log, out)) + (want_log, want), (got_log, got) = runs + assert got_log == want_log + assert [entry[0] for entry in got_log] == ["empty_topk", "experts"] + assert want.shape == got.shape == (0, 4) + + def _st(route=False): return SimpleNamespace(slots=plan.Slots(route=route)) @@ -934,7 +956,7 @@ def test_the_draft_mkm_uses_the_complete_slot_schema(monkeypatch): assert not slots.route_lin and not slots.route_pre_topk -def test_the_route_slot_is_not_built_for_a_layer_it_cannot_serve(monkeypatch): +def test_the_route_slot_is_not_built_for_a_layer_it_cannot_serve(monkeypatch, caplog): from sgl_nvidia_mega.qwen3_5.moe import route mlp = _moe_block(_log(), torch.zeros(7, 4)) @@ -957,10 +979,12 @@ def test_the_route_slot_is_not_built_for_a_layer_it_cannot_serve(monkeypatch): ) want = "layer 0: router gate torch.float16 (512, 4096)" assert route.build(model, 7, fused=False).reason == want - with pytest.raises( - RuntimeError, match="requested slots \\['route'\\] built no width" - ): - runtime.build(model, ctx) + with caplog.at_level("INFO"): + built = runtime.build(model, ctx) + assert built.slots == {7: plan.Slots()} + assert "none (stock)" in built.summary() + assert "requested slots ['route'] use stock" in caplog.text + assert want in caplog.text def _gdn_build_model(): @@ -3210,10 +3234,12 @@ def forward(self, x, forward_batch=None, defer_finalize=False): return "moe_out" -def _decoder_layer(kind, log, defer): +def _decoder_layer(kind, log, defer, rows=7, tuple_input=False): layer = kind.__new__(kind) torch.nn.Module.__init__(layer) - a_in, m_in = torch.zeros(7, 4), torch.zeros(7, 4) + a_in, m_in = torch.zeros(rows, 4), torch.zeros(rows, 4) + if tuple_input: + a_in = (a_in, "scale") log.known.update(a_in=a_in, m_in=m_in) layer.mlp = _MoeBlock(log) layer.__dict__.update( @@ -3235,12 +3261,17 @@ def _decoder_layer(kind, log, defer): "kind", [qwen3_5.Qwen3_5LinearDecoderLayer, qwen3_5.Qwen3_5AttentionDecoderLayer] ) @pytest.mark.parametrize("defer", [False, True]) -def test_a_layer_runs_what_sglangs_decoder_layer_runs(kind, defer): - batch = SimpleNamespace(forward_mode=SimpleNamespace(is_idle=lambda: False)) +@pytest.mark.parametrize( + "rows,idle,tuple_input", [(7, False, True), (0, False, False), (7, True, False)] +) +def test_a_layer_runs_what_sglangs_decoder_layer_runs( + kind, defer, rows, idle, tuple_input +): + batch = SimpleNamespace(forward_mode=SimpleNamespace(is_idle=lambda: idle)) runs = [] for which in ("sglang", "plugin"): log = _log(hidden="hidden", batch=batch, positions="pos", h1="h1", part="part") - layer = _decoder_layer(kind, log, defer) + layer = _decoder_layer(kind, log, defer, rows=rows, tuple_input=tuple_input) if which == "sglang": out = layer.forward( positions="pos", @@ -3265,6 +3296,7 @@ def test_a_layer_runs_what_sglangs_decoder_layer_runs(kind, defer): runs.append((log, out)) (want_log, want), (got_log, got) = runs assert got_log == want_log and want == got == "layer_out" + assert any(entry[0] == "mixer" for entry in got_log) == (rows > 0 and not idle) def test_the_layer_loop_opens_the_stream_scopes_each_layer_and_closes_it(monkeypatch): From fef99c91bdc9cb4bec171338f85babc16fce5851 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 2 Nov 2026 04:40:59 -0800 Subject: [PATCH 50/62] perf(qwen3.5): balance fused CP CTAs by KV length through batch 64 Uneven KV lengths leave the equal per-request split plan waiting for its longest request. Select the existing KV-balanced plan by default and admit up to 64 requests for both verify and Q1, using the existing wide plan. Keep the equal plan available through the existing test override; add no environment switch and change no CUDA kernel arithmetic or precision. Add CPU tests for admission boundaries, complete per-rank tile coverage, and the longest CTA tile range on a ragged batch. Different split points can change floating-point accumulation order; bitwise equality is not expected. Validation: 3 CPU unit tests; 48 GPU cases against FP64 at local query heads 4/28, batches 1/8/32/64, mixed KV up to 131072; same error magnitude. Full GSM8K at T=1: 1289/1319 balanced versus 1287/1319 equal, zero empty or garbled-heuristic responses. Concurrent TEP8 c130 AgentX, one 1200s arm each: TPS/user P50 +9.57%, P90 +58.15%; TTFT P90 +3.34%. Arms ran in separate cabinets with 91.38% request-window overlap. --- .../qwen3_5/sm103/attention/q35_cp_fused.py | 24 +++++++++++-------- test/unit/test_qwen3_5_runtime.py | 12 ++++------ test/unit/test_qwen3_5_steps.py | 2 +- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py index 6d81a690..7ac9b62d 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py @@ -47,11 +47,13 @@ _LANE_PLAN_MAX_BATCH = 32 -# Keep the production selector on R11's unbalanced lane plan. The wide and -# two-slot kernels remain available to explicit test hooks, but R11 did not -# select them when SGLANG_PLB_CP_BALANCED was disabled. -MAX_BATCH = _LANE_PLAN_MAX_BATCH -MAX_BATCH_Q1 = _LANE_PLAN_MAX_BATCH +# Production default: distribute the grid's CTAs by owned KV length, including +# batches <= 32. Match R11 TEP8's 64-request admission for verify and draft. +# This is a plugin policy, not an environment switch; explicit split-mode test +# overrides below retain their separate counter requirement. +_DEFAULT_BALANCED = True +MAX_BATCH = _WIDE_PLAN_MAX_BATCH +MAX_BATCH_Q1 = _WIDE_PLAN_MAX_BATCH # Pack every two-slot call; the unpacked form remains available as a test hook. _TWO_SLOT_PACK = True @@ -275,12 +277,13 @@ def q35_cp_fused_supported( and batch <= sm_count ): return False - if batch > _LANE_PLAN_MAX_BATCH: + if _DEFAULT_BALANCED: return batch < _num_ctas( sm_count ) # the plan gives every request a CTA; it plans <= 128 requests return ( - cp_partial_geometry(rows, _num_splits(batch, sm_count))["max_cov_rows"] + batch <= _LANE_PLAN_MAX_BATCH + and cp_partial_geometry(rows, _num_splits(batch, sm_count))["max_cov_rows"] <= _MAX_COMB_ROWS ) @@ -562,7 +565,8 @@ def q35_cp_fused_partial( rows, hq] FP32 (base 2). Test hooks: debug (int64 [b * splits, 16] %globaltimer stamps; [CTAs, 16] when balanced), unsafe_skip_wait (no dependency waits), n_splits (must come with its own ``counters``: every launch on a counter row must use the same split count; balanced: the grid's CTAs), fresh_margin, balanced (a test override of the - split mode; the wide runtime path selects it automatically and an override needs its own ``counters``), plan + split mode; production defaults to KV-balanced allocation at every admitted batch, + and an override needs its own ``counters``), plan ("lane", "wide", "quad" or "two": the wide split's plan variant instead of the one the batch selects; lane / wide / quad cover up to 32 / 64 / 128 requests and compute the same one-slot plan, "two" (up to 128) the two-slot plan, with [CTAs, 24] debug rows), pack (a test override of the packed two-slot plan). @@ -620,7 +624,7 @@ def q35_cp_fused_partial( device = q_all.device sms = sm_count or _sm_count(device) state = _state(device, sms) - bal = b > _LANE_PLAN_MAX_BATCH if balanced is None else bool(balanced) + bal = _DEFAULT_BALANCED if balanced is None else bool(balanced) assert bal or b <= _LANE_PLAN_MAX_BATCH, "the equal split handles <= 32 requests" # plan variant: lane (<= 32 requests), wide (33..64), two-slot (65..128); `plan` overrides it in tests (also # "quad": the one-slot plan for 65..128) @@ -677,7 +681,7 @@ def q35_cp_fused_partial( if balanced is None and not state.balanced_logged: state.balanced_logged = True logger.info( - f"CP-lite first wide-plan call (q_len {q_len}, batch {b}, " + f"CP-lite first balanced call (q_len {q_len}, batch {b}, " f"cp_rank {cp_rank}, {splits} CTAs)" ) if wide and balanced is None and not state.wide_logged: diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 609de1a5..8debe050 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -242,7 +242,7 @@ def test_requested_cp_state_does_not_restore_a_device_length_gate(): assert [cp._gate_thr(batch) for batch in (2, 4, 16)] == [None, None, None] -def test_cp_exact_selector_keeps_r11s_unbalanced_batch_cap(): +def test_cp_exact_selector_admits_verify_and_q1_64(): from sgl_nvidia_mega.qwen3_5.sm103.attention import q35_cp_fused as fused from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite @@ -251,13 +251,9 @@ def test_cp_exact_selector_keeps_r11s_unbalanced_batch_cap(): cp.hq = 4 cp.sms = 132 - assert (fused.MAX_BATCH, fused.MAX_BATCH_Q1) == (32, 32) - assert [cp._exact_supported(7, batch) for batch in (32, 64, 128)] == [ - True, - False, - False, - ] - assert [cp._exact_supported(1, batch) for batch in (32, 64)] == [True, False] + assert (fused.MAX_BATCH, fused.MAX_BATCH_Q1) == (64, 64) + assert [cp._exact_supported(7, batch) for batch in (64, 65)] == [True, False] + assert [cp._exact_supported(1, batch) for batch in (64, 65)] == [True, False] def test_cp_call_override_pairs_the_exact_partial_and_merge_order(monkeypatch): diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index b55cce5a..91124d4c 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -2637,7 +2637,7 @@ def test_draft_only_cp_slots_do_not_enable_target_multi_request(monkeypatch): @pytest.mark.parametrize( ("requests", "eligible"), - [(5, True), (32, True), (33, False), (64, False), (128, False)], + [(5, True), (32, True), (64, True), (65, False)], ) def test_target_cp_multi_request_honors_fused_batch_cap( monkeypatch, requests, eligible From f037dfdef81668aca5ab4d7befea4912710a2705 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Mon, 2 Nov 2026 11:40:52 -0800 Subject: [PATCH 51/62] perf(qwen3.5): admit fused CP verify through batch 128 Keep verify batches 65..128 on the existing packed two-slot CP plan instead of falling back above 64 requests. Retain the Q1 limit of 64, the SM-count admission check, and rejection above 128. This changes plugin admission, not scheduler capacity or kernel arithmetic. Extend CPU coverage for admission boundaries and complete per-rank tile ownership at batches 65/96/128 in packed and unpacked plans. Update the runtime and step selector expectations for the new limit. Validation: four CPU tests pass; the GPU FP64-reference matrix covers 72 cases, including mixed KV lengths through 128K and batch 128. Full GSM8K with eight-way expert parallelism and max-running 128 scores 1290/1319 versus 1289/1319 with the cap-64 parent. Both use temperature 1 and have no empty responses or request errors. Runtime epoch deltas confirm batch >64 fused CP execution on all eight ranks. --- sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py | 5 +++-- test/unit/test_qwen3_5_runtime.py | 6 +++--- test/unit/test_qwen3_5_steps.py | 2 +- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py index 7ac9b62d..d2bbfd32 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/attention/q35_cp_fused.py @@ -48,11 +48,12 @@ # Production default: distribute the grid's CTAs by owned KV length, including -# batches <= 32. Match R11 TEP8's 64-request admission for verify and draft. +# batches <= 32. Cover verify through 128 requests with the existing packed +# two-slot plan; keep draft/Q1 admission at 64, matching the TEP8 recipe. # This is a plugin policy, not an environment switch; explicit split-mode test # overrides below retain their separate counter requirement. _DEFAULT_BALANCED = True -MAX_BATCH = _WIDE_PLAN_MAX_BATCH +MAX_BATCH = _TWO_SLOT_MAX_BATCH MAX_BATCH_Q1 = _WIDE_PLAN_MAX_BATCH # Pack every two-slot call; the unpacked form remains available as a test hook. _TWO_SLOT_PACK = True diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 8debe050..1b59d04b 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -242,7 +242,7 @@ def test_requested_cp_state_does_not_restore_a_device_length_gate(): assert [cp._gate_thr(batch) for batch in (2, 4, 16)] == [None, None, None] -def test_cp_exact_selector_admits_verify_and_q1_64(): +def test_cp_exact_selector_admits_verify_128_and_q1_64(): from sgl_nvidia_mega.qwen3_5.sm103.attention import q35_cp_fused as fused from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite @@ -251,8 +251,8 @@ def test_cp_exact_selector_admits_verify_and_q1_64(): cp.hq = 4 cp.sms = 132 - assert (fused.MAX_BATCH, fused.MAX_BATCH_Q1) == (64, 64) - assert [cp._exact_supported(7, batch) for batch in (64, 65)] == [True, False] + assert (fused.MAX_BATCH, fused.MAX_BATCH_Q1) == (128, 64) + assert [cp._exact_supported(7, batch) for batch in (128, 129)] == [True, False] assert [cp._exact_supported(1, batch) for batch in (64, 65)] == [True, False] diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 91124d4c..2073ac93 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -2637,7 +2637,7 @@ def test_draft_only_cp_slots_do_not_enable_target_multi_request(monkeypatch): @pytest.mark.parametrize( ("requests", "eligible"), - [(5, True), (32, True), (64, True), (65, False)], + [(5, True), (32, True), (128, True), (129, False)], ) def test_target_cp_multi_request_honors_fused_batch_cap( monkeypatch, requests, eligible From 4da63824019025d2411fdf1a689235bac12337eb Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 2 Nov 2026 12:52:37 -0800 Subject: [PATCH 52/62] perf(qwen3.5): run route_lin in the TP8 plan Add topology.plan: the capability selection plus the slots a topology's plan turns on with the requested ones. The TP8 plan, aggregated and PD decode, now runs route_lin whenever the policy holds the slots it runs through (route, mkm, hop, gdn_push, mka, mka_push, and lowm_se with route_shared), with no environment switch. TEP8, DEP8, TP2 and TP4 keep exactly their selected slots. The numerics claim covers every slot a plan can add. Validation: on aggregated TP8 static decode with the P27 slots, route_lin takes B1 from 2996.2 to 2849.8 us/step (verify 2262.8 to 2117.8), recovering 145 of the 152 us B1 gap to the frozen TP8 reference; B32 is 11306.1 vs 11243.5 us. Router logits in serving stay at 0.989x stock's RMS error against FP64 on the same quantized operands (GDN 0.989x, attention 0.990x; tolerance 1.05x). GSM10 passes 10/10 on both arms. --- README.md | 6 +++++ sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 14 ++++++++--- sgl_nvidia_mega/qwen3_5/spec.py | 6 ++--- sgl_nvidia_mega/qwen3_5/topology.py | 27 ++++++++++++++++++++ test/unit/test_qwen3_5_runtime.py | 28 +++++++++++++++++++++ test/unit/test_qwen3_5_spec.py | 18 +++++++++++++ 6 files changed, 92 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 4a9b0093..6d63ce5e 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,12 @@ includes `hop`. It remains opt-in rather than part of the bitwise default polic an explicit `SGLANG_NVIDIA_MEGAKERNEL_SLOTS` value replaces, rather than extends, that default list. +The TP8 plan (`topology.plan`, aggregated and PD decode) also runs `route_lin` +whenever the policy holds the slots it runs through (`route`, `mkm`, `hop`, +`gdn_push`, `mka`, `mka_push`, and `lowm_se` with `route_shared`): its router logits +stay within 1.05x stock's RMS error against FP64 on the same quantized operands. +TEP8, DEP8, TP2 and TP4 run exactly the requested slots their topology supports. + The router-only `route`, `route_topk` and `route_gemm` slots support TP2, TP4, TP8 and TEP8 with the TRTLLM runner and no all-to-all. PD decode TP2 uses the same local row counts: batch 6/10 gives 42/70 verify rows. These fronts leave the expert GEMMs with the stock runner, so TP2 needs no fused expert engine. `route_shared` keeps its TP8/TEP8 projection bounds; DEP8 keeps stock MoE dispatch and routing. `mka` uses the fused projection kernels through 14 verify rows on TP8 and 8 on diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 82d427b6..d98ad9b3 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -195,13 +195,21 @@ def apply_expert_parallel(ep_size: int) -> frozenset: def apply_topology(parallel) -> frozenset: - """Keep only each slot's declared capabilities before package builds run.""" + """Keep only each slot's declared capabilities, plus the slots the topology's plan + adds (``topology.plan``), before package builds run; returns the dropped names.""" from sgl_nvidia_mega.qwen3_5 import topology global SLOTS - selected = topology.select(SLOTS, topology.from_parallel(parallel)) - dropped = SLOTS - selected + current = topology.from_parallel(parallel) + selected = topology.plan(SLOTS, current) + dropped, added = SLOTS - selected, selected - SLOTS SLOTS = selected + if added: + logger.info( + "NVIDIA megakernel: Qwen3.5 %s plan: slots %s run with the requested ones", + current.name, + sorted(added), + ) return dropped diff --git a/sgl_nvidia_mega/qwen3_5/spec.py b/sgl_nvidia_mega/qwen3_5/spec.py index a84c3d6f..b76adb3b 100644 --- a/sgl_nvidia_mega/qwen3_5/spec.py +++ b/sgl_nvidia_mega/qwen3_5/spec.py @@ -213,7 +213,7 @@ def numerics(slots) -> Numerics: return Numerics(kind=NumericsKind.TOLERANCE, tolerances=tolerances) -NUMERICS = numerics(runtime.SLOTS) +NUMERICS = numerics(topology.plannable(runtime.SLOTS)) def text_config_reasons(hf_config) -> list[str]: @@ -298,9 +298,7 @@ def _config_disable_reasons(cfg, hf_config) -> list[str]: topology_reasons + text_config_reasons(hf_config) + binding_reasons() - + runtime.slot_reasons( - topology.select(runtime.SLOTS, topology.from_config(cfg)) - ) + + runtime.slot_reasons(topology.plan(runtime.SLOTS, topology.from_config(cfg))) + pd_reasons(cfg) ) diff --git a/sgl_nvidia_mega/qwen3_5/topology.py b/sgl_nvidia_mega/qwen3_5/topology.py index f4d96e9a..6da76594 100644 --- a/sgl_nvidia_mega/qwen3_5/topology.py +++ b/sgl_nvidia_mega/qwen3_5/topology.py @@ -214,3 +214,30 @@ def select(requested: frozenset, topology: Topology) -> frozenset: for name in requested if topology.name in SLOT_TOPOLOGIES.get(name, frozenset()) ) + + +# The slots route_lin runs through (runtime.slot_reasons); with route_shared, lowm_se +# replaces it at one request. +ROUTE_LIN_NEEDS = frozenset({"route", "mkm", "hop", "gdn_push", "mka", "mka_push"}) + + +def _route_lin_fits(slots: frozenset) -> bool: + return ROUTE_LIN_NEEDS <= slots and ( + "route_shared" not in slots or "lowm_se" in slots + ) + + +def plan(requested: frozenset, topology: Topology) -> frozenset: + """``select``, plus the slots a topology's plan turns on with the requested ones: + the TP8 plan (aggregated and PD decode) adds route_lin wherever its slots are + selected. Every other topology keeps exactly its selected slots.""" + selected = select(requested, topology) + if topology.name == "tp8" and _route_lin_fits(selected): + selected |= {"route_lin"} + return selected + + +def plannable(requested: frozenset) -> frozenset: + """``requested`` plus every slot some topology's plan can add to it (the numerics + claim made before the topology is known).""" + return requested | ({"route_lin"} if _route_lin_fits(requested) else frozenset()) diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 1b59d04b..0ddbceac 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -13,6 +13,7 @@ from sgl_nvidia_mega import plugin from sgl_nvidia_mega.qwen3_5 import ( runtime, + topology, ) from sgl_nvidia_mega.qwen3_5.runtime import executor, hooks, parity from sgl_nvidia_mega.qwen3_5.runtime.state import StepState @@ -1767,6 +1768,33 @@ def test_every_slot_declares_its_topologies(): assert set(topology.SLOT_TOPOLOGIES) == known +# The router-overlap slots the TP8 serving policy requests, without route_lin +ROUTE_LIN_SLOTS = frozenset( + "route route_shared lowm_se mkm hop gdn_verify gdn_push mka mka_push".split() +) + + +@pytest.mark.parametrize( + "tp, dp, ep, on", + [(8, 1, 1, True), (8, 1, 8, False), (8, 8, 8, False), (4, 1, 1, False)], +) +def test_only_the_tp8_plan_turns_route_lin_on(monkeypatch, tp, dp, ep, on): + monkeypatch.setattr(runtime, "SLOTS", ROUTE_LIN_SLOTS) + parallel = SimpleNamespace(tp_size=tp, attn_dp_size=dp, moe_ep_size=ep) + runtime.apply_topology(parallel) + selected = topology.select(ROUTE_LIN_SLOTS, topology.Topology(tp, dp, ep)) + assert runtime.SLOTS == selected | ({"route_lin"} if on else set()) + if on: + assert runtime.slot_reasons() == [] + + +@pytest.mark.parametrize("missing", [*sorted(topology.ROUTE_LIN_NEEDS), "lowm_se"]) +def test_the_tp8_plan_leaves_route_lin_off_without_one_of_its_slots(missing): + slots = ROUTE_LIN_SLOTS - {missing} + assert "route_lin" not in topology.plan(slots, topology.Topology(8)) + assert "route_lin" not in topology.plannable(slots) + + def test_unsupported_specialization_does_not_disable_a_supported_slot(monkeypatch): from sgl_nvidia_mega.qwen3_5 import topology diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index db4cb043..7ae8dbbc 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -24,6 +24,7 @@ plan, runtime, spec, + topology, ) @@ -665,6 +666,23 @@ def test_topology_admission_is_driven_by_plugin_capabilities(tp, pd): assert _refusal(tp_size=tp, nnodes=1, disaggregation_mode=pd) == "" +# The router-overlap slots the TP8 serving policy requests, without route_lin +ROUTE_LIN_SLOTS = frozenset( + "route route_shared lowm_se mkm hop gdn_verify gdn_push mka mka_push".split() +) + + +@pytest.mark.parametrize("pd", ["null", "decode"]) +def test_the_tp8_route_lin_plan_is_accepted_in_both_roles(monkeypatch, pd): + monkeypatch.setattr(runtime, "SLOTS", ROUTE_LIN_SLOTS) + assert _refusal(disaggregation_mode=pd) == "" + + +def test_the_numerics_claim_covers_the_slots_a_plan_can_add(): + claim = spec.numerics(topology.plannable(ROUTE_LIN_SLOTS)) + assert claim == spec.numerics(ROUTE_LIN_SLOTS | {"route_lin"}) + + def test_dep8_accepts_the_recipe_stock_moe_and_dp_lm_head(): assert ( _refusal( From 33aa1ecf66969d8826cc94bc498b4ad1ea304635 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 2 Nov 2026 16:30:55 -0800 Subject: [PATCH 53/62] perf(qwen3.5): balance TP8 wide MoE work by default Add the existing mkm_wide_bal variant to the TP8 topology plan whenever mkm_m256 or mkm_m512 is selected. Keep other topology plans unchanged and include this bitwise variant in the pre-topology numerics claim. No new environment switch or kernel change. Validation: the unchanged wide kernels passed 16/16 bitwise comparisons. Same-holder full-graph TP8 decode step time improves 3.06% at B32 and 2.35% at B64; GSM10 passes 10/10 on both arms. Focused CPU tests cover TP8 selection, unchanged TEP8 selection, and the existing numerics-plan claim. --- README.md | 4 +++- sgl_nvidia_mega/qwen3_5/moe/route.py | 2 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 2 +- sgl_nvidia_mega/qwen3_5/topology.py | 10 ++++++++-- test/unit/test_qwen3_5_spec.py | 2 ++ test/unit/test_qwen3_5_wide_balance_plan.py | 13 +++++++++++++ 6 files changed, 28 insertions(+), 5 deletions(-) create mode 100644 test/unit/test_qwen3_5_wide_balance_plan.py diff --git a/README.md b/README.md index 6d63ce5e..97d0437e 100644 --- a/README.md +++ b/README.md @@ -59,7 +59,7 @@ Qwen3.5 (`qwen3_5/runtime`) runs the same way: per layer the stage boundary (`bo | `multi_request` | 2–128 requests | The target graphs keep every offered request count instead of padding to a coarser width | | `draft_loop`, `draft_extend`, `draft_multi`, `draft_embed` | 1–128 requests | EAGLE draft entries, exact sharded top-1 and embedding exchange; the pre-FC norms write directly into the FC input | -The remaining kernel families are opt-in through that one slot policy. These change the arithmetic within the tolerances `spec.py` declares; `mkm_wide_bal` also stays opt-in because its serving evidence covers TP8 only: +The remaining kernel families are opt-in through that one slot policy. These change the arithmetic within the tolerances `spec.py` declares: | Slot | Rows | Kernel | | --- | --- | --- | @@ -83,6 +83,8 @@ The TP8 plan (`topology.plan`, aggregated and PD decode) also runs `route_lin` whenever the policy holds the slots it runs through (`route`, `mkm`, `hop`, `gdn_push`, `mka`, `mka_push`, and `lowm_se` with `route_shared`): its router logits stay within 1.05x stock's RMS error against FP64 on the same quantized operands. +The same TP8 plan adds `mkm_wide_bal` for `mkm_m256` / `mkm_m512`: balanced +CTA assignment at 129–512 verify rows, with bitwise-identical outputs. TEP8, DEP8, TP2 and TP4 run exactly the requested slots their topology supports. The router-only `route`, `route_topk` and `route_gemm` slots support TP2, TP4, TP8 and TEP8 with the TRTLLM runner and no all-to-all. PD decode TP2 uses the same local row counts: batch 6/10 gives 42/70 verify rows. These fronts leave the expert GEMMs with the stock runner, so TP2 needs no fused expert engine. `route_shared` keeps its TP8/TEP8 projection bounds; DEP8 keeps stock MoE dispatch and routing. diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index 76dafc68..1551c43c 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -153,7 +153,7 @@ class Bound(msgspec.Struct, frozen=True): m128: bool = False # slot mkm_m128: the 128-token engine built (on m32's maps) m256: bool = False # slot mkm_m256: the 256-token engine built (on m32's maps) m512: bool = False # slot mkm_m512: the 512-token engine built (on m32's maps) - wide_bal: bool = False # opt-in balanced work split for mkm_m256 / mkm_m512 + wide_bal: bool = False # balanced work split for mkm_m256 / mkm_m512 # A caller-specific upper bound for the 16-token engine. The draft uses 12 so # its alternate-stream shared expert can overlap rows 13-16 with the 32-token engine. m16_max_rows: int = 16 diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index d98ad9b3..15b9a868 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -62,7 +62,7 @@ # shape instead of padding. OPT_IN_COVERAGE: tuple[str, ...] = ("multi_request",) # Graph or build variants of a slot, requested in the same variable: no width of their -# own. BAL stays opt-in because its serving evidence covers TP8 only. +# own. The TP8 plan enables BAL for its wide engines; elsewhere it stays explicit. # skv2_cp: one long request past the measured crossover runs CP-lite in place of SKV2, # in a graph variant of its own (attention/route.py) OPT_IN_VARIANTS: tuple[str, ...] = ("skv2_cp", "mkm_wide_bal") diff --git a/sgl_nvidia_mega/qwen3_5/topology.py b/sgl_nvidia_mega/qwen3_5/topology.py index 6da76594..e1c53f1e 100644 --- a/sgl_nvidia_mega/qwen3_5/topology.py +++ b/sgl_nvidia_mega/qwen3_5/topology.py @@ -230,14 +230,20 @@ def _route_lin_fits(slots: frozenset) -> bool: def plan(requested: frozenset, topology: Topology) -> frozenset: """``select``, plus the slots a topology's plan turns on with the requested ones: the TP8 plan (aggregated and PD decode) adds route_lin wherever its slots are - selected. Every other topology keeps exactly its selected slots.""" + selected, and balances the mkm_m256/mkm_m512 work split when offered. + Every other topology keeps exactly its selected slots.""" selected = select(requested, topology) if topology.name == "tp8" and _route_lin_fits(selected): selected |= {"route_lin"} + if topology.name == "tp8" and {"mkm_m256", "mkm_m512"} & selected: + selected |= {"mkm_wide_bal"} return selected def plannable(requested: frozenset) -> frozenset: """``requested`` plus every slot some topology's plan can add to it (the numerics claim made before the topology is known).""" - return requested | ({"route_lin"} if _route_lin_fits(requested) else frozenset()) + planned = requested | ({"route_lin"} if _route_lin_fits(requested) else frozenset()) + if {"mkm_m256", "mkm_m512"} & requested: + planned |= {"mkm_wide_bal"} + return planned diff --git a/test/unit/test_qwen3_5_spec.py b/test/unit/test_qwen3_5_spec.py index 7ae8dbbc..68b2b71f 100644 --- a/test/unit/test_qwen3_5_spec.py +++ b/test/unit/test_qwen3_5_spec.py @@ -681,6 +681,8 @@ def test_the_tp8_route_lin_plan_is_accepted_in_both_roles(monkeypatch, pd): def test_the_numerics_claim_covers_the_slots_a_plan_can_add(): claim = spec.numerics(topology.plannable(ROUTE_LIN_SLOTS)) assert claim == spec.numerics(ROUTE_LIN_SLOTS | {"route_lin"}) + shipped = frozenset(runtime.SHIPPED_SLOTS) + assert topology.plannable(shipped) == shipped | {"mkm_wide_bal"} def test_dep8_accepts_the_recipe_stock_moe_and_dp_lm_head(): diff --git a/test/unit/test_qwen3_5_wide_balance_plan.py b/test/unit/test_qwen3_5_wide_balance_plan.py new file mode 100644 index 00000000..cf25eab7 --- /dev/null +++ b/test/unit/test_qwen3_5_wide_balance_plan.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from sgl_nvidia_mega.qwen3_5 import topology + +WIDE = frozenset({"route_topk", "multi_request", "mkm_m256", "mkm_m512"}) + + +def test_tp8_plan_adds_balanced_wide_assignment(): + assert topology.plan(WIDE, topology.Topology(8)) == WIDE | {"mkm_wide_bal"} + + +def test_tep8_plan_keeps_the_existing_assignment(): + assert topology.plan(WIDE, topology.Topology(8, moe_ep_size=8)) == WIDE From b83045bcf3ae3c4b2251bca8bc372e312d8ead4f Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 2 Nov 2026 05:10:24 -0800 Subject: [PATCH 54/62] perf(qwen3.5): fall back to stock for unavailable default-plan slots --- sgl_nvidia_mega/qwen3_5/boundary/hop.py | 27 ++-- sgl_nvidia_mega/qwen3_5/draft/__init__.py | 10 +- sgl_nvidia_mega/qwen3_5/draft/attention.py | 29 ++-- sgl_nvidia_mega/qwen3_5/draft/moe.py | 8 +- sgl_nvidia_mega/qwen3_5/draft/parts.py | 133 +++++++++++++----- sgl_nvidia_mega/qwen3_5/draft/xchg.py | 15 +- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 103 ++++++++++++-- sgl_nvidia_mega/qwen3_5/runtime/executor.py | 26 ++-- sgl_nvidia_mega/qwen3_5/sm103/draft/embed.py | 5 +- sgl_nvidia_mega/qwen3_5/sm103/draft/top1.py | 2 +- sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py | 7 +- sgl_nvidia_mega/qwen3_5/sm103/hop/__init__.py | 5 +- test/unit/test_qwen3_5_runtime.py | 132 +++++++++++++++++ 13 files changed, 403 insertions(+), 99 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/boundary/hop.py b/sgl_nvidia_mega/qwen3_5/boundary/hop.py index ada47647..1f2ee62e 100644 --- a/sgl_nvidia_mega/qwen3_5/boundary/hop.py +++ b/sgl_nvidia_mega/qwen3_5/boundary/hop.py @@ -29,6 +29,8 @@ import msgspec import torch +from sgl_nvidia_mega.qwen3_5.runtime import slot_failure + logger = logging.getLogger(__name__) # The SGLang names this module reads (test/unit/test_contract.py pins them) @@ -173,14 +175,18 @@ def build(model, ctx) -> Bound: megakernel_agree(reasons, what="hop refusals") if reasons: return Bound(reason="; ".join(reasons)) - hop = Hop( - max_m=rows, - eps=m.flashinfer_mnnvl_cutedsl_fusion.rms_epsilon, - group=group.device_group, - device=torch.device("cuda", torch.cuda.current_device()), - deadline_s=ctx.deadline_s, - route_lin="route_lin" in runtime.SLOTS, - ) + try: + hop = Hop( + max_m=rows, + eps=m.flashinfer_mnnvl_cutedsl_fusion.rms_epsilon, + group=group.device_group, + device=torch.device("cuda", torch.cuda.current_device()), + deadline_s=ctx.deadline_s, + route_lin="route_lin" in runtime.SLOTS, + ) + except runtime.SlotUnavailable as exc: + runtime.unavailable_slots(exc.slots, exc.reason) + return Bound(reason=exc.reason) global _ACTIVE _ACTIVE = hop return Bound(hop=hop, mid=rows > plan.HOP_MAX_ROWS) @@ -384,7 +390,10 @@ def moe_input(hop, layer, hidden, st): lid, h0 = st.scratch.get("rlin.h0", (None, None)) if lid != layer.layer_id or h0 is None or h0.data_ptr() != residual.data_ptr(): raise RuntimeError( - f"route-lin: layer {layer.layer_id}'s residual is not the one its router partials read" + slot_failure( + "route_lin", + f"layer {layer.layer_id}'s residual is not the one its router partials read", + ) ) _engage("MoE input + router (route-lin)", rows) (norm, residual, *fp8), rl = hop.all_reduce_residual_rms_norm_rl( diff --git a/sgl_nvidia_mega/qwen3_5/draft/__init__.py b/sgl_nvidia_mega/qwen3_5/draft/__init__.py index f807735a..f9a6b624 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/draft/__init__.py @@ -190,7 +190,10 @@ def _parts(self): def _draft_loop(self, *, num_steps: int, max_bs: int, draft_tokens): from sgl_nvidia_mega.qwen3_5.draft.loop import DraftLoop - self._loop = DraftLoop(self._parts(), num_steps=num_steps) + parts = self._parts() + if parts is None: + return lambda bs: None + self._loop = DraftLoop(parts, num_steps=num_steps) return self._loop.step_for def _begin_draft_round(self, io) -> None: @@ -202,7 +205,10 @@ def _draft_extend( ): from sgl_nvidia_mega.qwen3_5.draft.extend import DraftExtend - extend = DraftExtend(self._parts(), tokens_per_request=tokens_per_request) + parts = self._parts() + if parts is None: + return lambda bs: None + extend = DraftExtend(parts, tokens_per_request=tokens_per_request) return extend.step_for diff --git a/sgl_nvidia_mega/qwen3_5/draft/attention.py b/sgl_nvidia_mega/qwen3_5/draft/attention.py index 8535bc75..e948c896 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/attention.py +++ b/sgl_nvidia_mega/qwen3_5/draft/attention.py @@ -385,17 +385,21 @@ def __init__(self, split_kv, route): @classmethod def build(cls, parts): + from sgl_nvidia_mega.qwen3_5.runtime import unavailable_slots + try: from sgl_nvidia_mega.qwen3_5.attention import route, split_kv except ImportError as e: - raise RuntimeError( - "NVIDIA megakernel: slot draft_skv2 needs the attention package's SKV2 " - f"(qwen3_5.attention.split_kv): {e}" - ) from e + unavailable_slots( + "draft_skv2", + f"needs the attention package's SKV2 (qwen3_5.attention.split_kv): {e}", + ) + return None width = int(parts.binding.tokens_per_request) why = split_kv.unsupported(parts.layer.attn, width) if why: - raise RuntimeError(f"NVIDIA megakernel: slot draft_skv2 cannot run: {why}") + unavailable_slots("draft_skv2", why) + return None return cls(split_kv, route) def long(self) -> bool: @@ -439,11 +443,15 @@ def build(cls, parts): why = split_kv.cp_reason() if why: + from sgl_nvidia_mega.qwen3_5.runtime import unavailable_slots + if multi and not (steps or extend): + logger.info( + "NVIDIA megakernel: draft_multi CP attention uses stock (%s)", why + ) return None - raise RuntimeError( - f"NVIDIA megakernel: slots draft_q1_cp / draft_skv2_cp cannot run: {why}" - ) + unavailable_slots(parts.slots & {"draft_q1_cp", "draft_skv2_cp"}, why) + return None return cls(split_kv, route, steps, extend, multi) def takes(self, rows_per_request: int, rows: int) -> bool: @@ -518,7 +526,10 @@ def build_q1(parts): device_index=device.index, max_context_len=attention.max_context_len ) if why is not None: - raise RuntimeError(f"NVIDIA megakernel: slot draft_q1 cannot run: {why}") + from sgl_nvidia_mega.qwen3_5.runtime import unavailable_slots + + unavailable_slots("draft_q1", why) + return None release = GATE_FOLD == 2 q1.build(device.index, heads=attention.heads, pdl=Q1_PDL, release=release) decode = q1.Q1Decode(pdl=Q1_PDL, release=release) diff --git a/sgl_nvidia_mega/qwen3_5/draft/moe.py b/sgl_nvidia_mega/qwen3_5/draft/moe.py index 2fcd4069..a279bb99 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/moe.py +++ b/sgl_nvidia_mega/qwen3_5/draft/moe.py @@ -151,11 +151,9 @@ def build(cls, mlp, device_index: int) -> Optional[DraftMoe]: ) ) if reason: - logger.warning( - "Single-launch draft MoE off for layer %d (%s); SGLang's MoE runs.", - mlp.layer_id, - reason, - ) + from sgl_nvidia_mega.qwen3_5.runtime import unavailable_slots + + unavailable_slots("draft_moe", reason) return None logger.info("Single-launch draft MoE active for layer %d.", mlp.layer_id) return cls(mlp, torch.device("cuda", device_index)) diff --git a/sgl_nvidia_mega/qwen3_5/draft/parts.py b/sgl_nvidia_mega/qwen3_5/draft/parts.py index 66061faf..0319e873 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/parts.py +++ b/sgl_nvidia_mega/qwen3_5/draft/parts.py @@ -78,28 +78,35 @@ def __init__(self, draft, layer): self.moe_mkm: Optional[Any] = None @classmethod - def build(cls, draft) -> Parts: + def build(cls, draft) -> Parts | None: from sgl_nvidia_mega.qwen3_5.draft.attention import Attention + from sgl_nvidia_mega.qwen3_5.runtime import unavailable_slots m = draft.lm.model layers = [m.layers[i] for i in range(m.start_layer, m.end_layer)] if len(layers) != 1: - raise RuntimeError( - f"NVIDIA megakernel: the MTP draft has {len(layers)} layers, not 1" + unavailable_slots( + draft.slots, f"the MTP draft has {len(layers)} layers, not 1" ) + return None + if getattr(draft, "context_len", 0) is None: + unavailable_slots(draft.slots, "the draft's HF config has no context_len") + return None parts = cls(draft, layers[0]) - parts.attention = Attention(parts) + try: + parts.attention = Attention(parts) + except ImportError as exc: + unavailable_slots( + draft.slots, f"draft attention dependency unavailable: {exc}" + ) + return None slots = parts.slots if "draft_fin" in slots: from sgl_nvidia_mega.qwen3_5.draft import finalize reason = finalize.enable(parts.layer.mlp) if reason: - logger.warning( - "NVIDIA megakernel: slot draft_fin cannot run (%s); " - "SGLang finalizes the draft MoE.", - reason, - ) + unavailable_slots("draft_fin", reason) slots = parts.slots = slots - {"draft_fin"} if "draft_moe" in slots: parts.moe = parts._build_moe() @@ -107,12 +114,18 @@ def build(cls, draft) -> Parts: slots = parts.slots = slots - {"draft_moe"} if "draft_hop" in slots: parts.hop = parts._build_hop() + if parts.hop is None: + slots = parts.slots = slots - {"draft_hop", "draft_fin"} if "draft_xchg" in slots: parts.mkd = parts._build_mkd() if "draft_multi" in slots: parts.top1x = parts._build_top1x() + if parts.top1x is None: + slots = parts.slots = slots - {"draft_multi", "draft_embed"} if "draft_embed" in slots: parts.embedx = parts._build_embedx() + if parts.embedx is None: + slots = parts.slots = slots - {"draft_embed"} if "draft_moe" in slots: from sgl_nvidia_mega.qwen3_5.draft import moe_mkm @@ -158,7 +171,10 @@ def _build_hop(self): early_head=True, ) if hop is None: - raise RuntimeError(f"NVIDIA megakernel: slot draft_hop cannot run: {why}") + from sgl_nvidia_mega.qwen3_5.runtime import unavailable_slots + + unavailable_slots("draft_hop", why) + return None logger.info( "NVIDIA megakernel: Qwen3.5 MTP draft hop releases at entry: the " "single-launch MoE %s, the draft head %s", @@ -204,9 +220,10 @@ def _build_mkd(self): num_steps=b.num_steps, ) if mkd is None: - raise RuntimeError( - "NVIDIA megakernel: slot draft_xchg cannot run: " - + ("; ".join(refused) or "see the MKD draft boundary warning") + from sgl_nvidia_mega.qwen3_5.runtime import unavailable_slots + + unavailable_slots( + "draft_xchg", "; ".join(refused) or "see the MKD draft boundary warning" ) return mkd @@ -218,22 +235,42 @@ def _build_top1x(self): start = int(lm.lm_head.shard_indices.org_vocab_start_index) vocab = int(lm.logits_processor.vocab_size) # the shard layout the all-gathered logits have: rank-major, no padding rows + reason = "" if start != b.tp_rank * rows or rows * b.tp_size != vocab: - raise RuntimeError( - "NVIDIA megakernel: slot draft_multi cannot run: the LM head's " + reason = ( + "the LM head's " f"{rows}-row shard at {start} is not rank {b.tp_rank}'s slice of " f"{vocab} ids" ) - return Top1Exchange( - world=b.tp_size, - rank=b.tp_rank, - vocab_local=rows, - vocab_start=start, - vocab_size=vocab, - group=b.tp_device_group, - cpu_group=b.tp_cpu_group, - device=self.device, - ) + weight = lm.lm_head.weight + if ( + weight.dtype != torch.bfloat16 + or weight.dim() != 2 + or not weight.is_contiguous() + ): + reason = f"LM head {weight.dtype} {tuple(weight.shape)}; needs a contiguous BF16 shard" + from sgl_nvidia_mega.qwen3_5.runtime import slot_supported + + if not slot_supported( + "draft_multi", reason, cpu_group=getattr(b, "tp_cpu_group", None) + ): + return None + from sgl_nvidia_mega.qwen3_5.runtime import SlotUnavailable, unavailable_slots + + try: + return Top1Exchange( + world=b.tp_size, + rank=b.tp_rank, + vocab_local=rows, + vocab_start=start, + vocab_size=vocab, + group=b.tp_device_group, + cpu_group=b.tp_cpu_group, + device=self.device, + ) + except SlotUnavailable as exc: + unavailable_slots("draft_multi", exc.reason) + return None def _build_embedx(self): from sgl_nvidia_mega.qwen3_5.sm103.draft.embed import EmbedExchange @@ -241,20 +278,42 @@ def _build_embedx(self): b = self.binding embed = self.model.embed_tokens shard = embed.shard_indices + reason = "" if shard.added_vocab_end_index != shard.added_vocab_start_index: - raise RuntimeError( - "NVIDIA megakernel: slot draft_embed cannot run: the embedding has " - "added vocabulary" + reason = "the embedding has added vocabulary" + table = embed.weight.data + local = int(shard.org_vocab_end_index - shard.org_vocab_start_index) + if ( + table.dtype != torch.bfloat16 + or table.dim() != 2 + or not table.is_contiguous() + or table.shape[1] % 2 + or table.shape[0] < local + ): + reason = f"embedding table {table.dtype} {tuple(table.shape)}, local vocabulary {local}; needs contiguous BF16 rows with an even hidden size" + from sgl_nvidia_mega.qwen3_5.runtime import slot_supported + + if not slot_supported( + "draft_embed", reason, cpu_group=getattr(b, "tp_cpu_group", None) + ): + return None + from sgl_nvidia_mega.qwen3_5.runtime import SlotUnavailable, unavailable_slots + + try: + return EmbedExchange( + world=b.tp_size, + table=embed.weight.data, + vocab_start=int(shard.org_vocab_start_index), + vocab_local=int( + shard.org_vocab_end_index - shard.org_vocab_start_index + ), + group=b.tp_device_group, + cpu_group=b.tp_cpu_group, + device=self.device, ) - return EmbedExchange( - world=b.tp_size, - table=embed.weight.data, - vocab_start=int(shard.org_vocab_start_index), - vocab_local=int(shard.org_vocab_end_index - shard.org_vocab_start_index), - group=b.tp_device_group, - cpu_group=b.tp_cpu_group, - device=self.device, - ) + except SlotUnavailable as exc: + unavailable_slots("draft_embed", exc.reason) + return None def embed(self, token_ids: torch.Tensor) -> torch.Tensor: """The embedding rows of ``token_ids``: slot draft_embed's exchange where it is diff --git a/sgl_nvidia_mega/qwen3_5/draft/xchg.py b/sgl_nvidia_mega/qwen3_5/draft/xchg.py index 7a9383a0..381f8a2a 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/xchg.py +++ b/sgl_nvidia_mega/qwen3_5/draft/xchg.py @@ -30,6 +30,7 @@ import torch +from sgl_nvidia_mega.qwen3_5.runtime import slot_failure from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels logger = logging.getLogger(__name__) @@ -185,7 +186,9 @@ def __init__( self.embed_weight, self.local_vocab = embed_shard, rows if vocab_start != rank * rows: raise RuntimeError( - "MKD draft boundary: vocabulary shard is not the rank-th equal slice" + slot_failure( + "draft_xchg", "vocabulary shard is not the rank-th equal slice" + ) ) # Full vocabulary: candidate position p (rank-major) is vocabulary id p. self.hot_ids = torch.arange(rows * world, dtype=torch.int64).to(dev) @@ -304,8 +307,10 @@ def _build_memo(self) -> torch.Tensor: dist.all_reduce(bad_t, group=self._device_group) if int(bad_t[0]) != 0: raise RuntimeError( - f"MKD memo table differs from the split fc on {int(bad_t[0])} of " - f"{int(bad_t[1])} rows" + slot_failure( + "draft_xchg", + f"MKD memo table differs from the split fc on {int(bad_t[0])} of {int(bad_t[1])} rows", + ) ) logger.info( "MKD memo table: %d rows on this rank (%.1f MB), equal to the split fc on all %d rows.", @@ -524,7 +529,9 @@ def finish_step( topk_index, topk_p = self.topk_index[step][:m], self.topk_p[step][:m] if not self.armed: raise RuntimeError( - "MKD draft boundary: only armed (one-row) steps take the exchange" + slot_failure( + "draft_xchg", "only armed (one-row) steps take the exchange" + ) ) self.exchange( local_logits, diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 15b9a868..6ef7cf08 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -122,6 +122,70 @@ def _requested_slots() -> frozenset: SLOTS = _requested_slots() +def slot_failure(slots, reason: str) -> str: + """A requested slot's failed capability, including the recovery action.""" + names = [slots] if isinstance(slots, str) else sorted(slots) + label = f"slot {names[0]}" if len(names) == 1 else f"slots {names}" + return ( + f"NVIDIA megakernel: {label} cannot run: {reason}; " + f"remove it from {SLOTS_ENV} or change the topology" + ) + + +class SlotUnavailable(RuntimeError): + """A build capability refusal, distinct from CUDA errors and kernel failures.""" + + def __init__(self, slots, reason: str): + self.slots, self.reason = slots, reason + super().__init__(slot_failure(slots, reason)) + + +def require_multicast(slots, pointer: int, group, device) -> None: + """All ranks reject an unavailable mapping before any consumer starts.""" + import torch + import torch.distributed as dist + + valid = torch.tensor( + [bool(pointer) and pointer % 16 == 0], dtype=torch.int32, device=device + ) + dist.all_reduce(valid, op=dist.ReduceOp.MIN, group=group) + if not int(valid.item()): + raise SlotUnavailable( + slots, "needs an aligned NVLink multicast mapping on every rank" + ) + + +def unavailable_slots(slots, reason: str) -> None: + """Keep defaults optional at a capability check that was already strict. + + Called at the capability check, before the unavailable implementation runs. + The caller keeps its existing stock implementation / unserved entry. + """ + names = [slots] if isinstance(slots, str) else sorted(slots) + if names and os.environ.get(SLOTS_ENV) is not None: + raise RuntimeError(slot_failure(slots, reason)) + logger.info( + "NVIDIA megakernel: requested slots %s use stock (default policy; %s)", + names, + reason, + ) + + +def slot_supported(slots, reason: str, *, cpu_group=None) -> bool: + """Agree on rank-local shard checks before an exchange's collective allocation.""" + if cpu_group is not None: + import torch.distributed as dist + + if dist.is_initialized(): + reasons = [None] * dist.get_world_size(cpu_group) + dist.all_gather_object(reasons, reason, group=cpu_group) + reason = "; ".join(dict.fromkeys(r for r in reasons if r)) + if not reason: + return True + unavailable_slots(slots, reason) + return False + + def apply_cp_lite_topology(attention_tp_size: int) -> frozenset: """Drop CP-lite slots when the attention-TP Q/KV topology cannot serve them.""" from sgl_nvidia_mega.qwen3_5 import cp_lite_topology @@ -165,10 +229,10 @@ def build_cp_lite(binding) -> None: what="Qwen3.5 CP page size and TP size", ) if binding.page_size != 64: - raise RuntimeError( - "NVIDIA megakernel: Qwen3.5 CP slots need page size 64 " - f"(got {binding.page_size})" + unavailable_slots( + requested, f"CP slots need page size 64 (got {binding.page_size})" ) + return import torch @@ -176,9 +240,8 @@ def build_cp_lite(binding) -> None: cp = CpLite.maybe_create(binding, torch.device("cuda", torch.cuda.current_device())) if cp is None: - raise RuntimeError( - "NVIDIA megakernel: Qwen3.5 CP slots need a supported attention-TP " - "Q/KV-head topology" + unavailable_slots( + requested, "CP slots need a supported attention-TP Q/KV-head topology" ) @@ -191,6 +254,13 @@ def apply_expert_parallel(ep_size: int) -> frozenset: return frozenset() dropped = SLOTS & EP_STOCK_SLOTS SLOTS = SLOTS - EP_STOCK_SLOTS + if dropped: + logger.info( + "NVIDIA megakernel: slots %s use stock under expert parallelism " + "(moe_ep_size %d)", + sorted(dropped), + ep_size, + ) return dropped @@ -210,6 +280,12 @@ def apply_topology(parallel) -> frozenset: current.name, sorted(added), ) + if dropped: + logger.info( + "NVIDIA megakernel: slots %s use stock (topology %s)", + sorted(dropped), + current.name, + ) return dropped @@ -444,6 +520,13 @@ def build(model, ctx): if getattr(parallel, "pp_size", 1) == 4: from sgl_nvidia_mega.qwen3_5.runtime.prefill import Qwen35Prefill + unsupported = SLOTS - {"gdn_prefill_split_conv"} + if unsupported: + logger.info( + "NVIDIA megakernel: slots %s use stock (PD prefill TP1/PP4 " + "only supports gdn_prefill_split_conv)", + sorted(unsupported), + ) return Qwen35Prefill(model, ctx, "gdn_prefill_split_conv" in SLOTS) dropped = apply_expert_parallel(parallel.moe_ep_size) @@ -471,9 +554,11 @@ def build(model, ctx): multi = slot_plan.multi_request(parallel.tp_size, SLOTS) widths = slot_plan.candidates(ctx.widths, ctx.tokens_per_request, multi) if not widths: - raise RuntimeError( - f"NVIDIA megakernel: {ctx.tokens_per_request} rows a request leave no " - "covered width" + unavailable_slots( + SLOTS, f"{ctx.tokens_per_request} rows a request leave no covered width" + ) + return Qwen35Verify( + model, {}, {}, (), binding=ctx.binding, summary="none (stock)" ) pctx = PackageContext( widths=widths, diff --git a/sgl_nvidia_mega/qwen3_5/runtime/executor.py b/sgl_nvidia_mega/qwen3_5/runtime/executor.py index fa8b4ed0..dfd3cb12 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/executor.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/executor.py @@ -38,7 +38,7 @@ from sgl_nvidia_mega.qwen3_5 import attention, boundary, gdn, logits, moe from sgl_nvidia_mega.qwen3_5.attention import route as _attention_route # noqa: F401 from sgl_nvidia_mega.qwen3_5.gdn import verify as _gdn_verify # noqa: F401 -from sgl_nvidia_mega.qwen3_5.runtime import embedding +from sgl_nvidia_mega.qwen3_5.runtime import embedding, slot_failure from sgl_nvidia_mega.qwen3_5.runtime.state import StepState # The SGLang names this module reads (test/unit/test_contract.py pins them) @@ -191,33 +191,23 @@ def _check_pools(self, view): view.mamba, self.binding.tokens_per_request ) if reasons: - raise RuntimeError( - "NVIDIA megakernel: slot gdn_verify cannot run: " - + "; ".join(reasons) - ) + raise RuntimeError(slot_failure("gdn_verify", "; ".join(reasons))) if any(s.gdn_conv for s in self.slots.values()): reasons = gdn.verify.conv_pool_reasons( view.mamba, self.binding.tokens_per_request ) if reasons: - raise RuntimeError( - "NVIDIA megakernel: slot gdn_conv cannot run: " + "; ".join(reasons) - ) + raise RuntimeError(slot_failure("gdn_conv", "; ".join(reasons))) if any(s.gdn_core for s in self.slots.values()): reasons = gdn.verify.core_pool_reasons(view.mamba) if reasons: - raise RuntimeError( - "NVIDIA megakernel: slot gdn_core cannot run: " + "; ".join(reasons) - ) + raise RuntimeError(slot_failure("gdn_core", "; ".join(reasons))) if any(s.gdn_front for s in self.slots.values()): from sgl_nvidia_mega.qwen3_5.gdn import front reasons = front.pool_reasons(view.mamba) if reasons: - raise RuntimeError( - "NVIDIA megakernel: slot gdn_front cannot run: " - + "; ".join(reasons) - ) + raise RuntimeError(slot_failure("gdn_front", "; ".join(reasons))) front.prepare( self.lm, view.mamba, @@ -229,8 +219,10 @@ def _check_pools(self, view): reasons = attention.pool_reasons(view, self.binding, skv2=skv2, mka=mka) if reasons: raise RuntimeError( - "NVIDIA megakernel: the attention slots cannot run: " - + "; ".join(reasons) + slot_failure( + {n for n, on in (("skv2", skv2), ("mka", mka)) if on}, + "; ".join(reasons), + ) ) self._pools_checked = True diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/embed.py b/sgl_nvidia_mega/qwen3_5/sm103/draft/embed.py index 5bb23a81..50a6b247 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/embed.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/embed.py @@ -92,7 +92,10 @@ def __init__( device, ) self._buf, self._handle, self.mc = symmetric_words( - SLOTS * MAX_ROWS * self.hidden // 2 + world, group, device + SLOTS * MAX_ROWS * self.hidden // 2 + world, + group, + device, + slot="draft_embed", ) self.uc = self._buf.data_ptr() # [epoch from 1, CTA arrivals] diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/top1.py b/sgl_nvidia_mega/qwen3_5/sm103/draft/top1.py index 9b14eb2a..25ce5d2c 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/top1.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/top1.py @@ -104,7 +104,7 @@ def __init__( self.vocab_local, self.vocab_start = vocab_local, vocab_start self.vocab_size = vocab_size self._buf, self._handle, self.mc = symmetric_words( - SLOTS * world * MAX_ROWS * 2, group, device + SLOTS * world * MAX_ROWS * 2, group, device, slot="draft_multi" ) self.uc = self._buf.data_ptr() # [epoch from 1, CTA arrivals] diff --git a/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py b/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py index e284229f..6604c167 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/draft/xchg.py @@ -714,7 +714,7 @@ def ordered_partial_sum( return out -def symmetric_words(words: int, group, device: torch.device): +def symmetric_words(words: int, group, device: torch.device, *, slot="draft_xchg"): """(tensor, handle, multicast pointer) of a zeroed symmetric int64 buffer on ``group``; a collective on every rank of it.""" import torch.distributed._symmetric_memory as symm_mem @@ -725,8 +725,9 @@ def symmetric_words(words: int, group, device: torch.device): buf = symm_mem.empty((words,), dtype=torch.int64, device=device) handle = symm_mem.rendezvous(buf, group) multicast = int(handle.multicast_ptr or 0) - if not multicast or multicast % 16: - raise RuntimeError("the draft step boundary needs an NVLink multicast mapping") + from sgl_nvidia_mega.qwen3_5.runtime import require_multicast + + require_multicast(slot, multicast, group, device) buf.zero_() return buf, handle, multicast diff --git a/sgl_nvidia_mega/qwen3_5/sm103/hop/__init__.py b/sgl_nvidia_mega/qwen3_5/sm103/hop/__init__.py index 9d2e143a..035fe002 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/hop/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/hop/__init__.py @@ -229,8 +229,9 @@ def __init__( ) handle = symm_mem.rendezvous(self._mailbox, group) multicast = int(handle.multicast_ptr or 0) - if not multicast or multicast % 16: - raise RuntimeError("the hop needs an NVLink multicast mapping") + from sgl_nvidia_mega.qwen3_5.runtime import require_multicast + + require_multicast("hop", multicast, group, self.device) self._mailbox.fill_(SENTINEL) self._ctrl = torch.zeros(4, dtype=torch.int32, device=self.device) self._handle = handle # keeps the mapping alive diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 0ddbceac..7adc7741 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -112,6 +112,7 @@ def test_cp_slots_build_the_plugin_state_before_validation(monkeypatch): def test_cp_slots_refuse_an_incompatible_page_or_topology(monkeypatch): + monkeypatch.setenv(runtime.SLOTS_ENV, "skv2_cp") from sglang.srt.kernel_plugins.nvidia import binding as plugin_binding from sgl_nvidia_mega.qwen3_5.sm103.attention.q35_cp_lite import CpLite @@ -1867,3 +1868,134 @@ def test_prefill_unsupported_view_returns_the_original_before_state_writes(): original = lambda *a, **k: marker step = Step(mode="prefill", rows=33, view=SimpleNamespace(kernel_id=0)) assert built.forward(step, original, None, None, None) is marker + + +# The frozen C-arm recipe lists (A167 SLOT-INTERSECTIONS.json) and what the +# topology filters keep of them, as evaluated on parent 1a4e8884; the TP8 plan +# also adds route_lin and mkm_wide_bal. +_RECIPE = frozenset( + """ + route route_topk route_gemm route_shared route_lin mkm mkm_wide mkm_wide_bal + mkm_m32 mkm_m64 mkm_m128 mkm_m256 mkm_m512 gdn_verify gdn_pdl gdn_conv gdn_core + gdn_push gdn_front gdn_trt_front hop hop_mid hop_fp8 skv2 skv2_cp mka mka_push + lowm_in lowm_out lowm_se lowm_push multi_request sharded_gumbel draft_loop + draft_extend draft_multi draft_embed draft_moe draft_q1 draft_xchg draft_hop + draft_fin draft_front draft_skv2 draft_q1_cp draft_skv2_cp + """.split() +) +_TP8_RECIPE = _RECIPE - {"route_lin", "mkm_wide_bal"} +_WIDE_MKM = {"mkm_wide", "mkm_m32", "mkm_m64", "mkm_m128", "mkm_m256", "mkm_m512"} +_RECIPES = { + "TP2": ( + (2, 2, 1, 1, 1), + _RECIPE, + {"gdn_verify", "gdn_pdl", "gdn_conv", "gdn_core", "mka", "multi_request"} + | {"route", "route_topk", "route_gemm", "sharded_gumbel"}, + ), + "TP4": ( + (4, 4, 1, 1, 1), + _RECIPE, + _RECIPE + - _WIDE_MKM + - {"mkm_wide_bal", "gdn_front", "gdn_trt_front", "route_lin"}, + ), + "TP8": ((8, 8, 1, 1, 1), _TP8_RECIPE, _RECIPE), + "TEP8-EP4": ( + (8, 8, 1, 4, 1), + _TP8_RECIPE, + _TP8_RECIPE - _WIDE_MKM - {"mkm", "draft_moe"}, + ), + "DEP8": ( + (8, 1, 8, 8, 1), + _RECIPE, + {"gdn_verify", "gdn_pdl", "gdn_conv", "mka", "multi_request", "sharded_gumbel"}, + ), + "P-TP1-PP4": ( + (1, 1, 1, 1, 4), + {"gdn_prefill_split_conv"}, + {"gdn_prefill_split_conv"}, + ), +} + + +def _parallel(tp, attn_tp, dp, ep, pp): + return SimpleNamespace( + tp_size=tp, + attn_tp_size=attn_tp, + attn_dp_size=dp, + moe_ep_size=ep, + pp_size=pp, + tp_group=SimpleNamespace(world_size=tp), + ) + + +def _empty_model(): + return SimpleNamespace( + model=SimpleNamespace( + start_layer=0, end_layer=0, layers=[], flashinfer_mnnvl_cutedsl_fusion=None + ) + ) + + +@pytest.mark.parametrize("name", _RECIPES) +def test_explicit_recipe_lists_raise_nothing_and_keep_the_parent_filter( + monkeypatch, caplog, name +): + profile, requested, expected = _RECIPES[name] + monkeypatch.setenv(runtime.SLOTS_ENV, ",".join(sorted(requested))) + monkeypatch.setattr(runtime, "SLOTS", runtime._requested_slots()) + parallel = _parallel(*profile) + with caplog.at_level("INFO"): + if parallel.pp_size == 4: + monkeypatch.setattr( + "sglang.srt.runtime_context.get_parallel", lambda: parallel + ) + ctx = SimpleNamespace(tokens_per_request=7) + assert runtime.build(_empty_model(), ctx) is not None + else: + runtime.apply_expert_parallel(parallel.moe_ep_size) + runtime.apply_cp_lite_topology(parallel.attn_tp_size) + runtime.apply_topology(parallel) + assert runtime.slot_reasons() == [] + assert runtime.SLOTS == expected + for slot in requested - expected: + assert slot in caplog.text + + +@pytest.mark.parametrize("name", _RECIPES) +def test_the_default_plan_keeps_stock_for_unavailable_slots(monkeypatch, caplog, name): + from sgl_nvidia_mega.qwen3_5 import moe + + monkeypatch.delenv(runtime.SLOTS_ENV, raising=False) + monkeypatch.setattr(runtime, "SLOTS", runtime._requested_slots()) + parallel = _parallel(*_RECIPES[name][0]) + monkeypatch.setattr("sglang.srt.runtime_context.get_parallel", lambda: parallel) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(moe.route, "topk_unsupported", lambda model: "fixture router") + ctx = SimpleNamespace( + widths=(7, 14), + tokens_per_request=7, + deadline_s=60, + binding=SimpleNamespace(page_size=64, tp_size=parallel.tp_size), + ) + with caplog.at_level("INFO"): + built = runtime.build(_empty_model(), ctx) + assert built is not None + if parallel.pp_size == 1: + assert not any(s.route_topk or s.hop for s in (built.slots or {}).values()) + assert "stock" in caplog.text + + +def test_an_explicit_unavailable_slot_raises_with_its_recovery(monkeypatch): + from sgl_nvidia_mega.qwen3_5.draft import moe as draft_moe + + monkeypatch.setenv(runtime.SLOTS_ENV, "draft_loop,draft_moe") + monkeypatch.setattr( + draft_moe, "unsupported_reason", lambda mlp: "a TRTLLM expert layout" + ) + with pytest.raises(RuntimeError) as exc: + draft_moe.DraftMoe.build(object(), 0) + assert str(exc.value) == ( + "NVIDIA megakernel: slot draft_moe cannot run: a TRTLLM expert layout; " + "remove it from SGLANG_NVIDIA_MEGAKERNEL_SLOTS or change the topology" + ) From 13ef5036bdfa9fc187b32cc9edbfc29580af75af Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 2 Nov 2026 17:35:11 -0800 Subject: [PATCH 55/62] refactor: declare absent speculative hooks on built objects --- sgl_nvidia_mega/deepseek_v41/runtime/executor.py | 1 + sgl_nvidia_mega/kimi_k3/spec.py | 1 + sgl_nvidia_mega/qwen3_5/draft/__init__.py | 1 + 3 files changed, 3 insertions(+) diff --git a/sgl_nvidia_mega/deepseek_v41/runtime/executor.py b/sgl_nvidia_mega/deepseek_v41/runtime/executor.py index 1e073193..130a0682 100644 --- a/sgl_nvidia_mega/deepseek_v41/runtime/executor.py +++ b/sgl_nvidia_mega/deepseek_v41/runtime/executor.py @@ -113,6 +113,7 @@ class MegakernelDecode: draft_extend = None begin_draft_round = None graph_variants = None + speculative_runtime_hooks = None def __init__( self, diff --git a/sgl_nvidia_mega/kimi_k3/spec.py b/sgl_nvidia_mega/kimi_k3/spec.py index c67e3a7c..e6b571e6 100644 --- a/sgl_nvidia_mega/kimi_k3/spec.py +++ b/sgl_nvidia_mega/kimi_k3/spec.py @@ -169,6 +169,7 @@ class _Built: draft_extend = None begin_draft_round = None graph_variants = None + speculative_runtime_hooks = None def __init__(self, model: Any, resources: Any, widths: tuple[int, ...]): from .executor import KimiK3MegakernelExecutor diff --git a/sgl_nvidia_mega/qwen3_5/draft/__init__.py b/sgl_nvidia_mega/qwen3_5/draft/__init__.py index f9a6b624..4f4f7fe9 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/draft/__init__.py @@ -117,6 +117,7 @@ class MtpDraft: widths: tuple[int, ...] = () draft_sampler = None # SGLang's draft proposal + speculative_runtime_hooks = None def __init__( self, From e7ed0d7aadf9d9522f2627cb692f8971c3580d28 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 2 Nov 2026 18:12:35 -0800 Subject: [PATCH 56/62] perf(qwen3.5): keep stock for slot refusals that depend on the weights A refusal that depends on the weights, the model's shapes or the environment (draft MoE expert layout, MTP layers, LM-head/embedding shards, FlashInfer, multicast mapping, MKD memo parity) now logs and keeps the stock path even for an explicit SGLANG_NVIDIA_MEGAKERNEL_SLOTS list. Only a server-argument contradiction (CP page size, no covered width) still raises, and only for an explicitly requested slot. MKD build refusals fall back like its other checks. Co-Authored-By: Claude Opus 5.5 --- sgl_nvidia_mega/qwen3_5/draft/xchg.py | 29 ++++----- sgl_nvidia_mega/qwen3_5/runtime/__init__.py | 26 ++++---- test/unit/test_qwen3_5_runtime.py | 70 +++++++++++++++++---- 3 files changed, 88 insertions(+), 37 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/draft/xchg.py b/sgl_nvidia_mega/qwen3_5/draft/xchg.py index 381f8a2a..00c5d33d 100644 --- a/sgl_nvidia_mega/qwen3_5/draft/xchg.py +++ b/sgl_nvidia_mega/qwen3_5/draft/xchg.py @@ -30,7 +30,7 @@ import torch -from sgl_nvidia_mega.qwen3_5.runtime import slot_failure +from sgl_nvidia_mega.qwen3_5.runtime import SlotUnavailable, slot_failure from sgl_nvidia_mega.qwen3_5.sm103.draft import xchg as kernels logger = logging.getLogger(__name__) @@ -265,8 +265,8 @@ def __init__( @classmethod def maybe_create(cls, *, vocab_size: int, **kwargs) -> Optional[MkdBoundary]: - """The boundary, or None (logged) when ``unsupported_reason`` refuses it; every - rank decides alike.""" + """The boundary, or None (logged) when ``unsupported_reason`` or a build check + (multicast mapping, memo parity) refuses it; every rank decides alike.""" fc_weight = kwargs["fc_weight"] reason = unsupported_reason( hidden=int(fc_weight.shape[0]), @@ -276,12 +276,15 @@ def maybe_create(cls, *, vocab_size: int, **kwargs) -> Optional[MkdBoundary]: embed_shard=kwargs["embed_shard"], lm_head_shard=kwargs["lm_head_shard"], ) - if reason is not None: - logger.warning( - "MKD draft boundary disabled (%s); using the stock path.", reason - ) - return None - return cls(**kwargs) + if reason is None: + try: + return cls(**kwargs) + except SlotUnavailable as exc: + reason = exc.reason + logger.warning( + "MKD draft boundary disabled (%s); using the stock path.", reason + ) + return None # ---------------------------------------------------------------- memo table @@ -306,11 +309,9 @@ def _build_memo(self) -> torch.Tensor: bad_t = torch.tensor([bad, rows], dtype=torch.int64, device=emb.device) dist.all_reduce(bad_t, group=self._device_group) if int(bad_t[0]) != 0: - raise RuntimeError( - slot_failure( - "draft_xchg", - f"MKD memo table differs from the split fc on {int(bad_t[0])} of {int(bad_t[1])} rows", - ) + raise SlotUnavailable( + "draft_xchg", + f"MKD memo table differs from the split fc on {int(bad_t[0])} of {int(bad_t[1])} rows", ) logger.info( "MKD memo table: %d rows on this rank (%.1f MB), equal to the split fc on all %d rows.", diff --git a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py index 6ef7cf08..cdcaedd0 100644 --- a/sgl_nvidia_mega/qwen3_5/runtime/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/runtime/__init__.py @@ -155,20 +155,18 @@ def require_multicast(slots, pointer: int, group, device) -> None: ) -def unavailable_slots(slots, reason: str) -> None: - """Keep defaults optional at a capability check that was already strict. +def unavailable_slots(slots, reason: str, *, config: bool = False) -> None: + """Log a refused slot; the caller keeps its stock implementation / unserved entry. - Called at the capability check, before the unavailable implementation runs. - The caller keeps its existing stock implementation / unserved entry. + A refusal that depends on the weights, the model's shapes or the environment + never stops the server, even for an explicit SLOTS list. Only a ``config`` + contradiction (a server argument the slot cannot serve) raises, and only for an + explicitly requested slot. """ names = [slots] if isinstance(slots, str) else sorted(slots) - if names and os.environ.get(SLOTS_ENV) is not None: + if config and names and os.environ.get(SLOTS_ENV) is not None: raise RuntimeError(slot_failure(slots, reason)) - logger.info( - "NVIDIA megakernel: requested slots %s use stock (default policy; %s)", - names, - reason, - ) + logger.info("NVIDIA megakernel: requested slots %s use stock (%s)", names, reason) def slot_supported(slots, reason: str, *, cpu_group=None) -> bool: @@ -230,7 +228,9 @@ def build_cp_lite(binding) -> None: ) if binding.page_size != 64: unavailable_slots( - requested, f"CP slots need page size 64 (got {binding.page_size})" + requested, + f"CP slots need page size 64 (got {binding.page_size})", + config=True, ) return @@ -555,7 +555,9 @@ def build(model, ctx): widths = slot_plan.candidates(ctx.widths, ctx.tokens_per_request, multi) if not widths: unavailable_slots( - SLOTS, f"{ctx.tokens_per_request} rows a request leave no covered width" + SLOTS, + f"{ctx.tokens_per_request} rows a request leave no covered width", + config=True, ) return Qwen35Verify( model, {}, {}, (), binding=ctx.binding, summary="none (stock)" diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 7adc7741..611d1149 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -111,7 +111,9 @@ def test_cp_slots_build_the_plugin_state_before_validation(monkeypatch): assert calls == [(binding, torch.device("cuda", 3))] -def test_cp_slots_refuse_an_incompatible_page_or_topology(monkeypatch): +def test_cp_slots_refuse_an_incompatible_page_and_skip_an_unbuilt_topology( + monkeypatch, caplog +): monkeypatch.setenv(runtime.SLOTS_ENV, "skv2_cp") from sglang.srt.kernel_plugins.nvidia import binding as plugin_binding @@ -121,13 +123,20 @@ def test_cp_slots_refuse_an_incompatible_page_or_topology(monkeypatch): monkeypatch.setattr( plugin_binding, "megakernel_agree", lambda *args, **kwargs: None ) - with pytest.raises(RuntimeError, match="page size 64"): + with pytest.raises(RuntimeError) as exc: runtime.build_cp_lite(SimpleNamespace(page_size=32, tp_size=8)) + assert str(exc.value) == ( + "NVIDIA megakernel: slot skv2_cp cannot run: CP slots need page size 64 " + "(got 32); remove it from SGLANG_NVIDIA_MEGAKERNEL_SLOTS or change the topology" + ) + # The model's Q/KV heads, not a server argument: stock, even when explicit. monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) monkeypatch.setattr(CpLite, "maybe_create", lambda binding, device: None) - with pytest.raises(RuntimeError, match="Q/KV-head topology"): + with caplog.at_level("INFO"): runtime.build_cp_lite(SimpleNamespace(page_size=64, tp_size=8)) + assert "['skv2_cp'] use stock" in caplog.text + assert "Q/KV-head topology" in caplog.text @pytest.mark.parametrize("attention_tp_size", [1, 2, 3, 6, 32]) @@ -1986,16 +1995,55 @@ def test_the_default_plan_keeps_stock_for_unavailable_slots(monkeypatch, caplog, assert "stock" in caplog.text -def test_an_explicit_unavailable_slot_raises_with_its_recovery(monkeypatch): +@pytest.mark.parametrize( + "layout", ["use_flashinfer_trtllm_moe", "w13_swiglu_interleaved"] +) +def test_an_explicit_slot_keeps_stock_for_a_weight_layout(monkeypatch, caplog, layout): + from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod + from sgl_nvidia_mega.qwen3_5.draft import moe as draft_moe monkeypatch.setenv(runtime.SLOTS_ENV, "draft_loop,draft_moe") - monkeypatch.setattr( - draft_moe, "unsupported_reason", lambda mlp: "a TRTLLM expert layout" + method = object.__new__(UnquantizedFusedMoEMethod) + method.use_flashinfer_trtllm_moe = layout == "use_flashinfer_trtllm_moe" + method.w13_swiglu_interleaved = layout == "w13_swiglu_interleaved" + mlp = SimpleNamespace( + experts=SimpleNamespace(quant_method=method), + shared_expert=object(), + shared_expert_gate=object(), + layer_id=0, ) - with pytest.raises(RuntimeError) as exc: - draft_moe.DraftMoe.build(object(), 0) - assert str(exc.value) == ( - "NVIDIA megakernel: slot draft_moe cannot run: a TRTLLM expert layout; " - "remove it from SGLANG_NVIDIA_MEGAKERNEL_SLOTS or change the topology" + with caplog.at_level("INFO"): + assert draft_moe.DraftMoe.build(mlp, 0) is None + assert ( + "requested slots ['draft_moe'] use stock " + "(expert weights in a kernel-specific layout)" in caplog.text ) + + +def test_an_explicit_slot_keeps_stock_for_a_refused_mkd_build(monkeypatch, caplog): + from sgl_nvidia_mega.qwen3_5.draft import xchg + + monkeypatch.setenv(runtime.SLOTS_ENV, "draft_xchg") + monkeypatch.setattr(xchg, "unsupported_reason", lambda **kwargs: None) + + def refuse(self, **kwargs): + raise runtime.SlotUnavailable( + "draft_xchg", "needs an aligned NVLink multicast mapping on every rank" + ) + + monkeypatch.setattr(xchg.MkdBoundary, "__init__", refuse) + fc = SimpleNamespace(shape=(4096, 8192)) + with caplog.at_level("WARNING"): + assert ( + xchg.MkdBoundary.maybe_create( + vocab_size=8, + fc_weight=fc, + world=8, + rank=0, + embed_shard=None, + lm_head_shard=None, + ) + is None + ) + assert "MKD draft boundary disabled (needs an aligned NVLink" in caplog.text From 4190d316cb670fadbdfae6087d5c12c9fa4150c0 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 3 Nov 2026 05:35:52 -0800 Subject: [PATCH 57/62] perf(qwen3.5): fix the TEP8 routed body's NVFP4 MoE tactics per bucket The routed body asked the autotuner for each width's trtllm_fp4_block_scale_moe tactic, and the online cold tune's pick between near-tied tactics varies from boot to boot. moe/tactics.py now fixes the tactic of each token bucket for FlashInfer 0.7.1, SM103 and the routed experts' MoERunner signature; any other version, architecture, signature or bucket asks the autotuner as before, and the log names each width's source. Only the routed body reads the table. SGLANG_NVIDIA_MEGAKERNEL_MOE_TACTICS switches it; unset, it follows SGLANG_ENABLE_NV_MEGA_RUNTIME_OPTIMIZATION. Co-Authored-By: Claude Opus 5.5 --- README.md | 6 ++ sgl_nvidia_mega/qwen3_5/moe/body.py | 19 +++- sgl_nvidia_mega/qwen3_5/moe/tactics.py | 138 +++++++++++++++++++++++++ test/unit/test_qwen3_5_moe_tactics.py | 56 ++++++++++ 4 files changed, 214 insertions(+), 5 deletions(-) create mode 100644 sgl_nvidia_mega/qwen3_5/moe/tactics.py create mode 100644 test/unit/test_qwen3_5_moe_tactics.py diff --git a/README.md b/README.md index 97d0437e..8cf59827 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,12 @@ The same TP8 plan adds `mkm_wide_bal` for `mkm_m256` / `mkm_m512`: balanced CTA assignment at 129–512 verify rows, with bitwise-identical outputs. TEP8, DEP8, TP2 and TP4 run exactly the requested slots their topology supports. +TEP8's routed body takes its NVFP4 MoE tactic per token bucket from a fixed table +(`moe/tactics.py`) on FlashInfer 0.7.1, SM103 and the routed experts' MoERunner +signature; any other configuration or bucket asks the autotuner, and the log names +each width's source. `SGLANG_NVIDIA_MEGAKERNEL_MOE_TACTICS` switches the table; unset, +it follows `SGLANG_ENABLE_NV_MEGA_RUNTIME_OPTIMIZATION`. + The router-only `route`, `route_topk` and `route_gemm` slots support TP2, TP4, TP8 and TEP8 with the TRTLLM runner and no all-to-all. PD decode TP2 uses the same local row counts: batch 6/10 gives 42/70 verify rows. These fronts leave the expert GEMMs with the stock runner, so TP2 needs no fused expert engine. `route_shared` keeps its TP8/TEP8 projection bounds; DEP8 keeps stock MoE dispatch and routing. `mka` uses the fused projection kernels through 14 verify rows on TP8 and 8 on diff --git a/sgl_nvidia_mega/qwen3_5/moe/body.py b/sgl_nvidia_mega/qwen3_5/moe/body.py index 83930ab7..e0eb002f 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/body.py +++ b/sgl_nvidia_mega/qwen3_5/moe/body.py @@ -62,6 +62,12 @@ "MoeRunnerInputs", "alloc_trtllm_moe_output", ), + # moe/tactics.py: the table's key + "flashinfer/autotuner/autotuner.py": ("AutoTuner._find_nearest_profile",), + "flashinfer/fused_moe/backends/trtllm/sm100_runner.py": ( + "MoERunner.get_cache_key_extras", + ), + "flashinfer/version.py": ("__version__",), } OP = "flashinfer::trtllm_fp4_block_scale_moe" # the routed op's autotuner entry @@ -179,8 +185,8 @@ def default_tile(m: int, top_k: int, local_experts: int) -> int: def width(experts, x: torch.Tensor) -> Width: """The width of ``x`` [m, hidden]: on first use its runner, the routed op's tactic at - m rows (the autotuner's; without one, FlashInfer's default, [-1, -1]) and that tile_n's - tables (the body buffers follow on the first ``run``).""" + m rows (the plugin table's, else the autotuner's; without one, FlashInfer's default, + [-1, -1]) and that tile_n's tables (the body buffers follow on the first ``run``).""" from flashinfer.autotuner import AutoTuner from flashinfer.fused_moe.core import get_trtllm_moe_sm100_module from flashinfer.tllm_enums import ( @@ -191,6 +197,7 @@ def width(experts, x: torch.Tensor) -> Width: ) from sglang.srt.utils import next_power_of_2 + from sgl_nvidia_mega.qwen3_5.moe import tactics from sgl_nvidia_mega.qwen3_5.sm103.moe import router m = x.shape[0] @@ -233,9 +240,11 @@ def width(experts, x: torch.Tensor) -> Width: use_cold_l2_cache=True, use_cuda_graph=True, ) - _, tactic = AutoTuner.get().choose_one( - OP, [w.runner], config, inputs.to_list(), **kw - ) + tactic = tactics.pick(m, w.runner, config, inputs.to_list()) + if tactic is None: + _, tactic = AutoTuner.get().choose_one( + OP, [w.runner], config, inputs.to_list(), **kw + ) if isinstance(tactic, int) or tactic[0] <= 0 or tactic[1] < 0: # the launcher's fallback: the default tile_n and its default config w.tactic = [-1, -1] diff --git a/sgl_nvidia_mega/qwen3_5/moe/tactics.py b/sgl_nvidia_mega/qwen3_5/moe/tactics.py new file mode 100644 index 00000000..2593006f --- /dev/null +++ b/sgl_nvidia_mega/qwen3_5/moe/tactics.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""The TEP8 routed body's default FP4 MoE tactics, which ``body.width`` takes before +asking the autotuner. + +The routed op's online autotune measures each token bucket once at startup, and its +pick between near-tied tactics varies from boot to boot. This table fixes the +``flashinfer::trtllm_fp4_block_scale_moe`` tactic of each bucket for the one +configuration it was measured on: the FlashInfer version, the architecture and the +MoERunner signature (``get_cache_key_extras``) below. Any other version, architecture +or signature, and any bucket the table lacks, asks the autotuner as before. Only the +routed body reads the table; every other op keeps its own tuning. + +``ENV`` switches the table; unset, it follows ``PARENT_ENV``, SGLang's switch of the +megakernel runtime optimizations. +""" + +from __future__ import annotations + +import logging +import os +from typing import Optional, Union + +logger = logging.getLogger(__name__) + +ENV = "SGLANG_NVIDIA_MEGAKERNEL_MOE_TACTICS" +PARENT_ENV = "SGLANG_ENABLE_NV_MEGA_RUNTIME_OPTIMIZATION" + +FLASHINFER_VERSION = "0.7.1" +ARCH = "sm103" +# Qwen3.5's routed experts under TEP8: top-10 of 512 experts, 128 a rank, NVFP4 (E2m1) +# activations and weights, hidden 4096, intermediate 512, SwiGLU, shuffled MajorK. +SIGNATURE = ( + "Module", + 10, + 128, + 512, + 0, + 17826818, + 17826818, + 0, + 4096, + 4096, + 512, + 3, + True, + 0, + False, + False, +) +# The autotuner's token bucket -> [tile_n, config]; -1 is the launcher's default tile_n +# and config. Best of repeated cold-L2 CUDA-graph tunings on GB300. +TACTICS: dict[int, Union[int, tuple[int, int]]] = { + 8: -1, + 16: -1, + 32: -1, + 64: (8, 317), + 128: (16, 309), + 256: (16, 312), + 512: -1, + 768: (32, 371), + 1024: (64, 6), +} + +_TRUE = ("1", "true", "yes", "on") + + +def enabled() -> bool: + raw = os.environ.get(ENV) + if raw is None: + raw = os.environ.get(PARENT_ENV, "") + return raw.strip().lower() in _TRUE + + +def mismatch(version: str, arch: str, signature: tuple) -> str: + """Why the table does not apply to this configuration ("" when it does).""" + if version != FLASHINFER_VERSION: + return f"flashinfer {version}, the table's is {FLASHINFER_VERSION}" + if arch != ARCH: + return f"{arch}, the table's is {ARCH}" + if tuple(signature) != SIGNATURE: + return f"MoERunner signature {tuple(signature)} is not the table's" + return "" + + +def lookup( + m: int, bucket: int, version: str, arch: str, signature: tuple +) -> Optional[Union[int, list]]: + """The table's tactic at ``bucket`` (-1 or [tile_n, config]), or None: the + caller asks the autotuner. Logs the tactic's source either way.""" + if not enabled(): + reason = f"{ENV} / {PARENT_ENV} off" + else: + reason = mismatch(version, arch, signature) + if not reason and bucket not in TACTICS: + reason = f"no entry for bucket {bucket}" + if reason: + logger.info( + "NVIDIA megakernel: TEP8 routed body at %d rows (bucket %d): tactic " + "from the autotuner (%s)", + m, + bucket, + reason, + ) + return None + tactic = TACTICS[bucket] + logger.info( + "NVIDIA megakernel: TEP8 routed body at %d rows (bucket %d): tactic %s " + "from the plugin table (flashinfer %s, %s)", + m, + bucket, + tactic, + version, + arch, + ) + return tactic if isinstance(tactic, int) else list(tactic) + + +def pick(m: int, runner, config, inputs: list) -> Optional[Union[int, list]]: + """``lookup`` for ``runner`` at the autotuner's bucket of ``inputs``.""" + import flashinfer + import torch + from flashinfer.autotuner import AutoTuner + + shapes = tuple( + tuple(t.size()) if isinstance(t, torch.Tensor) else (0,) for t in inputs + ) + spec = config.dynamic_tensor_specs[0] + profile = AutoTuner._find_nearest_profile(shapes, config) + bucket = profile[spec.input_idx[0]][spec.dim_idx[0]] + major, minor = torch.cuda.get_device_capability() + return lookup( + m, + bucket, + flashinfer.__version__, + f"sm{major}{minor}", + runner.get_cache_key_extras(inputs), + ) diff --git a/test/unit/test_qwen3_5_moe_tactics.py b/test/unit/test_qwen3_5_moe_tactics.py new file mode 100644 index 00000000..cb11cd93 --- /dev/null +++ b/test/unit/test_qwen3_5_moe_tactics.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""The TEP8 routed body's default FP4 MoE tactic table on CPU: a hit, the autotuner +on any mismatch, and its switch.""" + +import pytest + +from sgl_nvidia_mega.qwen3_5.moe import tactics + +KEY = (tactics.FLASHINFER_VERSION, tactics.ARCH, tactics.SIGNATURE) + + +@pytest.fixture(autouse=True) +def switch(monkeypatch): + monkeypatch.delenv(tactics.ENV, raising=False) + monkeypatch.setenv(tactics.PARENT_ENV, "1") + + +def test_a_bucket_in_the_table_takes_its_tactic(): + assert tactics.lookup(100, 128, *KEY) == [16, 309] + assert tactics.lookup(1000, 1024, *KEY) == [64, 6] + # the launcher's default stays an int, as the autotuner returns it + assert tactics.lookup(8, 8, *KEY) == -1 + + +@pytest.mark.parametrize( + "version, arch, signature, bucket", + [ + ("0.7.2", tactics.ARCH, tactics.SIGNATURE, 128), + (tactics.FLASHINFER_VERSION, "sm100", tactics.SIGNATURE, 128), + ( + tactics.FLASHINFER_VERSION, + tactics.ARCH, + ("Module", 8) + tactics.SIGNATURE[2:], + 128, + ), + (*KEY, 2048), + ], + ids=["version", "arch", "signature", "bucket"], +) +def test_anything_else_asks_the_autotuner(version, arch, signature, bucket): + assert tactics.lookup(bucket, bucket, version, arch, signature) is None + + +def test_the_switch_follows_the_runtime_optimizations_unless_set(monkeypatch): + assert tactics.enabled() + monkeypatch.setenv(tactics.ENV, "0") + assert not tactics.enabled() + assert tactics.lookup(100, 128, *KEY) is None + monkeypatch.setenv(tactics.PARENT_ENV, "0") + monkeypatch.setenv(tactics.ENV, "1") + assert tactics.lookup(100, 128, *KEY) == [16, 309] + monkeypatch.delenv(tactics.ENV) + assert not tactics.enabled() + monkeypatch.delenv(tactics.PARENT_ENV) + assert not tactics.enabled() From a1a418e42f100a82d55ec2759150c655f87e8978 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 3 Nov 2026 21:47:50 -0800 Subject: [PATCH 58/62] perf(qwen3.5): fuse eligible TEP8 shared activation quantization --- sgl_nvidia_mega/qwen3_5/moe/__init__.py | 42 +++++++++++++- sgl_nvidia_mega/qwen3_5/moe/route.py | 6 +- test/unit/test_qwen3_5_steps.py | 76 ++++++++++++++++++++++++- 3 files changed, 117 insertions(+), 7 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/moe/__init__.py b/sgl_nvidia_mega/qwen3_5/moe/__init__.py index 5cb94e69..e1b4d382 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/__init__.py +++ b/sgl_nvidia_mega/qwen3_5/moe/__init__.py @@ -63,7 +63,12 @@ ), "kernels/ops/elementwise/elementwise.py": ("fused_gate_sigmoid_mul_add",), "srt/model_executor/runner/__init__.py": ("get_is_capture_mode",), - "srt/runtime_context.py": ("get_parallel", '"tp_size"', '"moe_ep_size"'), + "srt/runtime_context.py": ( + "get_parallel", + '"tp_size"', + '"attn_dp_size"', + '"moe_ep_size"', + ), } # The SGLang functions whose body this module mirrors, by the digest of their code @@ -150,7 +155,38 @@ def build(model, ctx) -> route.Bound: router.build_big( torch.cuda.current_device(), mlp.experts.num_local_experts ) - return bound + return _build_shared_act_quant(model, bound) + + +def _build_shared_act_quant(model, bound): + """Prebuild the ordinary TEP8 shared MLP's tail, independently of route_shared.""" + import msgspec + from sglang.srt.runtime_context import get_parallel + + if not bound.pre: + return bound + p = get_parallel() + if (p.tp_size, p.attn_dp_size) != (8, 1) or p.moe_ep_size <= 1: + return bound + from sgl_nvidia_mega.qwen3_5.sm103.moe import shared_act_quant + + m = model.model + layers = [m.layers[i].mlp for i in range(m.start_layer, m.end_layer)] + if not layers or any(route.shared_unsupported(mlp) for mlp in layers): + return bound + for columns in sorted( + {int(mlp.shared_expert.gate_up_proj.weight.shape[1]) for mlp in layers} + ): + shared_act_quant.build(torch.cuda.current_device(), columns) + return msgspec.structs.replace(bound, shared_act_quant=True) + + +def _shared_expert(mlp, x, st): + """Keep the stock projection and gate; fuse only the eligible activation tail.""" + if st.bound["moe"].shared_act_quant: + gate_up, _ = mlp.shared_expert.gate_up_proj(x) + return route.shared_expert_tail(mlp.shared_expert, gate_up) + return mlp._forward_shared_experts(x, apply_gate=False) def step(layer, x: torch.Tensor, st, defer: bool, reduce=None): @@ -238,7 +274,7 @@ def _dual_stream(mlp, x, use_fused_gate: bool, defer: bool, st): x8 = lowm.take_fp8(st, x) if st.slots.hop_fp8 else None shared = lowm.shared_expert(mlp, x, st.bound["lowm"], x8) else: - shared = mlp._forward_shared_experts(x, apply_gate=False) + shared = _shared_expert(mlp, x, st) if shared is not None: shared = mlp._gate_shared_output_out_of_place(x, shared) current.wait_stream(mlp.alt_stream) diff --git a/sgl_nvidia_mega/qwen3_5/moe/route.py b/sgl_nvidia_mega/qwen3_5/moe/route.py index 1551c43c..8adec94a 100644 --- a/sgl_nvidia_mega/qwen3_5/moe/route.py +++ b/sgl_nvidia_mega/qwen3_5/moe/route.py @@ -142,6 +142,7 @@ class Bound(msgspec.Struct, frozen=True): # [160] int32 0, 1, ...: slot mkm's rows are in expert order identity: object = None shared: bool = False # slot route_shared: the shared gate/up in the router launch + shared_act_quant: bool = False # TEP8's ordinary shared MLP uses the fused tail topk: bool = False # slot route_topk: the top-10 + NVFP4 input launch built gemm: bool = ( False # slot route_gemm: the router front's 16-row tiles stand in for the GEMM @@ -430,12 +431,11 @@ def shared_expert_tail(expert, gate_up: torch.Tensor) -> torch.Tensor: quantized = silu_and_mul_static_quant_fp8( gate_up, expert.down_proj.input_scale.reshape(1) ) - key = ("route_shared_act_quant", gate_up.shape[0]) + key = ("shared_act_quant", gate_up.shape[0]) if key not in _ENGAGED: _ENGAGED.add(key) logger.info( - "NVIDIA megakernel: Qwen3.5 slot route_shared shared_act_quant " - "runs at %d rows", + "NVIDIA megakernel: Qwen3.5 shared_act_quant runs at %d rows", gate_up.shape[0], ) x, _ = expert.down_proj(quantized) diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 2073ac93..beaab3b2 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -269,7 +269,80 @@ def test_empty_dp_rank_skips_gate_and_shared_but_runs_experts(capture): def _st(route=False): - return SimpleNamespace(slots=plan.Slots(route=route)) + return SimpleNamespace( + slots=plan.Slots(route=route), bound={"moe": moe.route.Bound()} + ) + + +@pytest.mark.parametrize( + "topology,pre,reason,want", + [ + ((8, 1, 8), True, "", True), + ((8, 1, 4), True, "", True), + ((8, 1, 8), False, "", False), + ((8, 1, 1), True, "", False), + ((8, 8, 8), True, "", False), + ((4, 1, 4), True, "", False), + ((8, 1, 8), True, "BF16 down projection", False), + ], +) +def test_ordinary_shared_act_quant_builds_only_for_eligible_tep8( + monkeypatch, topology, pre, reason, want +): + from sglang.srt import runtime_context + + from sgl_nvidia_mega.qwen3_5.sm103.moe import shared_act_quant + + monkeypatch.setattr( + runtime_context, + "get_parallel", + lambda: SimpleNamespace( + tp_size=topology[0], attn_dp_size=topology[1], moe_ep_size=topology[2] + ), + ) + expert = SimpleNamespace( + gate_up_proj=SimpleNamespace(weight=_CudaMeta((4096, 256), torch.float8_e4m3fn)) + ) + mlp = SimpleNamespace(shared_expert=expert) + model = SimpleNamespace( + model=SimpleNamespace( + start_layer=0, end_layer=2, layers=[SimpleNamespace(mlp=mlp)] * 2 + ) + ) + monkeypatch.setattr(moe.route, "shared_unsupported", lambda _: reason) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + built = [] + monkeypatch.setattr(shared_act_quant, "build", lambda *args: built.append(args)) + bound = moe._build_shared_act_quant(model, moe.route.Bound(pre=pre)) + assert bound.shared_act_quant is want + assert built == ([(0, 256)] if want else []) + + +@pytest.mark.parametrize("capture", [False, True]) +@pytest.mark.parametrize("enabled", [False, True]) +def test_ordinary_shared_act_quant_preserves_the_gate_and_streams( + monkeypatch, capture, enabled +): + x = torch.zeros(448, 4) + log = _log(x=x) + mlp = _moe_block(log, x) + mlp.shared_expert.gate_up_proj = log.op("shared_up", ("gate_up", None)) + monkeypatch.setattr( + moe.route, + "shared_expert_tail", + lambda se, gu: log.op("act_quant_tail", "shared_out")(gu), + ) + st = _st() + st.bound["moe"] = moe.route.Bound(shared_act_quant=enabled) + with _streams(log, capture), _moe_kernels(log): + assert moe.step(SimpleNamespace(mlp=mlp), x, st, defer=True) == "handoff" + names = [e[0] for e in log] + assert ("act_quant_tail" in names) == (capture and enabled) + assert ("shared" in names) == (not (capture and enabled)) + assert names.count("gate_mul") == 1 + assert [e for e in log if e[0] == "wait"] == ( + [("wait", "alt", "main"), ("wait", "main", "alt")] if capture else [] + ) def _router_front(log, mlp): @@ -818,6 +891,7 @@ def combine(combine_input, log=log, combined=combined): mkm={3: "maps"}, identity=torch.arange(70, dtype=torch.int32), m16_max_rows=16, + shared_act_quant=False, ) st = SimpleNamespace(slots=slots, bound={"moe": bound}) fused = patch.object( From f60f99035a31a1d025f393d295c122d04d3bf066 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 3 Nov 2026 23:35:31 -0800 Subject: [PATCH 59/62] perf(qwen3.5): use the GDN three-CTA cap for one-wave grids --- sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py | 6 +++--- test/unit/test_qwen3_5_steps.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py b/sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py index 7b74ea03..4fdf4ac2 100644 --- a/sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py +++ b/sgl_nvidia_mega/qwen3_5/sm103/gdn/core.py @@ -155,9 +155,9 @@ def _sms(device_index: int) -> int: def _min_blocks(device_index: int, ctas: int) -> int: - """Three CTAs an SM (168 registers) once two an SM take more than two waves (128 requests - at TP8): past that the third resident CTA wins; below it the uncapped registers do.""" - return 3 if ctas > 4 * _sms(device_index) else 1 + """Cap registers where three CTAs/SM fit one wave or save later waves.""" + sms = _sms(device_index) + return 3 if 2 * sms < ctas <= 3 * sms or ctas > 4 * sms else 1 def _slot_stride(t: Optional[torch.Tensor]) -> int: diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index beaab3b2..0158fa6a 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -1888,6 +1888,25 @@ def test_the_gdn_front_leaves_out_its_l2_prefetch_ahead_of_a_pinned_or_early_cor assert seen["prefetch"] is prefetch +def test_the_gdn_core_caps_registers_for_the_three_cta_wave(): + from sgl_nvidia_mega.qwen3_5.sm103.gdn import core + + with patch.object(core, "_sms", return_value=148): + for batch in range(1, 129): + ctas = batch * 8 + old = 3 if ctas > 4 * 148 else 1 + assert core._min_blocks(0, ctas) == (3 if 38 <= batch <= 55 else old) + for ctas, expected in ( + (296, 1), + (297, 3), + (444, 3), + (445, 1), + (592, 1), + (593, 3), + ): + assert core._min_blocks(0, ctas) == expected + + @pytest.mark.parametrize("trt, wants", [(False, False), (True, True)]) def test_the_hop_quantizes_a_gdn_layers_input_for_the_trt_front_only(trt, wants): """Slot hop_fp8 writes the input consumed by slot gdn_trt_front.""" From 5edd8e993a16b9b9c41cc801fb99c6e86084d483 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Wed, 4 Nov 2026 13:22:29 -0800 Subject: [PATCH 60/62] fix(qwen3.5): audit GDN mirrors after padding-sink staging Upstream sgl-project/sglang#48004 stages strided prefill state at rows 1..n and reserves row 0 for padding. Keep the existing capability boundary: strided prefill returns to the native backend before writes, contiguous prefill keeps native indices, and target verification never stages prefill state. Exercise the native forward branch with CPU kernel stand-ins, compare contiguous request/state identity, and reject stage/unstage calls during target verification. Update the two mirror digests after these checks; no executable kernel or serving path changes. --- sgl_nvidia_mega/qwen3_5/gdn/prefill.py | 5 +- sgl_nvidia_mega/qwen3_5/gdn/verify.py | 2 + test/unit/test_qwen3_5_runtime.py | 203 +++++++++++++++++++++++++ test/unit/test_qwen3_5_steps.py | 16 ++ 4 files changed, 224 insertions(+), 2 deletions(-) diff --git a/sgl_nvidia_mega/qwen3_5/gdn/prefill.py b/sgl_nvidia_mega/qwen3_5/gdn/prefill.py index 379559c8..039fa6b9 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/prefill.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/prefill.py @@ -75,11 +75,12 @@ } SGLANG_MIRRORS = { - # v3 routes committed state via STATE_PART. LinearPrefillView already carries - # those exact committed conv/temporal buffers; no recurrence/order change. + # Upstream #48004 reserves staging row 0 for padding; this slot rejects strided + # pools before writes and passes the native pool indices unchanged. "srt/layers/attention/linear/gdn_backend.py:GDNAttnBackend.forward_extend": ( "12a835fc4f78af33", "6ed4eb6a818f7e20", + "bfd997695f32d409", ), "srt/layers/attention/hybrid_linear_attn_backend.py:MambaAttnBackendBase._track_mamba_state_extend": "84c447ca707eba64", } diff --git a/sgl_nvidia_mega/qwen3_5/gdn/verify.py b/sgl_nvidia_mega/qwen3_5/gdn/verify.py index 44ae3d31..17ade29d 100644 --- a/sgl_nvidia_mega/qwen3_5/gdn/verify.py +++ b/sgl_nvidia_mega/qwen3_5/gdn/verify.py @@ -135,10 +135,12 @@ SGLANG_MIRRORS = { # Without and with the early return of the GDN extend in breakable prefill graphs, # which target verify never takes; the first goes once all supported SGLangs have it + # Upstream #48004 changes strided prefill staging, which target verify never takes. "srt/layers/attention/linear/gdn_backend.py:GDNAttnBackend.forward_extend": ( "ab7b048fafb07d0e", "12a835fc4f78af33", "6ed4eb6a818f7e20", + "bfd997695f32d409", ), # Importing the ReplaySSM kernel from fla/ and, after sgl-project/sglang#45112, from # linear/; the first goes once all supported SGLangs have the move diff --git a/test/unit/test_qwen3_5_runtime.py b/test/unit/test_qwen3_5_runtime.py index 611d1149..8ac75d61 100644 --- a/test/unit/test_qwen3_5_runtime.py +++ b/test/unit/test_qwen3_5_runtime.py @@ -1879,6 +1879,209 @@ def test_prefill_unsupported_view_returns_the_original_before_state_writes(): assert built.forward(step, original, None, None, None) is marker +def _gdn_prefill_state_case(monkeypatch, strided): + from sglang.srt.layers.attention.linear import gdn_backend + + indices = torch.tensor([5, 2, -1, 0], dtype=torch.int32) + conv = torch.zeros(1, 8, 6, 6)[..., ::2] + temporal = torch.zeros(1, 8, 1, 2, 4)[..., ::2] + if strided != "conv": + conv = conv.contiguous() + if strided != "temporal": + temporal = temporal.contiguous() + for row in range(1, 8): + conv[:, row].fill_(row) + temporal[:, row].fill_(row) + view = SimpleNamespace( + kernel_id=1, + is_graph=False, + has_mis=False, + has_track=False, + state=SimpleNamespace(layer_ids=(0,), conv=(conv,), temporal=temporal), + state_indices=indices, + query_start_loc=torch.arange(5, dtype=torch.int32), + conv_block_table=None, + track_ssm_recompute_dst=None, + ) + fb = SimpleNamespace( + forward_mode=SimpleNamespace(is_target_verify=lambda: False), + extend_prefix_lens=torch.ones(4, dtype=torch.int32), + extend_seq_lens_cpu=[1] * 4, + ) + layer = SimpleNamespace( + layer_id=0, + num_q_heads=1, + num_k_heads=1, + num_v_heads=1, + head_q_dim=2, + head_k_dim=2, + head_v_dim=2, + q_dim=2, + k_dim=2, + v_dim=2, + conv_weights=torch.zeros(6, 4), + bias=None, + activation="silu", + A_log=torch.zeros(1), + dt_bias=torch.zeros(1), + ) + calls = [] + + def update(state, rows, amount): + calls.append(rows.clone()) + for row in rows.tolist(): + if row > 0: + state[row].add_(amount) + + def convolve(x, weight, bias, **kw): + update(kw["conv_states"], kw["cache_indices"], 1) + return x + + def recurrence(**kw): + update(kw["ssm_states"], kw["cache_indices"], 2) + return kw["v"], None, None + + monkeypatch.setattr(gdn_backend, "causal_conv1d_fn", convolve) + for name in ("is_cpu", "is_cuda", "is_hip", "is_xpu", "is_npu"): + monkeypatch.setattr(gdn_backend, name, lambda: False) + meta = SimpleNamespace( + query_start_loc=view.query_start_loc, + mamba_cache_indices=indices, + retrieve_next_token=None, + retrieve_next_sibling=None, + retrieve_parent_token=None, + conv_block_table=None, + has_mamba_track_mask=False, + state_checkpoint_cu_starts=None, + num_state_checkpoints=None, + state_checkpoint_every_n_tokens=None, + ) + backend = SimpleNamespace( + prefill_graph_metadata=None, + forward_metadata=meta, + mis_metadata=None, + requires_contiguous_prefill_state=True, + state_layer_cache=lambda _: SimpleNamespace( + conv=(conv[0],), temporal=temporal[0] + ), + kernel_dispatcher=SimpleNamespace(extend=recurrence), + _prefill_gates=lambda *args: (torch.zeros(1, 4, 1), torch.zeros(1, 4, 1)), + ) + backend._convolve_prefill = lambda *args: ( + gdn_backend.GDNAttnBackend._convolve_prefill(backend, *args) + ) + + def native(): + return gdn_backend.GDNAttnBackend.forward_extend( + backend, layer, fb, torch.zeros(4, 6), torch.zeros(4, 1), torch.zeros(4, 1) + ) + + return view, fb, layer, calls, native, convolve, recurrence + + +@pytest.mark.parametrize("strided", ["conv", "temporal"]) +def test_strided_prefill_falls_back_before_writes_and_preserves_the_sink( + monkeypatch, strided +): + from sgl_nvidia_mega.qwen3_5.runtime.prefill import Qwen35Prefill + + view, fb, _, calls, native, _, _ = _gdn_prefill_state_case(monkeypatch, strided) + conv, temporal = view.state.conv[0], view.state.temporal + before = (conv.clone(), temporal.clone()) + built = Qwen35Prefill( + SimpleNamespace(model=SimpleNamespace(start_layer=0, end_layer=0, layers=[])), + SimpleNamespace(tokens_per_request=7), + True, + ) + monkeypatch.setattr( + built, "_forward", lambda *args: pytest.fail("strided plugin write") + ) + entered = [] + + def original(*args, **kwargs): + assert calls == [] + torch.testing.assert_close(conv, before[0]) + torch.testing.assert_close(temporal, before[1]) + entered.append(True) + return native() + + built.forward(Step(mode="prefill", rows=4, view=view), original, None, None, fb) + assert entered == [True] + assert len(calls) == 2 + assert all(torch.equal(rows, torch.tensor([1, 2, 0, 0])) for rows in calls) + for state, old, amount in ((conv, before[0], 1), (temporal, before[1], 2)): + expected = old.clone() + expected[:, 5].add_(amount) + expected[:, 2].add_(amount) + torch.testing.assert_close(state, expected) + + +def test_contiguous_prefill_keeps_main_request_indices_without_zero_based_staging( + monkeypatch, +): + from sglang.srt.layers.attention.linear.kernels import state_pool_io + + from sgl_nvidia_mega.qwen3_5.gdn import prefill + + view, fb, layer, calls, native, convolve, recurrence = _gdn_prefill_state_case( + monkeypatch, None + ) + conv, temporal = view.state.conv[0], view.state.temporal + before = (conv.clone(), temporal.clone()) + native_out = native() + expected = (conv.clone(), temporal.clone()) + conv.copy_(before[0]) + temporal.copy_(before[1]) + calls.clear() + for name in ("stage_pool_states", "unstage_pool_states"): + monkeypatch.setattr( + state_pool_io, name, lambda *args: pytest.fail("unexpected staging") + ) + monkeypatch.setattr( + "sglang.kernels.ops.mamba.causal_conv1d_triton.causal_conv1d_fn", convolve + ) + monkeypatch.setattr( + "sglang.kernels.ops.attention.triton_gdn_fused_proj.qwen3_5_gdn_prefill_projection_views", + lambda *args: ( + torch.zeros(4, 6), + torch.zeros(4, 1, 2), + torch.zeros(4, 1), + torch.zeros(4, 1), + ), + ) + monkeypatch.setattr( + "sglang.kernels.ops.attention.linear.fused_gdn_gating.fused_gdn_gating", + lambda *args: (torch.zeros(1, 4, 1), torch.zeros(1, 4, 1)), + ) + monkeypatch.setattr( + "sglang.kernels.ops.attention.fla.chunk.chunk_gated_delta_rule", + lambda **kw: recurrence( + ssm_states=kw["initial_state"], + cache_indices=kw["initial_state_indices"], + v=kw["v"], + ), + ) + gdn = SimpleNamespace( + attn=layer, + num_k_heads=1, + num_v_heads=1, + attn_tp_size=1, + head_k_dim=2, + head_v_dim=2, + _forward_input_proj=lambda _: (torch.zeros(4, 8), torch.zeros(4, 2)), + norm=lambda core, z: core, + out_proj=lambda core: (core, None), + ) + assert prefill.unsupported(view) == "" + got = prefill.step(gdn, torch.zeros(4, 2), view, fb) + torch.testing.assert_close(got, native_out.reshape(4, 2)) + torch.testing.assert_close(conv, expected[0]) + torch.testing.assert_close(temporal, expected[1]) + assert len(calls) == 4 + assert all(torch.equal(rows, view.state_indices) for rows in calls) + assert calls[0][0].item() == 5 + + # The frozen C-arm recipe lists (A167 SLOT-INTERSECTIONS.json) and what the # topology filters keep of them, as evaluated on parent 1a4e8884; the TP8 plan # also adds route_lin and mkm_wide_bal. diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 0158fa6a..75f82e2a 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -1390,6 +1390,22 @@ def symmetric(group, disabled=False): yield with ( + patch( + "sglang.srt.layers.attention.linear.kernels.state_pool_io.stage_pool_states", + side_effect=AssertionError("target verify must not stage prefill state"), + ), + patch( + "sglang.srt.layers.attention.linear.kernels.state_pool_io.unstage_pool_states", + side_effect=AssertionError("target verify must not unstage prefill state"), + ), + patch( + "sglang.srt.layers.attention.linear.gdn_backend.stage_pool_states", + side_effect=AssertionError("target verify must not stage prefill state"), + ), + patch( + "sglang.srt.layers.attention.linear.gdn_backend.unstage_pool_states", + side_effect=AssertionError("target verify must not unstage prefill state"), + ), patch( "sglang.kernels.ops.attention.triton_gdn_fused_proj.qwen3_5_gdn_prefill_projection_views", views, From d04aa78fed1b5355348c59b0e8223f931fdf2496 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Wed, 4 Nov 2026 16:18:17 -0800 Subject: [PATCH 61/62] [CI] test: Pin SGLang to a main commit with the speculative runtime hooks sgl-project/sglang#47201 is merged; the runtime tests import HooksFactory from it. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/unit-test.yml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 33b3cca8..2c3cd823 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -14,10 +14,11 @@ concurrency: cancel-in-progress: true env: - # Includes the hybrid preparation/views API from sgl-project/sglang#45406. + # Includes the hybrid preparation/views API from sgl-project/sglang#45406 and the + # speculative runtime hooks from sgl-project/sglang#47201. # Pin source until a published nightly contains this interface; updating the # pin is a normal PR and must pass the static mirror/contract tests. - SGLANG_COMMIT: aa3e34aaebf880d72482cea74dbbfef54da728e2 + SGLANG_COMMIT: 8df0cd51e1afbcf2a67358ab559f4a9046ff5782 jobs: unit: From b5c85d93abf7f84a90ef9e7e7acb3b3b2fd8320a Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Wed, 4 Nov 2026 16:26:25 -0800 Subject: [PATCH 62/62] test(qwen3.5): check the prequantized GDN out_proj on the plugin alone The test compared against an SGLang FP8 reference keyword that main does not have. Check the plugin path directly: FP32 raises before any BMM, and BF16/FP16 call the BMM once with the layer dtype on the same input, weight and scales. Co-Authored-By: Claude Opus 5.5 --- test/unit/test_qwen3_5_steps.py | 28 +++++++--------------------- 1 file changed, 7 insertions(+), 21 deletions(-) diff --git a/test/unit/test_qwen3_5_steps.py b/test/unit/test_qwen3_5_steps.py index 75f82e2a..f00c2fb0 100644 --- a/test/unit/test_qwen3_5_steps.py +++ b/test/unit/test_qwen3_5_steps.py @@ -1535,34 +1535,20 @@ def bmm(x, weight, x_scale, weight_scale, output_dtype): ) if model_dtype == torch.float32: - with pytest.raises(ValueError, match="explicit BF16 or FP16"): - fp8_utils.apply_fp8_linear_bmm_flashinfer( - fp8, - layer.weight, - layer.weight_scale, - layer.input_scale, - pre_quant_output_dtype=model_dtype, - ) with pytest.raises(ValueError, match="explicit BF16 or FP16"): verify._out_proj(layer, fp8, recurrence_dtype) assert calls == [] return - expected = fp8_utils.apply_fp8_linear_bmm_flashinfer( - fp8, - layer.weight, - layer.weight_scale, - layer.input_scale, - pre_quant_output_dtype=model_dtype, - ) actual = verify._out_proj(layer, fp8, recurrence_dtype) - assert actual.dtype == expected.dtype == model_dtype - assert torch.equal(actual, expected) - assert len(calls) == 2 - assert calls[0][-1] == calls[1][-1] == model_dtype - assert calls[0][0].data_ptr() == calls[1][0].data_ptr() == fp8.data_ptr() - assert all(left is right for left, right in zip(calls[0][1:4], calls[1][1:4])) + assert actual.dtype == model_dtype + assert len(calls) == 1 + x, weight, x_scale, weight_scale, output_dtype = calls[0] + assert output_dtype == model_dtype + assert x.data_ptr() == fp8.data_ptr() + assert weight is layer.weight + assert x_scale is layer.input_scale and weight_scale is layer.weight_scale def test_gdn_bf16_out_proj_can_use_lowm(monkeypatch):