feat(dsv41 mono): serve the widths that spill by dropping K2's index stage, not the width

atom/models/deepseek_v41/mono/builds.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+"""What a V4.1 mono width was built as, fused or unfused.
+
+One record a width (``WidthBuilds``, keyed by the step's row count), holding
+the launchers and -- the reason it is a record and not a row count -- the K2
+build key they came from. A step's shape is not a function of its rows alone,
+so ``runner`` reads it off the build in hand rather than deriving it again.
+
+A width is normally built *fused*: K2 carries the indexer, so a selecting
+layer's scoring, its top-k and the pool rows the attention reads all happen
+inside the one launch (``kernels.index_score``). At some widths that build
+needs a byte of scratch a lane -- at TP 4 a 32- or 48-row K2 on the MFMA FP4
+index layout comes out at 256 VGPRs instead of 255 -- and a persistent kernel
+whose waves need scratch cannot be dispatched whole, so ``WidthBuilds`` refuses
+the width and its steps take the original model's path.
+
+The unfused variant is what ``WidthBuilds``' ``lean`` seam rebuilds such a
+width as: K2 without its index stage (``index_bound_max`` 0, the shape the
+DSpark draft already builds), with ATOM's own indexer run between K1 and K2
+and its selection published into the buffers K2 would otherwise have filled
+itself (``runner.V41MonoDecodeRunner._publish_selection``). The same
+configuration, on kernels byte-identical to these, is what the out-of-tree
+AITER provider drives, so this is a build choice, not a new kernel.
+
+``strict`` is the other half of the same story: a width refused even after the
+lean rebuild is still a silent fallback to the original model's path, and
+``ATOM_DSV41_MONO_STRICT`` makes the runner say so rather than let a
+measurement report it as mono.
+"""
+
+from dataclasses import dataclass, replace
+
+from atom.utils import envs
+
+
+@dataclass(frozen=True)
+class RowsBuild:
+    """The kernels of an S-row step: K1's build a layer and their launchers,
+    and K2 (or, with ``check``, K2a and K2b apart).
+
+    ``post_key`` is the K2 build key the launchers came from: the one place a
+    step's ``tok_par`` and its K2 scratch layout are recorded, so the forward
+    and the debug paths ask the build what it is instead of recomputing it
+    from the row count and hoping the two answers agree. ``unfused``: K2 was
+    built without its indexer and the runner scores between K1 and K2."""
+
+    keys: list
+    launchers: dict
+    layer_post: object = None
+    attn_post: object = None
+    moe: object = None
+    post_key: object = None
+    unfused: bool = False
+
+    @property
+    def tok_par(self) -> bool:
+        """Whether this build's kernels run token parallel. Read it off the
+        build, never off the row count: the unfused variant of a width that
+        would run token parallel does not."""
+        return bool(self.post_key is not None and self.post_key.tok_par)
+
+
+def unfused_post_key(key):
+    """``key`` with K2's index stage out.
+
+    ``index_bound_max`` 0 builds the K2 that reads a selection instead of
+    making one; ``tok_par`` off because the indexer is what a token-parallel
+    step splits (each rank selects its own rows and writes every rank's copy,
+    ``index_score._st_out``) -- without it in the kernel every rank scores
+    every row on the host side and keeps its own copy.
+    """
+    return replace(key, index_bound_max=0, tok_par=False)
+
+
+def unfused_keys(keys: list, post_key):
+    """A width's (K1 build keys, K2 build key), unfused.
+
+    K1 loses token parallel with K2: a token-parallel K1 writes only its own
+    rank's rows (``token_par.own_rows``), which is whose selection that rank's
+    K2 indexer would have made. Nothing else about K1 changes -- it still
+    projects the indexer query, which the host-side scorer's FP8 form reads.
+    """
+    key = unfused_post_key(post_key)
+    return [replace(k, tok_par=key.tok_par) for k in keys], key
+
+
+def lean_width(built: RowsBuild, spilled: list, rebuild):
+    """``built`` rebuilt unfused when one of its K2 launchers needs scratch,
+    else None (``WidthBuilds``: nothing leaner, so the width is refused).
+
+    ``rebuild()``: the same width's unfused build. Only a fused build has a
+    leaner form, and only a K2 spill is what the unfused form fixes -- a K1
+    that spills on its own is rebuilt by nothing here, nor is a K2b, which the
+    unfused key does not change (``LayerPostBuild.halves``), nor a build whose
+    K2 has no index stage to drop (one already unfused, or the DSpark draft's).
+    """
+    key = built.post_key
+    if built.unfused or key is None or key.index_bound_max == 0:
+        return None
+    k2 = [f for f in (built.layer_post, built.attn_post) if f is not None]
+    if not any(f in spilled for f in k2):
+        return None
+    return rebuild()
+
+
+def strict() -> bool:
+    """``ATOM_DSV41_MONO_STRICT``: refuse to fall back silently."""
+    return bool(envs.ATOM_DSV41_MONO_STRICT)
+
+
+def refuse(rows: int) -> None:
+    """Raise for a step width mono will not serve, under ``strict``.
+
+    Off (the default) this does nothing and the step takes the original
+    model's path, as every unserved width does. On, it is an error: the only
+    reason to ask for a width is a measurement, and a measurement that quietly
+    reports the original path is worse than no measurement.
+    """
+    if not strict():
+        return
+    raise RuntimeError(
+        f"V4.1 mono: width {rows} refused (the warning above names why); its "
+        "steps would take the original path (ATOM_DSV41_MONO_STRICT=1)"
+    )

atom/models/deepseek_v41/mono/check.py

@@ -330,7 +330,7 @@ def check_moe(runner, block):
     # sums, added in part order and rounded to bf16 as ``route`` does
     x = runner.ffn_normed.reshape(-1, runner.ffn_normed.shape[-1])
     # K2b's regions where either launch shape puts them (``layer_post.moe_base``)
-    layout = k2.scratch_layout(runner._post_key(x.shape[0]))
+    layout = k2.scratch_layout(runner.build.post_key)
     off, nbytes = layout["logit"]
     pairs = runner.scratch[off : off + nbytes].view(torch.float32).view(-1, 2)
     parts = pairs[:, 0].view(x.shape[0], -1, k2b.ROUTER_PARTS)

atom/models/deepseek_v41/mono/runner.py

@@ -16,7 +16,6 @@
 
 import logging
 import struct
-from dataclasses import dataclass
 
 import torch
 from aiter.dist.parallel_state import (
@@ -34,6 +33,12 @@
 from atom.model_ops.v4_kernels.paged_decode import LOG2E
 from atom.models.deepseek_v41.config import AttentionMode
 from atom.models.deepseek_v41.mono import index_plan as ip
+from atom.models.deepseek_v41.mono.builds import (
+    RowsBuild,
+    lean_width,
+    refuse,
+    unfused_keys,
+)
 from atom.models.deepseek_v41.mono.check import (
     check_front,
     check_moe,
@@ -55,7 +60,11 @@
 from atom.models.deepseek_v41.mono.kernels.moe_shape import EXPERTS, TOPK
 from atom.models.deepseek_v41.mono.kernels.peer_layout import peer_bytes, plain_layout
 from atom.models.deepseek_v41.mono.kernels.token_par import own_rows
-from atom.models.deepseek_v41.mono.step_meta import KMETA, write_step_meta
+from atom.models.deepseek_v41.mono.step_meta import (
+    KMETA,
+    write_selection_rows,
+    write_step_meta,
+)
 from atom.models.deepseek_v41.mono.weights import bind_model
 from atom.mono.device import queue
 from atom.mono.plan.execution import BLOCKS
@@ -105,18 +114,6 @@ def _rows(x, dim, s):
     return x.narrow(dim, 0, s)
 
 
-@dataclass(frozen=True)
-class RowsBuild:
-    """The kernels of an S-row step: K1's build a layer and their launchers,
-    and K2 (or, with ``check``, K2a and K2b apart)."""
-
-    keys: list
-    launchers: dict
-    layer_post: object = None
-    attn_post: object = None
-    moe: object = None
-
-
 # the layers whose attention norm has other readers (the compressor, the indexer)
 _FEEDS_INDEX = (AttentionMode.FULL, AttentionMode.REINDEX)
 
@@ -179,7 +176,11 @@ def _bind_agreed(self, bind) -> None:
         self._full = {name: getattr(self, name) for name, _ in _ROW_BUFFERS}
         self.mailboxes = StepMailboxes(self.peers, self.scratch, debug=self.debug)
         self.widths = WidthBuilds(
-            self._rows_build, self._row_kernels, group, "V4.1 mono"
+            self._rows_build,
+            self._row_kernels,
+            group,
+            "V4.1 mono",
+            lean=self._lean_width,
         )
 
     def _place_index_outputs(self, plain: dict) -> None:
@@ -200,8 +201,17 @@ def region(name, *shape):
         self.icand = region("icand", self.max_rows, ip.TOPK_BLOCKS)
 
     def prepare(self, rows: int) -> bool:
-        """Whether every rank holds the ``rows``-row kernels (``WidthBuilds``)."""
-        return 0 < rows <= self.max_rows and self.widths.prepare(rows)
+        """Whether every rank holds the ``rows``-row kernels (``WidthBuilds``).
+
+        ``ATOM_DSV41_MONO_STRICT`` turns a refused width into an error instead:
+        its steps would otherwise run the original model's path without saying
+        so, which an A/B measures as if it were mono."""
+        if not 0 < rows <= self.max_rows:
+            return False
+        if self.widths.prepare(rows):
+            return True
+        refuse(rows)
+        return False
 
     def _bind(self, model, drafter) -> None:
         self.model = model
@@ -346,9 +356,10 @@ def _tok_par(self, s: int) -> bool:
         )
 
     def _own_rows(self, s: int) -> tuple[int, int]:
-        """(first row, rows) of an s-row step this rank's indexer selects
-        (``index_score._own_rows``)."""
-        return own_rows(s, self.tp, self.rank) if self._tok_par(s) else (0, s)
+        """(first row, rows) of this step's rows this rank's indexer selects
+        (``index_score._own_rows``), off the build in hand rather than off
+        ``s``: a width's shape is the build's answer, not the row count's."""
+        return own_rows(s, self.tp, self.rank) if self.build.tok_par else (0, s)
 
     def _layer_keys(self, s: int) -> list:
         """K1's build of every layer for an s-row step."""
@@ -389,20 +400,38 @@ def _post_key(self, s: int) -> k2.LayerPostBuild:
             kv_fp8=self.kv_fp8,
         )
 
