[ROCm][Perf][DSv4.1] Mono decode for the DSpark draft layers

tests/models/test_dsv41_mono_decode.py

@@ -22,9 +22,10 @@
 class _Runner:
     def __init__(self):
         self.calls = []
+        self.draft = None
 
     def forward(self, *args, **kwargs):
-        self.calls.append((args, kwargs))
+        self.calls.append((args, {**kwargs, "draft_runner": self.draft}))
         return ("mono outputs",)
 
     def first_pre_mix(self, rows):
@@ -86,6 +87,7 @@ def _run(
         ffn_only=False,
         entry=None,
         layer=None,
+        mono=None,
         **inputs,
     ):
         if metadata is None:
@@ -97,10 +99,15 @@ def _run(
         runner = _Runner()
         monkeypatch.setattr(md, "is_forward_context_available", lambda: True)
         monkeypatch.setattr(md, "get_forward_context", lambda: fc)
-        monkeypatch.setattr(md, "_mono_runner", lambda device: runner)
+
+        def pick(device, draft=None):
+            runner.draft = draft
+            return runner
+
+        monkeypatch.setattr(md, "_mono_runner", pick)
         monkeypatch.setattr(md.MonoDecodeLayer, "weights", lambda self, layer: None)
         args = {**_inputs(), **inputs}
-        mono = md.MonoDecodeLayer(ffn_only)
+        mono = mono or md.MonoDecodeLayer(ffn_only)
         if (entry or ("ffn" if ffn_only else "layer")) == "ffn":
             # x stands for wo_b's unreduced output; the first layer's attention
             # has no compressed cache
@@ -135,6 +142,49 @@ def test_first_layer_runs_mono_from_the_embedding(run):
     assert args[2].shape == (M, 4, 5120) and not kwargs["fold"]
 
 
+DRAFT_WIDTH = 192  # the DSpark draft's non-causal SWA rows: window 128 + 5 drafts
+
+
+def _draft(**overrides):
+    """A DSpark draft layer's mono path, and the metadata its step carries."""
+    layer = _layer(comp=None)
+    layer.attn.compress_ratio = 0
+    swa = _swa(
+        decode_swa_width=DRAFT_WIDTH,
+        decode_swa_indices=torch.zeros(64, 1, DRAFT_WIDTH, dtype=torch.int32),
+    )
+    mono = md.MonoDecodeLayer(draft=True, experts=128, topk=3, swa_width=DRAFT_WIDTH)
+    return dict(layer=layer, mono=mono, metadata={"swa": swa}, **overrides)
+
+
+def test_draft_layer_runs_mono_on_its_own_runner(run):
+    out, calls = run(**_draft())
+    assert out == ("mono outputs",) and len(calls) == 1
+    assert calls[0][0][9].stride(0) == DRAFT_WIDTH  # the non-causal SWA rows
+    assert calls[0][1]["draft_runner"] is not None  # not the target's runner
+
+
+def test_draft_first_layer_takes_the_broadcast_embedding(run):
+    """The draft hands its first layer the embedding already repeated to the
+    hc copies: K1 takes it as the settled residual."""
+    x = torch.zeros(M, 4, 5120, dtype=torch.bfloat16)
+    out, calls = run(
+        **_draft(residual=None, post_mix=None, res_mix=None, pre_mix=None, x=x)
+    )
+    assert out == ("mono outputs",)
+    args, kwargs = calls[0]
+    assert args[1].shape == (M, 5120) and args[2] is x and not kwargs["fold"]
+
+
+def test_draft_and_target_layers_take_only_their_own_swa_rows(run):
+    causal = {"swa": _swa(), "comp": SimpleNamespace(block_size=128)}
+    out, calls = run(**{**_draft(), "metadata": causal})
+    assert out is None and not calls
+    target = {"swa": _draft()["metadata"]["swa"], "comp": causal["comp"]}
+    out, calls = run(metadata=target)
+    assert out is None and not calls
+
+
 def test_index_source_selects_between_the_kernels(run):
     """An index source hands the runner its compressor and indexer to run
     between K1 and K2."""

vllm/models/deepseek_v41/amd/mono/README.md

@@ -20,7 +20,9 @@ new, reading vLLM's KV records and weights as loaded.
   their compressor; 24, 28, 32, 36) K1 also hands out its normed input and bf16
   qr, vLLM's compressor and indexer run between K1 and K2, and K2 builds the
   key table from that top-k. The FFN launch serves a layer outside these (an
-  MXFP8 KV cache, another indexer); the DSpark draft layers keep vLLM's path.
+  MXFP8 KV cache, another indexer). The DSpark draft's three layers run the same
+  K1 / K2 on a runner of their own (128 experts / top-3, window-only attention over
+  the block's non-causal SWA rows).
   The paths interleave freely (same tensors at the layer boundary).
 - Steps: decode only (FULL graph or eager), M <= 48 rows (8 requests x 6
   DSpark tokens), TP2 / TP4, `fp8_ds_mla` KV records (ROCm), text tokens (the

vllm/models/deepseek_v41/amd/mono/attention/front.py

@@ -106,6 +106,7 @@ class FrontBuild:
     tp: int
     ratio: int  # the layer's compress ratio (0: no top-k keys)
     timeline: bool = False
+    swa_width: int = WINDOW  # the SWA index rows' width (the DSpark draft's: wider)
 
 
 def plus(t0, v):
@@ -202,8 +203,8 @@ def stage_kt(c, t):
     kv_len = ntopk + nswa
     for i in range_constexpr((KEYS + THREADS - 1) // THREADS):
         k = fresh(tid) + THREADS * i
-        sidx = fx.min(fx.max(k - ntopk, 0), WINDOW - 1)
-        s_slot = ld_i32(a["swa_idx"], t * WINDOW + sidx)
+        sidx = fx.min(fx.max(k - ntopk, 0), c["swa_width"] - 1)
+        s_slot = ld_i32(a["swa_idx"], t * c["swa_width"] + sidx)
         val = (k < kv_len).select(s_slot, fx.Int32(-1))
         if const_expr(ratio > 0):
             local = ld_i32(a["topk"], t * TOPK + fx.min(k, TOPK - 1))
@@ -531,6 +532,7 @@ def front_context(key, lds, args: dict, scratch, tag, layout=None) -> dict:
         "mb": Mailbox(tag),
         "d": d,
         "ratio": key.ratio,
+        "swa_width": key.swa_width,
         "ldbf": lambda ptr, i: fx.Float32(
             fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16))
         ),

vllm/models/deepseek_v41/amd/mono/layer.py

@@ -25,15 +25,15 @@
 from .attention import back as mb_back
 from .attention import front as mb_front
 from .attention.device import BLOCKS, CM_DEV, THREADS, gstore, kernel_symbol, rsrc
-from .attention.plan import HIDDEN, Dims, back_scratch, front_scratch
+from .attention.plan import HIDDEN, WINDOW, Dims, back_scratch, front_scratch
 from .common.mx import clamp_fp8
 from .common.ops import butterfly, fp8_pack4, peer_bases, traced
 from .common.plan import WAVES, first_task, key_tuple
 from .common.sync import Mailbox, preg, publish, sreg
 from .sources import SOURCES
 from .stages import moe, seam
 from .stages.dims import Dims as StageDims
-from .stages.moe_shape import SORT_NETS, MoeBuild, route_shape
+from .stages.moe_shape import EXPERTS, SORT_NETS, TOPK, MoeBuild, route_shape
 
 _STREAM = fx.Stream(None)
 MAX_TOKENS = 48
@@ -64,21 +64,49 @@ class MonoBuild:
     # the layer selects its own top-k (an index source): K1 hands its normed
     # input and qr out for vLLM's indexer, K2 builds the key table
     index: bool = False
+    # a layer of the DSpark draft: its routed experts and top-k, the FP4 groups
+    # of an ug task, and the width of its (non-causal) SWA index rows
+    experts: int = EXPERTS
+    topk: int = TOPK
+    ug_groups: int = 2
+    swa_width: int = WINDOW
 
 
 def _moe_key(key: MonoBuild) -> MoeBuild:
-    return MoeBuild(tokens=key.tokens, tp=key.tp)
+    return MoeBuild(
+        tokens=key.tokens,
+        experts=key.experts,
+        topk=key.topk,
+        tp=key.tp,
+        ug_groups=key.ug_groups,
+    )
 
 
 def _front_key(key: MonoBuild):
-    return mb_front.FrontBuild(key.tokens, key.tp, key.ratio, key.timeline)
+    return mb_front.FrontBuild(
+        key.tokens, key.tp, key.ratio, key.timeline, key.swa_width
+    )
+
+
+def _symbol_params(key: MonoBuild) -> dict:
+    """The kernel symbol's draft-only parameters (none for a target layer)."""
+    out = {}
+    if key.experts != EXPERTS:
+        out["e"] = key.experts
+    if key.topk != TOPK:
+        out["k"] = key.topk
+    if key.ug_groups != 2:
+        out["g"] = key.ug_groups
+    if key.swa_width != WINDOW:
+        out["w"] = key.swa_width
+    return out
 
 
 def _back_key(key: MonoBuild):
     return mb_back.BackBuild(key.tokens, key.tp, key.ratio, key.timeline)
 
 
-def scratch_layout(s: int, tp: int) -> dict:
+def scratch_layout(s: int, tp: int, experts: int = EXPERTS, topk: int = TOPK) -> dict:
     """Every region of step width s's launches -> (byte offset, bytes), disjoint:
     K1's front and seam, K2's back, the MoE's regions, the normed rows and their
     flags. Each width has a scratch of its own (``DSV41MonoLayer.scratch``):
@@ -101,7 +129,9 @@ def place(regions):
     place(front_scratch(s, d))
     place(seam.scratch_layout(s))
     place(back_scratch(s, d, start=0))
-    regions = moe.scratch_layout(_moe_key(MonoBuild(s, tp, 1))).items()
+    regions = moe.scratch_layout(
+        _moe_key(MonoBuild(s, tp, 1, experts=experts, topk=topk))
+    ).items()
     mr = {n: v for n, v in regions if n not in COUNTERS}
     lo = min(o for o, _ in mr.values())
     place({n: (o - lo, n_) for n, (o, n_) in mr.items()})
@@ -110,8 +140,10 @@ def place(regions):
     return out
 
 
-def scratch_bytes(s: int = MAX_TOKENS, tp: int = 2) -> int:
-    return max(o + n for o, n in scratch_layout(s, tp).values())
+def scratch_bytes(
+    s: int = MAX_TOKENS, tp: int = 2, experts: int = EXPERTS, topk: int = TOPK
+) -> int:
+    return max(o + n for o, n in scratch_layout(s, tp, experts, topk).values())
 
 
 def peer_half_bytes(tp: int) -> int:
@@ -164,7 +196,7 @@ def publish_x8(c, t):
 def build_mono_k1(key: MonoBuild):
     s, tp = key.tokens, key.tp
     assert 1 <= s <= MAX_TOKENS
-    layout = scratch_layout(s, tp)
+    layout = scratch_layout(s, tp, key.experts, key.topk)
     fkey = _front_key(key)
     FrontLds = mb_front.front_lds(s, Dims(tp))
     NORM0 = SLICES  # the norms on the CTAs past the slices: the wqkv waits on them
@@ -190,6 +222,7 @@ class K1Lds:
         tl=key.timeline,
         f=key.fold,
         i=key.index,
+        **_symbol_params(key),
     )
     keyed = key_tuple(key, SOURCES)
 
@@ -458,7 +491,7 @@ def run_ffn(key: MonoBuild, lds, a: dict, amb, bases, own, rank, ep, part=None):
     the residual), the MoE and its all-reduce into ``a["out"]``, then the
     launch pair's epoch moved on."""
     s, tp = key.tokens, key.tp
-    layout = scratch_layout(s, tp)
+    layout = scratch_layout(s, tp, key.experts, key.topk)
     mkey = _moe_key(key)
     rs = route_shape(mkey)
     GATE0 = SLICES
@@ -554,7 +587,7 @@ def reset():
 def build_mono_k2(key: MonoBuild):
     s, tp = key.tokens, key.tp
     assert 1 <= s <= MAX_TOKENS
-    layout = scratch_layout(s, tp)
+    layout = scratch_layout(s, tp, key.experts, key.topk)
     bkey = _back_key(key)
     mkey = _moe_key(key)
     rs = route_shape(mkey)
@@ -574,7 +607,13 @@ class K2Lds:
         moe: MoeLds  # type: ignore[valid-type]
 
     name = kernel_symbol(
-        "dsv41_mono_k2", s=s, tp=tp, r=key.ratio, tl=key.timeline, i=key.index
+        "dsv41_mono_k2",
+        s=s,
+        tp=tp,
+        r=key.ratio,
+        tl=key.timeline,
+        i=key.index,
+        **_symbol_params(key),
     )
     keyed = key_tuple(key, SOURCES)
 
@@ -688,6 +727,7 @@ def push(t, col, v0, v1):
                 "tid": c["tid"],
                 "dev": True,
                 "ratio": key.ratio,
+                "swa_width": key.swa_width,
                 "args": {
                     "slot": slot,
                     "swa_idx": swa_idx,

vllm/models/deepseek_v41/amd/mono/runner.py

@@ -22,7 +22,7 @@
 
 import torch
 
-from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Q_RANK, Dims
+from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Q_RANK, WINDOW, Dims
 from .common.plan import BLOCKS
 from .layer import (
     EPOCH_WORDS,
@@ -35,7 +35,7 @@
     scratch_bytes,
 )
 from .stages.dims import Dims as MoeDims
-from .stages.moe_shape import EXPERTS
+from .stages.moe_shape import EXPERTS, TOPK
 
 __all__ = ["BLOCKS", "MAX_TOKENS", "AttnWeights", "DSV41MonoLayer", "MonoLayerWeights"]
 
@@ -125,23 +125,23 @@ class MonoLayerWeights:
     sw2: torch.Tensor  # shared down [5120, inter] e4m3
     sw2_s: torch.Tensor
 
-    def check(self, tp: int) -> None:
-        """The MoE's tensors: 384 experts, TP-sharded intermediates (the routed
-        one padded as the loader pads it), AITER's A8W4 layout, 32 x 32 E8M0
-        blocks for the shared expert."""
+    def check(self, tp: int, experts: int = EXPERTS) -> None:
+        """The MoE's tensors: ``experts`` routed experts (384; the DSpark draft's
+        128), TP-sharded intermediates (the routed one padded as the loader pads
+        it), AITER's A8W4 layout, 32 x 32 E8M0 blocks for the shared expert."""
         d, fp4, e4m3, u8 = (
             MoeDims(tp),
             torch.float4_e2m1fn_x2,
             torch.float8_e4m3fn,
             torch.uint8,
         )
         want = {
-            "gate_w": ((EXPERTS, HIDDEN), torch.bfloat16),
-            "bias": ((EXPERTS,), torch.float32),
-            "w13": ((EXPERTS, 2 * d.inter, HIDDEN // 2), fp4),
-            "w13_s": ((EXPERTS * 2 * d.inter, HIDDEN // 32), u8),
-            "w2": ((EXPERTS, HIDDEN, d.inter // 2), fp4),
-            "w2_s": ((EXPERTS * HIDDEN, d.down_scale_cols), u8),
+            "gate_w": ((experts, HIDDEN), torch.bfloat16),
+            "bias": ((experts,), torch.float32),
+            "w13": ((experts, 2 * d.inter, HIDDEN // 2), fp4),
+            "w13_s": ((experts * 2 * d.inter, HIDDEN // 32), u8),
+            "w2": ((experts, HIDDEN, d.inter // 2), fp4),
+            "w2_s": ((experts * HIDDEN, d.down_scale_cols), u8),
             "sgu": ((2 * d.sh_inter, HIDDEN), e4m3),
             "sgu_s": ((2 * d.sh_inter // 32, HIDDEN // 32), u8),
             "sw2": ((HIDDEN, d.sh_inter), e4m3),
@@ -154,11 +154,25 @@ class DSV41MonoLayer:
     """The mono decode runner of one TP rank (``group``: its TP group, for the
     peer-memory handle exchange; a gloo / CPU group)."""
 
-    def __init__(self, tp: int, rank: int, group, device: torch.device | str = "cuda"):
+    def __init__(
+        self,
+        tp: int,
+        rank: int,
+        group,
+        device: torch.device | str = "cuda",
+        experts: int = EXPERTS,
+        topk: int = TOPK,
+        ug_groups: int = 2,
+        swa_width: int = WINDOW,
+    ):
         _ensure_writable_flydsl_cache()
         from .common.peer_memory import PeerBuffer
 
         self.tp, self.rank = tp, rank
+        # the layers it runs: the target's (384 experts, top-6, 128-wide causal SWA
+        # rows) or the DSpark draft's, which has a runner of its own
+        self.experts, self.topk, self.ug_groups = experts, topk, ug_groups
+        self.swa_width = swa_width
         self.d = Dims(tp)
         dev = self.device = torch.device(device)
         self._scratch: dict[int, torch.Tensor] = {}
@@ -205,7 +219,9 @@ def scratch(self, tokens: int) -> torch.Tensor:
                     "a CUDA graph capture"
                 )
             buf = torch.zeros(
-                scratch_bytes(tokens, self.tp), dtype=torch.uint8, device=self.device
+                scratch_bytes(tokens, self.tp, self.experts, self.topk),
+                dtype=torch.uint8,
+                device=self.device,
             )
             self._scratch[tokens] = buf
         return buf
@@ -220,7 +236,17 @@ def supports(tokens: int) -> bool:
     def kernels(self, tokens: int, ratio: int, fold: bool = True, index: bool = False):
         key = (tokens, ratio, fold, index)
         if key not in self._kernels:
-            b = MonoBuild(tokens, self.tp, ratio, fold=fold, index=index)
+            b = MonoBuild(
+                tokens,
+                self.tp,
+                ratio,
+                fold=fold,
+                index=index,
+                experts=self.experts,
+                topk=self.topk,
+                ug_groups=self.ug_groups,
+                swa_width=self.swa_width,
+            )
             self._kernels[key] = (build_mono_k1(b), build_mono_k2(b))
         return self._kernels[key]
 
@@ -282,7 +308,9 @@ def forward(
             assert t_.is_contiguous()
         assert positions.dtype == slot_mapping.dtype == torch.int64
         assert swa_indices.dtype == swa_lens.dtype == token_to_req.dtype == torch.int32
-        assert swa_indices.stride(0) == 128 and swa_cache.dtype == torch.uint8
+        assert (
+            swa_indices.stride(0) == self.swa_width and swa_cache.dtype == torch.uint8
+        )
         assert swa_cache.shape[-1] == 584 and swa_cache.stride(1) == 584
         swa_block = swa_cache.shape[1]
         if ratio: