[Feat] deepseek_v41: Score SGLang's page8 index-K pool in K2i

README.md

@@ -59,9 +59,7 @@ SGLANG_USE_AITER=1 AITER_BF16_FP8_MOE_BOUND=0 python -m sglang.launch_server \
   --monokernel-backend amd
 ```
 
-On gfx950 the decode graph capture fails in AITER's MoE without
-`AITER_BF16_FP8_MOE_BOUND=0`. The flags are those of the serving tests
-(`test/e2e/deepseek_v41_server.py`).
+The flags are those of the serving tests (`test/e2e/deepseek_v41_server.py`).
 
 GLM-5.3, with MTP (its NextN layer drafts one chain; a verify step of 4 rows a
 request runs the kernels at batch sizes 1 and 2 at TP8):
@@ -130,7 +128,8 @@ MXFP8 activations, as the kernels do.
 For MoRI PD, enable monokernels only on the decode server
 (`--disaggregation-mode decode --disaggregation-transfer-backend mori`): TP8 for
 DeepSeek-V4-Pro, TP2 or TP4 for DeepSeek-V4.1-Flash. Leave the flag off on prefill;
-both peers must use the same KV layout.
+both peers must use the same KV layout, `SGLANG_HIP_DSV41_INDEX_K_PAGE8` included
+(SGLang's PD handshake does not compare it).
 
 If the model, GPU or configuration is not supported, the server refuses to start and lists every reason.
 

sgl_amd_mono/deepseek_v41/gfx950/indexer/k2i.py

@@ -25,11 +25,11 @@
 MAX_ROWS = 48  # a step's rows at most (8 DSpark requests of 6)
 
 
-def build_k2i(tokens: int):
-    """K2i for a step of ``tokens`` rows."""
+def build_k2i(tokens: int, page8: bool = False):
+    """K2i for a step of ``tokens`` rows; ``page8``: on a page8 pool's plane (``score_fp4``)."""
     assert 0 < tokens <= MAX_ROWS, tokens
-    name = kernel_symbol("dsv41_mono_k2i", s=tokens)
-    keyed = (("tokens", tokens), ("sources", SOURCES))
+    name = kernel_symbol("dsv41_mono_k2i", s=tokens, **({"p8": 1} if page8 else {}))
+    keyed = (("tokens", tokens), ("page8", page8), ("sources", SOURCES))
 
     @fx.struct
     class RowsLds:
@@ -55,7 +55,7 @@ def k2i(
         }  # fmt: skip
         c = {
             "tid": tid, "lane": tid % 64, "wave": tid // 64, "args": args,
-            "rows_lds": lds.rows.ptr,
+            "rows_lds": lds.rows.ptr, "page8": page8,
         }  # fmt: skip
         run_scores4(c, tokens, fx.block_idx.x)
 

sgl_amd_mono/deepseek_v41/gfx950/indexer/runner.py

@@ -92,13 +92,14 @@ def __init__(self, device: torch.device, bound_max: int):
         )
         self._kernels: dict = {}
 
-    def kernel(self, rows: int):
+    def kernel(self, rows: int, page8: bool = False):
         from ..runner import no_capture
 
-        if rows not in self._kernels:
-            no_capture(f"K2i width {rows}")
-            self._kernels[rows] = build_k2i(rows)
-        return self._kernels[rows]
+        key = (rows, page8)
+        if key not in self._kernels:
+            no_capture(f"K2i width {rows} (page8 {page8})")
+            self._kernels[key] = build_k2i(rows, page8)
+        return self._kernels[key]
 
     def walk_plan(
         self, *, ratio: int, bounds: torch.Tensor, rows: int, per_req: int
@@ -121,7 +122,7 @@ def select(
         per_req: int,
         ratio: int,
         query: tuple,  # [rows, 32, 64] E2M1 pairs, [rows, 1, 4, 16, 4] E8M0, [rows, 32] bf16
-        plane: tuple,  # the layer's index plane (``check_plane``)
+        plane: tuple,  # the layer's index plane (``check_plane``), whether a page8 pool's
         page_table: torch.Tensor,  # int32 [rows, pages] of 64-row pages
         bounds: torch.Tensor,  # int32 [rows]: the columns each row sees
         raw_out: torch.Tensor,  # int32 [rows, 512]: sparse_raw_indices
@@ -235,9 +236,9 @@ def _score(
         ibound, lstride, reindex,
     ) -> None:  # fmt: skip
         q_fp4, q_scale, weights = query
-        payload, scale = plane
+        payload, scale, page8 = plane
         _run_compiled(
-            self.kernel(rows),
+            self.kernel(rows, page8),
             q_fp4.data_ptr(), q_scale.data_ptr(), weights.data_ptr(), payload.data_ptr(), scale.data_ptr(),
             plan.data_ptr(), itab.data_ptr(), itab_stride, itab_len, ibatch.data_ptr(), ishift,
             ibound.data_ptr(), self.ilog.data_ptr(), lstride, reindex,

sgl_amd_mono/deepseek_v41/gfx950/indexer/score_fp4.py

@@ -11,6 +11,9 @@
     iq  [S, HEADS, DIM / 2] E2M1, iqs [S, 4 (K chunk), 16 (head % 16), 4 (head // 16, padded)]
         E8M0, iw [S, HEADS] bf16 (widening to f32 is exact)
 
+A page8 pool (``SGLANG_HIP_DSV41_INDEX_K_PAGE8``) holds kv and ks in aiter's 8-row pages instead,
+which are the virtual pages: only ``_keys`` reads the two layouts differently.
+
 Lane (g, l) of tile nt holds key 4 l + nt, so a step's page is one lookup and its four tiles'
 scales one dword. The logits match aiter ``pa_mqa_logits_fp4_rowgroup`` only in its arithmetic
 order: relu as x + |x| against half the weight, the lane's heads in one fma chain, then the
@@ -40,6 +43,7 @@
 from . import index_plan as ip
 
 PAGE = ip.BLOCK_ROWS
+# a page8 pool's page, in aiter's ``pack_kv_cache`` layout
 PAGE_BYTES = PAGE * ip.DIM // 2
 PAGE_SCALES = PAGE * ip.DIM // 32
 # SGLang's FP4 index plane (``deepseek_v4_memory_pool``): 64-row pages, payload and scales apart
@@ -216,6 +220,13 @@ def _keys(c, page, nt):
     REINDEX row, whose request's other rows read the same candidate pages."""
     a, lane = c["args"], c["lane"]
     l, g = lane % 16, lane // 16
+    if const_expr(c["page8"]):
+        # payload [key % 4][chunk][key // 4][16 B], scales [chunk][key]: this lane's 4 (l % 2) + nt
+        page = fx.Int64(page)
+        base = a["plane"] + page * PAGE_BYTES + fx.Int64(g * 32 + (l % 2) * 16)
+        kv = [global_load(base + t * 128, 4, nt) for t in range(4)]
+        at = a["pscale"] + page * PAGE_SCALES + fx.Int64(g * 8 + (l % 2) * 4)
+        return kv, fx.Int32(global_load(at))
     # ``page`` is virtual: 8 x SGLang's 64-row page + the 8-row sub page
     sub = page & 7
     phys = fx.Int64(page >> 3)

sgl_amd_mono/deepseek_v41/gfx950/runner.py

@@ -456,20 +456,19 @@ def precompile(
         k1=(),
         k2=(),
         ffn: bool = False,
-        k2i: bool = False,
+        k2i=(),
         tok_par: bool = False,
     ) -> None:
         """Allocate this width's scratch and compile its launches without running them.
 
         K1 for each (fold, index, iq) in ``k1``, K2 for each (compressed, index) in
-        ``k2``, the FFN launch, and K2i (``enable_k2i`` first)."""
+        ``k2``, the FFN launch, and K2i for each page8 in ``k2i`` (``enable_k2i`` first)."""
         self.scratch(tokens)
         kernels = [self.k1_kernel(tokens, *v, tok_par=tok_par) for v in sorted(k1)]
         kernels += [self.k2_kernel(tokens, *v, tok_par=tok_par) for v in sorted(k2)]
         if ffn:
             kernels.append(self.ffn_kernel(tokens))
-        if k2i:
-            kernels.append(self.k2i.kernel(tokens))
+        kernels += [self.k2i.kernel(tokens, page8) for page8 in sorted(k2i)]
         for launch in kernels:
             compile_only(launch)
 

sgl_amd_mono/deepseek_v41/spec.py

@@ -140,6 +140,7 @@
         "disaggregation_transfer_backend",
     ),
     "srt/environ.py": (
+        "envs",
         "SGLANG_DSV4_FP4_EXPERTS",
         "SGLANG_DSV4_FP4_DEQUANT",
         "SGLANG_USE_AITER_MOE_GU_ITLV",
@@ -224,6 +225,7 @@
         "get_index_k_fp4_payload_buffer",
         "get_index_k_fp4_scale_buffer",
         "low_ratio_index_k_is_split",
+        "low_ratio_index_k_page8",
         "get_index_k_page_size",
     ),
     "srt/model_executor/forward_context.py": (
@@ -564,7 +566,7 @@ class _Layer(msgspec.Struct):
     swa_block: int = 0
     comp: Any = None  # the compressed cache, the same layout
     comp_block: int = 0
-    plane: Any = None  # an indexer's FP4 index plane: (payload, scales)
+    plane: Any = None  # an indexer's FP4 index plane: (payload, scales, page8)
 
 
 def _keep(copies: dict, t: torch.Tensor, *key) -> torch.Tensor:
@@ -923,7 +925,7 @@ def _index_plane(*, L, M: int, meta, raw_idx) -> tuple:
     payload = pool.get_index_k_fp4_payload_buffer(L.i)
     scale = pool.get_index_k_fp4_scale_buffer(L.i)
     check_plane(payload=payload, scale=scale)
-    return payload.view(torch.uint8), scale
+    return payload.view(torch.uint8), scale, pool.low_ratio_index_k_page8(L.i)
 
 
 class _Built:
@@ -1161,6 +1163,7 @@ def _longest_context(hf_config) -> int:
 
 
 def _build(model, hf_config, ctx) -> _Built:
+    from sglang.srt.environ import envs
     from sglang.srt.runtime_context import get_parallel
 
     from sgl_amd_mono.deepseek_v41.gfx950.runner import (
@@ -1185,7 +1188,14 @@ def _build(model, hf_config, ctx) -> _Built:
         shuffled=bool(shuffled_layout(L.module.self_attn for L in layers.values())),
     )
     runner.set_deadline(ctx.deadline_s)
-    k2i = any(L.module.self_attn.indexer is not None for L in layers.values())
+    # the plane layouts K2i reads: SGLang's page8 pool is ratio 1's unless aiter lacks the
+    # row-group kernel (each layer's plane picks its build, ``_index_plane``)
+    page8 = envs.SGLANG_HIP_DSV41_INDEX_K_PAGE8.get()
+    k2i = {
+        page8 and L.ratio == 1
+        for L in layers.values()
+        if L.module.self_attn.indexer is not None
+    }
     if k2i:
         # the indexer's buffers, sized here: before SGLang sizes the KV cache
         runner.enable_k2i(_longest_context(hf_config))
@@ -1240,8 +1250,6 @@ def _build(model, hf_config, ctx) -> _Built:
             "SGLANG_DSV4_KV_LAYOUT": "v4",
             # K1 quantizes the SWA nope payload from f32, as SGLang's fused writer does
             "SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE": False,
-            # the scorer reads the ratio-1 index-K pages in the page-64 byte layout
-            "SGLANG_HIP_DSV41_INDEX_K_PAGE8": False,
             # benchmark routing overrides: K2 routes with the model's own router
             "SGLANG_SIMULATE_UNIFORM_EXPERTS": False,
             "SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS": False,