-    def _rows_build(self, s: int) -> RowsBuild:
-        """The s-row step's kernels."""
+    def _rows_build(self, s: int, *, unfused: bool = False) -> RowsBuild:
+        """The s-row step's kernels, fused or (``unfused``) with K2's indexer
+        out and token parallel off (``builds.unfused_post_key``): the runner
+        then scores between K1 and K2 (``_score``, ``_publish_selection``)."""
+        post_key = self._post_key(s)
         keys = self._layer_keys(s)
+        if unfused:
+            keys, post_key = unfused_keys(keys, post_key)
         launchers = {key: k1.build_attn_pre(key) for key in set(keys)}
-        post_key = self._post_key(s)
+        rest = {"post_key": post_key, "unfused": unfused}
         if self.split_k2:
             key_a, key_b = post_key.halves()
             return RowsBuild(
                 keys,
                 launchers,
                 attn_post=k2a.build_attn_post(key_a),
                 moe=k2b.build_moe(key_b),
+                **rest,
             )
-        return RowsBuild(keys, launchers, layer_post=k2.build_layer_post(post_key))
+        return RowsBuild(
+            keys, launchers, layer_post=k2.build_layer_post(post_key), **rest
+        )
+
+    def _lean_width(self, built: RowsBuild, spilled: list) -> RowsBuild | None:
+        """``WidthBuilds``' lean seam: a width whose K2 needs scratch rebuilt
+        unfused, which has no index stage to spill on. Anything else has
+        nothing leaner and the width is refused as before."""
+        return lean_width(
+            built,
+            spilled,
+            lambda: self._rows_build(built.post_key.tokens, unfused=True),
+        )
 
     @staticmethod
     def _row_kernels(build: RowsBuild) -> list:
@@ -534,11 +563,14 @@ def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Ten
             cache.geometry.window(0, cache.num_pages),
             self.ring_rel,
             self.key_meta,
-            self.walk_plans,
+            # an unfused width's scorer is ATOM's own, which plans its own work
+            {} if self.build.unfused else self.walk_plans,
             # a token-parallel step's indexer scores this rank's rows alone
             self._own_rows(rows),
         )
         state = SinglePassHCState.from_embeddings(hidden, model.config.hc_mult)
+        # an unfused width: the selection owners this step has published
+        published = set()
         for spec, block, weights, key in zip(
             model.topology, model.layers, self.layers, self.build.keys
         ):
@@ -565,7 +597,12 @@ def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Ten
                 )
             if block.attn.indexer is not None:
                 block.attn._compress_batch(self.normed, cache, step, rope)
-                if self.check:
+                if self.build.unfused:
+                    # this width's K2 has no index stage: ATOM's own indexer
+                    # makes the selection and K2 reads it
+                    self._score(block, cache, step, rope)
+                    self._publish_selection(spec, cache, step, published)
+                elif self.check:
                     # the original scorer, the reference the selection is held to
                     block.attn.indexer.score_quantized(
                         self.index_query(), self.iw, cache, step
@@ -588,7 +625,9 @@ def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Ten
                     rows_per_block=ip.BLOCK_ROWS,
                 )
             if self.check:
-                if block.attn.indexer is not None:
+                # an unfused width makes no selection of its own to hold to the
+                # original scorer: the one K2 read is already that scorer's
+                if block.attn.indexer is not None and not self.build.unfused:
                     check_selection(self, spec, step)
                 # the original index build, for the reference attention alone
                 keys = decode_keys(cache, spec, step)
@@ -598,7 +637,7 @@ def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Ten
             pre, post, comb = self.gates_ffn
             state = SinglePassHCState(self.res_ffn, pre, self.ffn_out, post, comb)
         hidden = state.collapse()
-        if self._tok_par(rows):
+        if self.build.tok_par:
             hidden = self._gather_rows(hidden, rows)
         model.end_forward(hidden, None)
         self._raise_on_given_up_waits(rows)
@@ -619,6 +658,47 @@ def _gather_rows(self, hidden, rows):
             buf[:rows].copy_(every[:, k + 1])
         return every[:, 0].contiguous().unsqueeze(0)
 
+    def _score(self, block, cache, step, rope) -> None:
+        """An unfused width's indexer, run between K1 and K2: this layer's
+        top-k (and a candidate producer's blocks) into ``step``, on K1's
+        outputs, as the AITER provider drives these same kernels."""
+        indexer = block.attn.indexer
+        if self.index_fp4:
+            # K1's FP4 query is laid out for K2's scorer, not this one: the
+            # original projection off K1's normed row and q latent instead
+            query, weights = indexer.project(
+                self.normed, self.qr, self.qr_scale, cache, step, rope
+            )
+            indexer.score(query, weights, cache, step)
+        else:
+            # K1's per-(row, head) FP8 query and its scale-folded head weights,
+            # as check mode scores them
+            indexer.score_quantized(self.iq, self.iw, cache, step)
+
+    def _publish_selection(self, spec, cache, step, published: set) -> None:
+        """An unfused width's selection where K2 reads it: owner
+        ``spec.layer_id``'s pool rows in ``irow`` and, on a candidate producer,
+        its blocks in ``icand`` -- what a fused build's K2 writes itself
+        (``index_score._final``). Once an owner a step: every layer that shares
+        a selection shares its KV owner, so its rows."""
+        owner = spec.layer_id
+        if owner in published:
+            return
+        published.add(owner)
+        geometry = cache.geometry
+        write_selection_rows(
+            self.irow[owner],
+            step.selected[owner][0],
+            step.batch_ids,
+            step.block_tables,
+            geometry.rows_per_page(spec.ratio),
+            geometry.page_bytes // geometry.row_bytes,
+            geometry.main_offset(spec.kv_owner),
+        )
+        if spec.produces_candidates:
+            # the REINDEX layers read them through ``candidate_block_table``
+            self.icand.copy_(step.candidates[owner])
+
     def _front(self, layer_id, key, weights, state, cache, step, rope):
         """K1 -> (the residual after the owed post, the seam's (pre, post, comb));
         the query, the window row and the indexer's inputs land in this runner's
@@ -750,7 +830,7 @@ def _attn_post_args(self, spec, block, weights, residual, gates, step, cache, ro
     def _moe(self, layer_id, weights) -> None:
         """K2b: router, routed and shared experts and their all-reduce ->
         ``ffn_out``."""
-        key_a, _ = self._post_key(self.ffn_normed.shape[1]).halves()
+        key_a, _ = self.build.post_key.halves()
         args = k2b.ABI.pack({
             "x": self.ffn_normed.data_ptr(), **self._moe_args(weights),
             **self._peer_args(), "layer": MOE_TAG0 + layer_id,
@@ -785,7 +865,7 @@ def _raise_on_given_up_waits(self, rows: int) -> None:
             return
         regions = {
             *k1.scratch_layout(rows),
-            *k2.scratch_layout(self._post_key(rows)),
+            *k2.scratch_layout(self.build.post_key),
         }
         diag = self.scratch[self.diag_off : self.diag_off + DIAG_BYTES]
         waits = given_up_waits(diag, kernel_debug.region_names(regions))