From b9ba7bbaee0d15bd77ef8d193a9829067d77f099 Mon Sep 17 00:00:00 2001
From: yihonglie <hyi@amd.com>
Date: Fri, 25 Sep 2026 22:48:35 -0500
Subject: [PATCH 1/5] perf(dsv4): adapt tuned AITER ASM FP8 decode

---
 atom/model_ops/v4_kernels/paged_decode.py    | 66 +++++++++++++++
 atom/models/deepseek_v4.py                   |  4 +
 tests/test_v4_prefill_asm_decode_dispatch.py | 86 +++++++++++++++++++-
 3 files changed, 155 insertions(+), 1 deletion(-)

diff --git a/atom/model_ops/v4_kernels/paged_decode.py b/atom/model_ops/v4_kernels/paged_decode.py
index 6cb5a10a4b..dd2d9d440d 100644
--- a/atom/model_ops/v4_kernels/paged_decode.py
+++ b/atom/model_ops/v4_kernels/paged_decode.py
@@ -84,6 +84,65 @@
 _FP8_DTYPE = torch.float8_e4m3fnuz
 
 
+@functools.lru_cache(maxsize=None)
+def _device_arch(device_index: int) -> str:
+    """Return the base GCN architecture name without feature suffixes."""
+    return getattr(
+        torch.cuda.get_device_properties(device_index), "gcnArchName", ""
+    ).split(":", 1)[0]
+
+
+def _v4_aiter_fp8_decode_splits(
+    q_packed: torch.Tensor,
+    *,
+    query_group: int,
+    kv_kind: str,
+) -> int | None:
+    """Return the qualified gfx950 H128/q7 AITER CSA split count.
+
+    AITER #5195's V4 NM wrapper maximizes a pure CU-occupancy score when the
+    caller leaves ``num_kv_splits`` unset. For q7 this creates a sawtooth
+    schedule: nearby batches can select 1 versus 15 splits, and every
+    multi-split launch writes FP32 partial outputs before a second reduction
+    kernel. The measured batch-only table below stays CUDA Graph safe because
+    it depends only on the captured Q shape. It is robust across K384/K640/
+    K1152 and avoids reading GPU-resident ``kv_indptr`` on the host.
+
+    HCA remains on AITER's ragged split planner: its heterogeneous long-KV
+    vectors have a different optimum from the CSA table.
+    """
+    shape = getattr(q_packed, "shape", None)
+    device = getattr(q_packed, "device", None)
+    if shape is None or len(shape) != 3 or device is None:
+        return None
+
+    device_index = device.index
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    tokens, heads, _ = shape
+    if (
+        _device_arch(device_index) != "gfx950"
+        or heads != 128
+        or query_group != 7
+        or kv_kind != "csa"
+        or tokens % query_group
+    ):
+        return None
+
+    requests = tokens // query_group
+    if requests <= 0:
+        return None
+    if requests == 1:
+        return 5
+    if requests == 2:
+        return 4
+    if requests == 3:
+        return 3
+    if requests == 4:
+        return 2
+    return 1
+
+
 @functools.lru_cache(maxsize=1)
 def _cu_count() -> int:
     """Compute-unit count of the active GPU, queried once via aiter.
@@ -1141,6 +1200,8 @@ def sparse_attn_v4_paged_decode(
     kv_last_page_lens: torch.Tensor | None = None,
     empty_kv_indptr: torch.Tensor | None = None,
     prefix: str = "",
+    query_group: int = 1,
+    kv_kind: str = "",
 ) -> torch.Tensor:
     """V4 decode sparse attention over a unified KV pool with paged indices.
 
@@ -1185,6 +1246,11 @@ def sparse_attn_v4_paged_decode(
             q_rope_in,
             qo_indptr=qo_indptr,
             kv_last_page_lens=kv_last_page_lens,
+            num_kv_splits=_v4_aiter_fp8_decode_splits(
+                q_packed_in,
+                query_group=query_group,
+                kv_kind=kv_kind,
+            ),
         )
     gfx = get_gfx()
     if gfx == "gfx1250" or gfx.startswith("gfx94"):
diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py
index 320198dde0..4691bed778 100644
--- a/atom/models/deepseek_v4.py
+++ b/atom/models/deepseek_v4.py
@@ -3336,6 +3336,10 @@ def _sparse_attention(
                 qo_indptr=attn_md.qo_indptr,
                 empty_kv_indptr=attn_md.empty_kv_indptr,
                 prefix=f"{self.layer_name}.sparse_attn_decode",
+                query_group=(
+                    7 if attn_md.min_seqlen_q == attn_md.max_seqlen_q == 7 else 1
+                ),
+                kv_kind="csa" if ratio == 4 else "hca" if ratio == 128 else "swa",
             )  # [S, H, head_dim]
         else:
             # Two-source paged prefill: prefix from `unified_kv` (per-ratio
diff --git a/tests/test_v4_prefill_asm_decode_dispatch.py b/tests/test_v4_prefill_asm_decode_dispatch.py
index 85ad316742..8ae5d33c6d 100644
--- a/tests/test_v4_prefill_asm_decode_dispatch.py
+++ b/tests/test_v4_prefill_asm_decode_dispatch.py
@@ -1,8 +1,10 @@
 # SPDX-License-Identifier: MIT
-"""Dispatch coverage for reusing the V4 H=128 prefill ASM in decode."""
+"""Dispatch coverage for V4 FP8 ASM decode paths and split tuning."""
 
 from __future__ import annotations
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -111,3 +113,85 @@ def fake_prefill_asm(**kwargs):
 )
 def test_ineligible_decode_keeps_dedicated_asm(monkeypatch, kwargs):
     assert _dispatch(monkeypatch, **kwargs) == "decode"
+
+
+@pytest.mark.parametrize(
+    "requests,expected",
+    [(1, 5), (2, 4), (3, 3), (4, 2), (5, 1), (32, 1), (64, 1), (256, 1)],
+)
+def test_aiter_csa_uses_graphsafe_batch_split_table(monkeypatch, requests, expected):
+    monkeypatch.setattr(paged_decode, "_device_arch", lambda _index: "gfx950")
+    q_packed = SimpleNamespace(
+        shape=(requests * 7, 128, 512), device=SimpleNamespace(index=0)
+    )
+
+    assert (
+        paged_decode._v4_aiter_fp8_decode_splits(
+            q_packed,
+            query_group=7,
+            kv_kind="csa",
+        )
+        == expected
+    )
+
+
+@pytest.mark.parametrize(
+    "tokens,heads,query_group,kv_kind,arch",
+    [
+        (6 * 7, 128, 7, "hca", "gfx950"),
+        (6 * 7, 64, 7, "csa", "gfx950"),
+        (6 * 4, 128, 4, "csa", "gfx950"),
+        (6 * 7, 128, 7, "csa", "gfx1250"),
+    ],
+)
+def test_aiter_split_table_rejects_unqualified_shapes(
+    monkeypatch, tokens, heads, query_group, kv_kind, arch
+):
+    monkeypatch.setattr(paged_decode, "_device_arch", lambda _index: arch)
+    q_packed = SimpleNamespace(
+        shape=(tokens, heads, 512), device=SimpleNamespace(index=0)
+    )
+
+    assert (
+        paged_decode._v4_aiter_fp8_decode_splits(
+            q_packed,
+            query_group=query_group,
+            kv_kind=kv_kind,
+        )
+        is None
+    )
+
+
+def test_dedicated_asm_receives_selected_split_count(monkeypatch):
+    captured = {}
+
+    monkeypatch.setattr(paged_decode.envs, "ATOM_USE_V4_PREFILL_ASM_FOR_DECODE", False)
+    monkeypatch.setattr(
+        paged_decode,
+        "_v4_aiter_fp8_decode_splits",
+        lambda q_packed, *, query_group, kv_kind: 3,
+    )
+
+    def fake_decode(*args, **kwargs):
+        captured.update(kwargs)
+        return "decode"
+
+    monkeypatch.setattr(paged_decode, "_sparse_attn_v4_paged_decode_asm", fake_decode)
+    n, heads = 21, 128
+    result = paged_decode.sparse_attn_v4_paged_decode(
+        q=None,
+        unified_kv=torch.empty((4, 512)),
+        kv_indices=torch.empty(0, dtype=torch.int32),
+        kv_indptr=torch.zeros(n + 1, dtype=torch.int32),
+        attn_sink=torch.empty(heads),
+        softmax_scale=512**-0.5,
+        unified_kv_rope=torch.empty((4, 64)),
+        q_packed_in=torch.empty((n, heads, 512)),
+        q_rope_in=torch.empty((n, heads, 64)),
+        qo_indptr=torch.arange(n + 1, dtype=torch.int32),
+        query_group=7,
+        kv_kind="csa",
+    )
+
+    assert result == "decode"
+    assert captured["num_kv_splits"] == 3

From 9bf852787184e70674db856cb6ba7fcddd46cfdf Mon Sep 17 00:00:00 2001
From: yihonglie <hyi@amd.com>
Date: Fri, 25 Sep 2026 22:51:23 -0500
Subject: [PATCH 2/5] style: use functools cache for device arch

---
 atom/model_ops/v4_kernels/paged_decode.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/atom/model_ops/v4_kernels/paged_decode.py b/atom/model_ops/v4_kernels/paged_decode.py
index dd2d9d440d..ab3d43b8ae 100644
--- a/atom/model_ops/v4_kernels/paged_decode.py
+++ b/atom/model_ops/v4_kernels/paged_decode.py
@@ -84,7 +84,7 @@
 _FP8_DTYPE = torch.float8_e4m3fnuz
 
 
-@functools.lru_cache(maxsize=None)
+@functools.cache
 def _device_arch(device_index: int) -> str:
     """Return the base GCN architecture name without feature suffixes."""
     return getattr(

From 87521b21599b112741fef0268685a5b6b83ea8e6 Mon Sep 17 00:00:00 2001
From: yihonglie <hyi@amd.com>
Date: Sun, 27 Sep 2026 08:20:51 -0500
Subject: [PATCH 3/5] perf(dsv4): steer AITER decode split pick with
 kv_len_hint

AITER's V4 NM wrapper can now pick num_kv_splits from a latency model when
told the per-token KV length to plan for (`kv_len_hint`). ATOM's decode
kv_indices buffers are capacity-sized, so AITER cannot derive that length
itself; the builder publishes the worst case per layer type and the decode
call forwards it:

- CSA: window + index_topk
- HCA: window + max_model_len // 128

Both depend only on config, so the pick stays fixed across CUDA graph
capture and replay. With an AITER build that predates `kv_len_hint`
(detected from the wrapper signature), CSA keeps the batch split table and
HCA keeps AITER's own auto pick.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
---
 atom/model_ops/attentions/deepseek_v4_attn.py |  7 +++
 atom/model_ops/v4_kernels/paged_decode.py     | 45 ++++++++++++++++++-
 atom/models/deepseek_v4.py                    |  6 +++
 tests/test_v4_prefill_asm_decode_dispatch.py  | 36 ++++++++++++---
 4 files changed, 86 insertions(+), 8 deletions(-)

diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py
index fb5214d4f2..87a8a07581 100644
--- a/atom/model_ops/attentions/deepseek_v4_attn.py
+++ b/atom/model_ops/attentions/deepseek_v4_attn.py
@@ -296,6 +296,11 @@ class AttentionMetaData_DSV4(AttentionMetaData):
     """[padded_T+1] int32 GPU — packed cumsum of per-token HCA kv_len
     (= `min(positions[t]+1, win) + (positions[t]+1)//128`). Padded tail = last
     value."""
+    kv_len_hint_csa: int | None = None
+    """Longest per-token CSA kv_len (`win + index_topk`); steers aiter's fp8
+    decode split planner."""
+    kv_len_hint_hca: int | None = None
+    """Longest per-token HCA kv_len (`win + max_model_len // 128`)."""
     envelope_rows: int = 0
     """Rows one V4 block takes across every layer of the pool — the stride from
     one block's compressed rows to the next in a layer's view of a plane. What
@@ -3799,6 +3804,8 @@ def _attach_v4_paged_decode_meta(
             win=win,
             index_topk=index_topk,
         )
+        attn_metadata.kv_len_hint_csa = win + index_topk
+        attn_metadata.kv_len_hint_hca = win + self.max_committed_hca
 
         # Expand block tables per query row so the unchanged aiter paged-MQA
         # kernels can run once with shape `[decode_rows, 1, ...]`. Source and
diff --git a/atom/model_ops/v4_kernels/paged_decode.py b/atom/model_ops/v4_kernels/paged_decode.py
index ab3d43b8ae..e82bcb366e 100644
--- a/atom/model_ops/v4_kernels/paged_decode.py
+++ b/atom/model_ops/v4_kernels/paged_decode.py
@@ -49,6 +49,7 @@
 from __future__ import annotations
 
 import functools
+import inspect
 import os
 
 import torch
@@ -110,6 +111,9 @@ def _v4_aiter_fp8_decode_splits(
 
     HCA remains on AITER's ragged split planner: its heterogeneous long-KV
     vectors have a different optimum from the CSA table.
+
+    Only used with aiter builds that predate the cost-model planner (see
+    ``_aiter_v4_nm_split_planner``).
     """
     shape = getattr(q_packed, "shape", None)
     device = getattr(q_packed, "device", None)
@@ -143,6 +147,39 @@ def _v4_aiter_fp8_decode_splits(
     return 1
 
 
+@functools.cache
+def _aiter_v4_nm_split_planner() -> bool:
+    """Whether aiter's V4 NM wrapper has the cost-model split planner.
+
+    That aiter takes ``kv_len_hint``; older builds only have the
+    occupancy-only auto pick, which the CSA table above works around.
+    """
+    import aiter.mla
+
+    return "kv_len_hint" in inspect.signature(aiter.mla.mla_decode_fwd_v4_nm).parameters
+
+
+def _v4_aiter_decode_split_kwargs(
+    q_packed: torch.Tensor,
+    *,
+    query_group: int,
+    kv_kind: str,
+    kv_len_hint: int | None,
+) -> dict:
+    """Split arguments for ``mla_decode_fwd_v4_nm``.
+
+    With the aiter cost-model planner, ``kv_len_hint`` (the longest per-token
+    kv_len the call can see) steers the split pick. Older aiter keeps the
+    measured CSA batch table.
+    """
+    if _aiter_v4_nm_split_planner():
+        return {} if kv_len_hint is None else {"kv_len_hint": kv_len_hint}
+    splits = _v4_aiter_fp8_decode_splits(
+        q_packed, query_group=query_group, kv_kind=kv_kind
+    )
+    return {} if splits is None else {"num_kv_splits": splits}
+
+
 @functools.lru_cache(maxsize=1)
 def _cu_count() -> int:
     """Compute-unit count of the active GPU, queried once via aiter.
@@ -1056,6 +1093,7 @@ def _sparse_attn_v4_paged_decode_asm(
     qo_indptr: torch.Tensor,
     kv_last_page_lens: torch.Tensor | None = None,  # unused on v4 nm (page_size=1)
     num_kv_splits: int | None = None,
+    kv_len_hint: int | None = None,
 ) -> torch.Tensor:
     """Native 2buff fp8 V4 decode via the aiter assembly kernel
     ``mla_decode_fwd_v4_nm`` (ROCm/aiter#3112, mi350/gfx950).
@@ -1165,6 +1203,8 @@ def _sparse_attn_v4_paged_decode_asm(
         (N, H, V4_DIM_NOPE + V4_DIM_ROPE), dtype=torch.bfloat16, device=device
     )
 
+    # Only forward `kv_len_hint` when set, so aiter builds without it work.
+    split_kwargs = {} if kv_len_hint is None else {"kv_len_hint": kv_len_hint}
     aiter.mla.mla_decode_fwd_v4_nm(
         q_packed,
         q_rope,
@@ -1178,6 +1218,7 @@ def _sparse_attn_v4_paged_decode_asm(
         sink=attn_sink,
         sm_scale=softmax_scale,
         num_kv_splits=num_kv_splits,
+        **split_kwargs,
     )
     # Drop padded heads. The slice is a non-contiguous view, so .contiguous()
     # gives downstream a dense tensor; no-op (and no copy) when H was unpadded.
@@ -1202,6 +1243,7 @@ def sparse_attn_v4_paged_decode(
     prefix: str = "",
     query_group: int = 1,
     kv_kind: str = "",
+    kv_len_hint: int | None = None,
 ) -> torch.Tensor:
     """V4 decode sparse attention over a unified KV pool with paged indices.
 
@@ -1246,10 +1288,11 @@ def sparse_attn_v4_paged_decode(
             q_rope_in,
             qo_indptr=qo_indptr,
             kv_last_page_lens=kv_last_page_lens,
-            num_kv_splits=_v4_aiter_fp8_decode_splits(
+            **_v4_aiter_decode_split_kwargs(
                 q_packed_in,
                 query_group=query_group,
                 kv_kind=kv_kind,
+                kv_len_hint=kv_len_hint,
             ),
         )
     gfx = get_gfx()
diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py
index 4691bed778..ae0e856001 100644
--- a/atom/models/deepseek_v4.py
+++ b/atom/models/deepseek_v4.py
@@ -3340,6 +3340,12 @@ def _sparse_attention(
                     7 if attn_md.min_seqlen_q == attn_md.max_seqlen_q == 7 else 1
                 ),
                 kv_kind="csa" if ratio == 4 else "hca" if ratio == 128 else "swa",
+                # Bridged metadata (e.g. sglang) may not carry the planner hints.
+                kv_len_hint=getattr(
+                    attn_md,
+                    {4: "kv_len_hint_csa", 128: "kv_len_hint_hca"}.get(ratio, ""),
+                    None,
+                ),
             )  # [S, H, head_dim]
         else:
             # Two-source paged prefill: prefix from `unified_kv` (per-ratio
diff --git a/tests/test_v4_prefill_asm_decode_dispatch.py b/tests/test_v4_prefill_asm_decode_dispatch.py
index 8ae5d33c6d..325e0c0b4e 100644
--- a/tests/test_v4_prefill_asm_decode_dispatch.py
+++ b/tests/test_v4_prefill_asm_decode_dispatch.py
@@ -162,22 +162,23 @@ def test_aiter_split_table_rejects_unqualified_shapes(
     )
 
 
-def test_dedicated_asm_receives_selected_split_count(monkeypatch):
+def _decode_split_kwargs(monkeypatch, *, planner, n=21, **kwargs):
     captured = {}
 
     monkeypatch.setattr(paged_decode.envs, "ATOM_USE_V4_PREFILL_ASM_FOR_DECODE", False)
+    monkeypatch.setattr(paged_decode, "_aiter_v4_nm_split_planner", lambda: planner)
     monkeypatch.setattr(
         paged_decode,
         "_v4_aiter_fp8_decode_splits",
         lambda q_packed, *, query_group, kv_kind: 3,
     )
 
-    def fake_decode(*args, **kwargs):
-        captured.update(kwargs)
+    def fake_decode(*args, **kw):
+        captured.update(kw)
         return "decode"
 
     monkeypatch.setattr(paged_decode, "_sparse_attn_v4_paged_decode_asm", fake_decode)
-    n, heads = 21, 128
+    heads = 128
     result = paged_decode.sparse_attn_v4_paged_decode(
         q=None,
         unified_kv=torch.empty((4, 512)),
@@ -190,8 +191,29 @@ def fake_decode(*args, **kwargs):
         q_rope_in=torch.empty((n, heads, 64)),
         qo_indptr=torch.arange(n + 1, dtype=torch.int32),
         query_group=7,
-        kv_kind="csa",
+        **kwargs,
     )
-
     assert result == "decode"
-    assert captured["num_kv_splits"] == 3
+    return {k: v for k, v in captured.items() if k in _SPLIT_KEYS}
+
+
+_SPLIT_KEYS = ("num_kv_splits", "kv_len_hint")
+
+
+def test_legacy_aiter_receives_csa_split_table(monkeypatch):
+    got = _decode_split_kwargs(
+        monkeypatch, planner=False, kv_kind="csa", kv_len_hint=1152
+    )
+    assert got == {"num_kv_splits": 3}
+
+
+def test_planner_aiter_receives_kv_len_hint(monkeypatch):
+    got = _decode_split_kwargs(
+        monkeypatch, planner=True, kv_kind="csa", kv_len_hint=1152
+    )
+    assert got == {"kv_len_hint": 1152}
+
+
+def test_planner_aiter_without_hint_uses_auto(monkeypatch):
+    got = _decode_split_kwargs(monkeypatch, planner=True, kv_kind="swa")
+    assert got == {}

From 2aaa907cdc4df3f0966e948dfaefcd57420b048b Mon Sep 17 00:00:00 2001
From: yihonglie <hyi@amd.com>
Date: Mon, 28 Sep 2026 01:05:33 -0500
Subject: [PATCH 4/5] perf(dsv4): pass AITER decode split plans through
 attention metadata

Replace kv_len_hint with an explicit split plan. The V4 metadata builder
calls AITER's get_mla_v4_nm_split_plan once per decode step for each sparse
layer type and stores (num_kv_splits, split_indptr) on the attention
metadata; the fp8 ASM decode forwards both to mla_decode_fwd_v4_nm.

- Plans are sized for the padded row count and the longest kv_len of the
  layer type (SWA: window, CSA: window + index_topk, HCA: window +
  max_model_len // 128), so capture and replay agree.
- split_indptr lives in persistent forward_vars buffers (per ubatch prefix)
  that the builder rewrites in place before each forward; graphs of
  different sizes share them.
- MTP draft steps clear the SWA plan (it is sized for the verify rows) and
  fall back to AITER's own pick.
- AITER builds without get_mla_v4_nm_split_plan keep the CSA batch table.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
---
 atom/model_ops/attentions/deepseek_v4_attn.py | 62 +++++++++++++--
 atom/model_ops/v4_kernels/paged_decode.py     | 44 ++++++-----
 atom/models/deepseek_v4.py                    |  8 +-
 tests/test_v4_prefill_asm_decode_dispatch.py  | 79 +++++++++++++++++--
 4 files changed, 155 insertions(+), 38 deletions(-)

diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py
index 87a8a07581..384838df46 100644
--- a/atom/model_ops/attentions/deepseek_v4_attn.py
+++ b/atom/model_ops/attentions/deepseek_v4_attn.py
@@ -296,11 +296,13 @@ class AttentionMetaData_DSV4(AttentionMetaData):
     """[padded_T+1] int32 GPU — packed cumsum of per-token HCA kv_len
     (= `min(positions[t]+1, win) + (positions[t]+1)//128`). Padded tail = last
     value."""
-    kv_len_hint_csa: int | None = None
-    """Longest per-token CSA kv_len (`win + index_topk`); steers aiter's fp8
-    decode split planner."""
-    kv_len_hint_hca: int | None = None
-    """Longest per-token HCA kv_len (`win + max_model_len // 128`)."""
+    split_plan_swa: tuple[int, torch.Tensor] | None = None
+    """`(num_kv_splits, split_indptr[padded_T+1])` for the fp8 decode ASM
+    kernel on SWA layers, from aiter's `get_mla_v4_nm_split_plan`."""
+    split_plan_csa: tuple[int, torch.Tensor] | None = None
+    """Same for CSA layers."""
+    split_plan_hca: tuple[int, torch.Tensor] | None = None
+    """Same for HCA layers."""
     envelope_rows: int = 0
     """Rows one V4 block takes across every layer of the pool — the stride from
     one block's compressed rows to the next in a layer's view of a plane. What
@@ -715,6 +717,15 @@ def __init__(self, model_runner):
         # Worst-case HCA per-token committed compress count
         # (= max_model_len // 128 for V4-Pro = 8192 at 1M context).
         self.max_committed_hca = model_runner.config.max_model_len // 128
+        # Heads the fp8 decode ASM kernel runs per rank: the model's local head
+        # count, which the aiter wrapper pads up to 16 when smaller.
+        from aiter.dist.parallel_state import get_tensor_model_parallel_world_size
+
+        self._decode_local_heads = max(
+            16,
+            model_runner.config.hf_config.num_attention_heads
+            // get_tensor_model_parallel_world_size(),
+        )
 
         # Sparse-attn + per-fwd metadata buffers (CG-A: pre-allocate for fixed
         # GPU pointers, prerequisite for CUDAGraph capture). All H2D copies in
@@ -2612,6 +2623,9 @@ def prepare_mtp_decode(
         attn_metadata.max_seqlen_q = 1
         attn_metadata.kv_indices_swa = swa_indices_buf
         attn_metadata.kv_indptr_swa = swa_indptr
+        # The verify step's SWA plan is sized for its row count, not the draft's;
+        # the draft falls back to aiter's own split pick.
+        attn_metadata.split_plan_swa = None
         attn_metadata.batch_id_per_q_token = batch_id_per_q_token
 
         # fp8 asm decode per-token index tensors. MTP draft step is 1-token-per-
@@ -3682,6 +3696,37 @@ def _attach_v4_per_fwd_meta(
             buf_prefix_ubatch=buf_prefix_ubatch,
         )
 
+    def _attach_v4_decode_split_plans(
+        self, attn_metadata, T_pad: int, buf_prefix_ubatch: str
+    ) -> None:
+        """Build aiter's fp8 decode KV split plan per sparse layer type.
+
+        The plan is a function of the padded row count and the longest kv_len
+        a token of that layer type can see, both fixed for a captured graph,
+        so capture and every replay agree. `split_indptr` lives in a
+        persistent buffer the graph reads; this rewrites it (one launch per
+        type, no host sync) before each forward, since graphs of different
+        `T_pad` share the buffer.
+        """
+        from atom.model_ops.v4_kernels.paged_decode import aiter_v4_nm_split_plan_fn
+
+        plan_fn = aiter_v4_nm_split_plan_fn()
+        if not self._kv_fp8 or plan_fn is None:
+            return
+        var = self.model_runner.forward_vars
+        win = self.window_size
+        for kind, kv_len in (
+            ("swa", win),
+            ("csa", win + self.index_topk),
+            ("hca", win + self.max_committed_hca),
+        ):
+            buf = var[f"{buf_prefix_ubatch}v4_split_indptr_{kind}"]
+            setattr(
+                attn_metadata,
+                f"split_plan_{kind}",
+                plan_fn(T_pad, self._decode_local_heads, kv_len, split_indptr=buf),
+            )
+
     def _attach_v4_paged_decode_meta(
         self,
         attn_metadata,
@@ -3804,8 +3849,7 @@ def _attach_v4_paged_decode_meta(
             win=win,
             index_topk=index_topk,
         )
-        attn_metadata.kv_len_hint_csa = win + index_topk
-        attn_metadata.kv_len_hint_hca = win + self.max_committed_hca
+        self._attach_v4_decode_split_plans(attn_metadata, T_pad, buf_prefix_ubatch)
 
         # Expand block tables per query row so the unchanged aiter paged-MQA
         # kernels can run once with shape `[decode_rows, 1, ...]`. Source and
@@ -4629,6 +4673,8 @@ def _alloc_v4_metadata_buffers(self) -> None:
         bufs["v4_kv_indptr_swa"] = torch.zeros(T_dec + 1, **i32)
         bufs["v4_kv_indptr_csa"] = torch.zeros(T_dec + 1, **i32)
         bufs["v4_kv_indptr_hca"] = torch.zeros(T_dec + 1, **i32)
+        for kind in ("swa", "csa", "hca"):
+            bufs[f"v4_split_indptr_{kind}"] = torch.zeros(T_dec + 1, **i32)
         bufs["v4_csa_n_committed_per_token"] = torch.zeros(T_dec, **i32)
         # Device-only, and as wide as the `block_tables` it gathers from: a host
         # mirror would be `T_dec * cols * 4` of pinned memory nothing writes.
@@ -4828,6 +4874,8 @@ def _alloc_v4_ubatch_metadata_buffers(
             bufs[f"{p}v4_kv_indptr_swa"] = torch.zeros(T_dec + 1, **i32)
             bufs[f"{p}v4_kv_indptr_csa"] = torch.zeros(T_dec + 1, **i32)
             bufs[f"{p}v4_kv_indptr_hca"] = torch.zeros(T_dec + 1, **i32)
+            for kind in ("swa", "csa", "hca"):
+                bufs[f"{p}v4_split_indptr_{kind}"] = torch.zeros(T_dec + 1, **i32)
             bufs[f"{p}v4_csa_n_committed_per_token"] = torch.zeros(T_dec, **i32)
             if not self._indexer_fp4:
                 bufs[f"{p}v4_block_tables_per_token"] = torch.zeros(
diff --git a/atom/model_ops/v4_kernels/paged_decode.py b/atom/model_ops/v4_kernels/paged_decode.py
index e82bcb366e..7a911af2ed 100644
--- a/atom/model_ops/v4_kernels/paged_decode.py
+++ b/atom/model_ops/v4_kernels/paged_decode.py
@@ -49,7 +49,6 @@
 from __future__ import annotations
 
 import functools
-import inspect
 import os
 
 import torch
@@ -112,8 +111,8 @@ def _v4_aiter_fp8_decode_splits(
     HCA remains on AITER's ragged split planner: its heterogeneous long-KV
     vectors have a different optimum from the CSA table.
 
-    Only used with aiter builds that predate the cost-model planner (see
-    ``_aiter_v4_nm_split_planner``).
+    Only used with aiter builds that predate ``get_mla_v4_nm_split_plan`` (see
+    ``aiter_v4_nm_split_plan_fn``).
     """
     shape = getattr(q_packed, "shape", None)
     device = getattr(q_packed, "device", None)
@@ -148,15 +147,15 @@ def _v4_aiter_fp8_decode_splits(
 
 
 @functools.cache
-def _aiter_v4_nm_split_planner() -> bool:
-    """Whether aiter's V4 NM wrapper has the cost-model split planner.
+def aiter_v4_nm_split_plan_fn():
+    """aiter's ``get_mla_v4_nm_split_plan``, or None on builds without it.
 
-    That aiter takes ``kv_len_hint``; older builds only have the
-    occupancy-only auto pick, which the CSA table above works around.
+    Those older builds only have the occupancy-only auto pick, which the CSA
+    table above works around.
     """
     import aiter.mla
 
-    return "kv_len_hint" in inspect.signature(aiter.mla.mla_decode_fwd_v4_nm).parameters
+    return getattr(aiter.mla, "get_mla_v4_nm_split_plan", None)
 
 
 def _v4_aiter_decode_split_kwargs(
@@ -164,16 +163,19 @@ def _v4_aiter_decode_split_kwargs(
     *,
     query_group: int,
     kv_kind: str,
-    kv_len_hint: int | None,
+    split_plan: tuple[int, torch.Tensor] | None,
 ) -> dict:
     """Split arguments for ``mla_decode_fwd_v4_nm``.
 
-    With the aiter cost-model planner, ``kv_len_hint`` (the longest per-token
-    kv_len the call can see) steers the split pick. Older aiter keeps the
-    measured CSA batch table.
+    Uses the metadata builder's plan when it covers this call's rows; the
+    kernel launch then reads the plan's persistent ``split_indptr``. Without a
+    plan, aiter builds that have the planner fall back to their own pick and
+    older ones keep the measured CSA batch table.
     """
-    if _aiter_v4_nm_split_planner():
-        return {} if kv_len_hint is None else {"kv_len_hint": kv_len_hint}
+    if split_plan is not None and split_plan[1].shape[0] >= q_packed.shape[0] + 1:
+        return {"num_kv_splits": split_plan[0], "split_indptr": split_plan[1]}
+    if aiter_v4_nm_split_plan_fn() is not None:
+        return {}
     splits = _v4_aiter_fp8_decode_splits(
         q_packed, query_group=query_group, kv_kind=kv_kind
     )
@@ -1093,7 +1095,7 @@ def _sparse_attn_v4_paged_decode_asm(
     qo_indptr: torch.Tensor,
     kv_last_page_lens: torch.Tensor | None = None,  # unused on v4 nm (page_size=1)
     num_kv_splits: int | None = None,
-    kv_len_hint: int | None = None,
+    split_indptr: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Native 2buff fp8 V4 decode via the aiter assembly kernel
     ``mla_decode_fwd_v4_nm`` (ROCm/aiter#3112, mi350/gfx950).
@@ -1203,8 +1205,10 @@ def _sparse_attn_v4_paged_decode_asm(
         (N, H, V4_DIM_NOPE + V4_DIM_ROPE), dtype=torch.bfloat16, device=device
     )
 
-    # Only forward `kv_len_hint` when set, so aiter builds without it work.
-    split_kwargs = {} if kv_len_hint is None else {"kv_len_hint": kv_len_hint}
+    # A uniform split plan's prefix is itself a valid plan, so it trims like
+    # qo_indptr.
+    if split_indptr is not None:
+        split_indptr = split_indptr[: N + 1]
     aiter.mla.mla_decode_fwd_v4_nm(
         q_packed,
         q_rope,
@@ -1218,7 +1222,7 @@ def _sparse_attn_v4_paged_decode_asm(
         sink=attn_sink,
         sm_scale=softmax_scale,
         num_kv_splits=num_kv_splits,
-        **split_kwargs,
+        split_indptr=split_indptr,
     )
     # Drop padded heads. The slice is a non-contiguous view, so .contiguous()
     # gives downstream a dense tensor; no-op (and no copy) when H was unpadded.
@@ -1243,7 +1247,7 @@ def sparse_attn_v4_paged_decode(
     prefix: str = "",
     query_group: int = 1,
     kv_kind: str = "",
-    kv_len_hint: int | None = None,
+    split_plan: tuple[int, torch.Tensor] | None = None,
 ) -> torch.Tensor:
     """V4 decode sparse attention over a unified KV pool with paged indices.
 
@@ -1292,7 +1296,7 @@ def sparse_attn_v4_paged_decode(
                 q_packed_in,
                 query_group=query_group,
                 kv_kind=kv_kind,
-                kv_len_hint=kv_len_hint,
+                split_plan=split_plan,
             ),
         )
     gfx = get_gfx()
diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py
index ae0e856001..f969879866 100644
--- a/atom/models/deepseek_v4.py
+++ b/atom/models/deepseek_v4.py
@@ -3340,10 +3340,12 @@ def _sparse_attention(
                     7 if attn_md.min_seqlen_q == attn_md.max_seqlen_q == 7 else 1
                 ),
                 kv_kind="csa" if ratio == 4 else "hca" if ratio == 128 else "swa",
-                # Bridged metadata (e.g. sglang) may not carry the planner hints.
-                kv_len_hint=getattr(
+                # Bridged metadata (e.g. sglang) may not carry split plans.
+                split_plan=getattr(
                     attn_md,
-                    {4: "kv_len_hint_csa", 128: "kv_len_hint_hca"}.get(ratio, ""),
+                    {0: "split_plan_swa", 4: "split_plan_csa"}.get(
+                        ratio, "split_plan_hca"
+                    ),
                     None,
                 ),
             )  # [S, H, head_dim]
diff --git a/tests/test_v4_prefill_asm_decode_dispatch.py b/tests/test_v4_prefill_asm_decode_dispatch.py
index 325e0c0b4e..6bd2c49d82 100644
--- a/tests/test_v4_prefill_asm_decode_dispatch.py
+++ b/tests/test_v4_prefill_asm_decode_dispatch.py
@@ -166,7 +166,11 @@ def _decode_split_kwargs(monkeypatch, *, planner, n=21, **kwargs):
     captured = {}
 
     monkeypatch.setattr(paged_decode.envs, "ATOM_USE_V4_PREFILL_ASM_FOR_DECODE", False)
-    monkeypatch.setattr(paged_decode, "_aiter_v4_nm_split_planner", lambda: planner)
+    monkeypatch.setattr(
+        paged_decode,
+        "aiter_v4_nm_split_plan_fn",
+        lambda: (lambda *a, **k: None) if planner else None,
+    )
     monkeypatch.setattr(
         paged_decode,
         "_v4_aiter_fp8_decode_splits",
@@ -197,23 +201,82 @@ def fake_decode(*args, **kw):
     return {k: v for k, v in captured.items() if k in _SPLIT_KEYS}
 
 
-_SPLIT_KEYS = ("num_kv_splits", "kv_len_hint")
+_SPLIT_KEYS = ("num_kv_splits", "split_indptr")
 
 
 def test_legacy_aiter_receives_csa_split_table(monkeypatch):
+    got = _decode_split_kwargs(monkeypatch, planner=False, kv_kind="csa")
+    assert got == {"num_kv_splits": 3}
+
+
+@pytest.mark.parametrize("planner", [True, False])
+def test_metadata_split_plan_is_forwarded(monkeypatch, planner):
+    n = 21
+    plan = (4, torch.arange(0, 4 * (n + 1), 4, dtype=torch.int32))
     got = _decode_split_kwargs(
-        monkeypatch, planner=False, kv_kind="csa", kv_len_hint=1152
+        monkeypatch, planner=planner, n=n, kv_kind="hca", split_plan=plan
     )
-    assert got == {"num_kv_splits": 3}
+    assert got["num_kv_splits"] == 4
+    assert got["split_indptr"] is plan[1]
+
+
+def test_plan_covering_more_rows_is_forwarded(monkeypatch):
+    # Eager forwards can run fewer rows than the padded grid the plan was built
+    # for; the ASM wrapper trims the uniform plan like qo_indptr.
+    plan = (2, torch.arange(0, 2 * 29, 2, dtype=torch.int32))
+    got = _decode_split_kwargs(
+        monkeypatch, planner=True, n=21, kv_kind="csa", split_plan=plan
+    )
+    assert got["num_kv_splits"] == 2
 
 
-def test_planner_aiter_receives_kv_len_hint(monkeypatch):
+def test_plan_for_fewer_rows_falls_back(monkeypatch):
+    plan = (2, torch.arange(0, 2 * 8, 2, dtype=torch.int32))
     got = _decode_split_kwargs(
-        monkeypatch, planner=True, kv_kind="csa", kv_len_hint=1152
+        monkeypatch, planner=True, n=21, kv_kind="csa", split_plan=plan
     )
-    assert got == {"kv_len_hint": 1152}
+    assert got == {}
 
 
-def test_planner_aiter_without_hint_uses_auto(monkeypatch):
+def test_planner_aiter_without_plan_uses_auto(monkeypatch):
     got = _decode_split_kwargs(monkeypatch, planner=True, kv_kind="swa")
     assert got == {}
+
+
+@pytest.mark.parametrize("kv_fp8,planner", [(True, True), (False, True), (True, False)])
+def test_builder_attaches_split_plans(monkeypatch, kv_fp8, planner):
+    from atom.model_ops.attentions import deepseek_v4_attn
+
+    calls = []
+
+    def fake_plan(num_seqs, num_heads, kv_len, *, split_indptr):
+        calls.append((num_seqs, num_heads, kv_len))
+        return 2, split_indptr[: num_seqs + 1]
+
+    monkeypatch.setattr(
+        paged_decode,
+        "aiter_v4_nm_split_plan_fn",
+        lambda: fake_plan if planner else None,
+    )
+    bufs = {
+        f"v4_split_indptr_{k}": torch.zeros(65, dtype=torch.int32)
+        for k in ("swa", "csa", "hca")
+    }
+    builder = SimpleNamespace(
+        _kv_fp8=kv_fp8,
+        _decode_local_heads=128,
+        window_size=128,
+        index_topk=1024,
+        max_committed_hca=8192,
+        model_runner=SimpleNamespace(forward_vars=bufs),
+    )
+    md = SimpleNamespace(split_plan_swa=None, split_plan_csa=None, split_plan_hca=None)
+    deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._attach_v4_decode_split_plans(
+        builder, md, 28, ""
+    )
+
+    if kv_fp8 and planner:
+        assert calls == [(28, 128, 128), (28, 128, 1152), (28, 128, 8320)]
+        assert md.split_plan_csa[0] == 2 and md.split_plan_csa[1].shape[0] == 29
+    else:
+        assert calls == [] and md.split_plan_csa is None

From 94d1b24f6682674a45a684b7512d9b0cd8427a1e Mon Sep 17 00:00:00 2001
From: yihonglie <hyi@amd.com>
Date: Mon, 28 Sep 2026 02:20:07 -0500
Subject: [PATCH 5/5] refactor(dsv4): one split-plan path, owned by whoever
 builds the decode CSR

- v4_decode_split_plan (paged_decode) is the single place that knows about
  aiter's get_mla_v4_nm_split_plan; it returns None on older aiter, which
  leaves the split pick to aiter.
- Each owner of a decode CSR builds its plan next to it, in its own
  persistent split_indptr buffer: the V4 metadata builder (SWA/CSA/HCA for
  verify, SWA for the MTP draft) and DSpark's DSparkIndexBuffers (built with
  the CSR at stage 0, read by every stage).
- sparse_attn_v4_paged_decode takes one split_plan argument; the model picks
  it in the same ratio branch as kv_indices / kv_indptr.
- Drop the legacy CSA batch table and the query_group / kv_kind parameters
  it needed.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
---
 atom/model_ops/attentions/deepseek_v4_attn.py |  89 ++++----
 atom/model_ops/v4_kernels/__init__.py         |   2 +
 .../v4_kernels/dspark_fp8_indices.py          |  10 +-
 atom/model_ops/v4_kernels/paged_decode.py     | 133 +++--------
 atom/models/deepseek_v4.py                    |  16 +-
 atom/models/deepseek_v4_dspark.py             |   3 +-
 tests/test_v4_prefill_asm_decode_dispatch.py  | 216 +++++++-----------
 7 files changed, 171 insertions(+), 298 deletions(-)

diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py
index 384838df46..d3538db3d4 100644
--- a/atom/model_ops/attentions/deepseek_v4_attn.py
+++ b/atom/model_ops/attentions/deepseek_v4_attn.py
@@ -48,6 +48,7 @@
 import numpy as np
 import torch
 from aiter import dtypes
+from aiter.dist.parallel_state import get_tensor_model_parallel_world_size
 from aiter.jit.utils.chip_info import get_gfx
 
 from atom.distributed.pcp_utils import (
@@ -119,6 +120,7 @@
     build_v4_paged_decode_indptr,
     fp4_indexer_enabled,
     plan_context_lens,
+    v4_decode_split_plan,
     write_v4_paged_decode_indices,
     write_v4_paged_prefill_indices,
 )
@@ -297,8 +299,8 @@ class AttentionMetaData_DSV4(AttentionMetaData):
     (= `min(positions[t]+1, win) + (positions[t]+1)//128`). Padded tail = last
     value."""
     split_plan_swa: tuple[int, torch.Tensor] | None = None
-    """`(num_kv_splits, split_indptr[padded_T+1])` for the fp8 decode ASM
-    kernel on SWA layers, from aiter's `get_mla_v4_nm_split_plan`."""
+    """aiter `MlaV4NmSplitPlan` (`num_kv_splits`, `split_indptr[padded_T+1]`)
+    for the fp8 decode ASM kernel on SWA layers. None: aiter picks."""
     split_plan_csa: tuple[int, torch.Tensor] | None = None
     """Same for CSA layers."""
     split_plan_hca: tuple[int, torch.Tensor] | None = None
@@ -717,14 +719,10 @@ def __init__(self, model_runner):
         # Worst-case HCA per-token committed compress count
         # (= max_model_len // 128 for V4-Pro = 8192 at 1M context).
         self.max_committed_hca = model_runner.config.max_model_len // 128
-        # Heads the fp8 decode ASM kernel runs per rank: the model's local head
-        # count, which the aiter wrapper pads up to 16 when smaller.
-        from aiter.dist.parallel_state import get_tensor_model_parallel_world_size
-
-        self._decode_local_heads = max(
-            16,
+        # Heads per rank the fp8 decode ASM kernel runs; sizes its split plans.
+        self._local_heads = (
             model_runner.config.hf_config.num_attention_heads
-            // get_tensor_model_parallel_world_size(),
+            // get_tensor_model_parallel_world_size()
         )
 
         # Sparse-attn + per-fwd metadata buffers (CG-A: pre-allocate for fixed
@@ -2623,9 +2621,9 @@ def prepare_mtp_decode(
         attn_metadata.max_seqlen_q = 1
         attn_metadata.kv_indices_swa = swa_indices_buf
         attn_metadata.kv_indptr_swa = swa_indptr
-        # The verify step's SWA plan is sized for its row count, not the draft's;
-        # the draft falls back to aiter's own split pick.
-        attn_metadata.split_plan_swa = None
+        attn_metadata.split_plan_swa = self._decode_split_plan(
+            running_bs, self.window_size, var["v4_draft_split_indptr"]
+        )
         attn_metadata.batch_id_per_q_token = batch_id_per_q_token
 
         # fp8 asm decode per-token index tensors. MTP draft step is 1-token-per-
@@ -3696,36 +3694,19 @@ def _attach_v4_per_fwd_meta(
             buf_prefix_ubatch=buf_prefix_ubatch,
         )
 
-    def _attach_v4_decode_split_plans(
-        self, attn_metadata, T_pad: int, buf_prefix_ubatch: str
-    ) -> None:
-        """Build aiter's fp8 decode KV split plan per sparse layer type.
-
-        The plan is a function of the padded row count and the longest kv_len
-        a token of that layer type can see, both fixed for a captured graph,
-        so capture and every replay agree. `split_indptr` lives in a
-        persistent buffer the graph reads; this rewrites it (one launch per
-        type, no host sync) before each forward, since graphs of different
-        `T_pad` share the buffer.
-        """
-        from atom.model_ops.v4_kernels.paged_decode import aiter_v4_nm_split_plan_fn
+    def _decode_split_plan(
+        self, rows: int, kv_len: int, split_indptr: torch.Tensor
+    ) -> tuple[int, torch.Tensor] | None:
+        """Split plan for the fp8 decode ASM kernel (see `v4_decode_split_plan`).
 
-        plan_fn = aiter_v4_nm_split_plan_fn()
-        if not self._kv_fp8 or plan_fn is None:
-            return
-        var = self.model_runner.forward_vars
-        win = self.window_size
-        for kind, kv_len in (
-            ("swa", win),
-            ("csa", win + self.index_topk),
-            ("hca", win + self.max_committed_hca),
-        ):
-            buf = var[f"{buf_prefix_ubatch}v4_split_indptr_{kind}"]
-            setattr(
-                attn_metadata,
-                f"split_plan_{kind}",
-                plan_fn(T_pad, self._decode_local_heads, kv_len, split_indptr=buf),
-            )
+        `rows` and `kv_len` are fixed for a captured graph, so capture and every
+        replay agree. `split_indptr` is a persistent buffer the graph reads;
+        graphs of different sizes share it, so every forward rebuilds its plan
+        (one in-place launch, no host sync) before running.
+        """
+        if not self._kv_fp8:
+            return None
+        return v4_decode_split_plan(rows, self._local_heads, kv_len, split_indptr)
 
     def _attach_v4_paged_decode_meta(
         self,
@@ -3849,7 +3830,19 @@ def _attach_v4_paged_decode_meta(
             win=win,
             index_topk=index_topk,
         )
-        self._attach_v4_decode_split_plans(attn_metadata, T_pad, buf_prefix_ubatch)
+        # aiter decode split plans: sized by the rows this forward runs and the
+        # longest kv_len a token of each layer type can see.
+        attn_metadata.split_plan_swa = self._decode_split_plan(
+            T_pad, win, var[f"{buf_prefix_ubatch}v4_split_indptr_swa"]
+        )
+        attn_metadata.split_plan_csa = self._decode_split_plan(
+            T_pad, win + index_topk, var[f"{buf_prefix_ubatch}v4_split_indptr_csa"]
+        )
+        attn_metadata.split_plan_hca = self._decode_split_plan(
+            T_pad,
+            win + self.max_committed_hca,
+            var[f"{buf_prefix_ubatch}v4_split_indptr_hca"],
+        )
 
         # Expand block tables per query row so the unchanged aiter paged-MQA
         # kernels can run once with shape `[decode_rows, 1, ...]`. Source and
@@ -4673,8 +4666,10 @@ def _alloc_v4_metadata_buffers(self) -> None:
         bufs["v4_kv_indptr_swa"] = torch.zeros(T_dec + 1, **i32)
         bufs["v4_kv_indptr_csa"] = torch.zeros(T_dec + 1, **i32)
         bufs["v4_kv_indptr_hca"] = torch.zeros(T_dec + 1, **i32)
-        for kind in ("swa", "csa", "hca"):
-            bufs[f"v4_split_indptr_{kind}"] = torch.zeros(T_dec + 1, **i32)
+        # Written in place by `_decode_split_plan` before each decode forward.
+        bufs["v4_split_indptr_swa"] = torch.zeros(T_dec + 1, **i32)
+        bufs["v4_split_indptr_csa"] = torch.zeros(T_dec + 1, **i32)
+        bufs["v4_split_indptr_hca"] = torch.zeros(T_dec + 1, **i32)
         bufs["v4_csa_n_committed_per_token"] = torch.zeros(T_dec, **i32)
         # Device-only, and as wide as the `block_tables` it gathers from: a host
         # mirror would be `T_dec * cols * 4` of pinned memory nothing writes.
@@ -4702,6 +4697,7 @@ def _alloc_v4_metadata_buffers(self) -> None:
         # the per-fwd cost is a slice + H2D.
         bufs["v4_qo_indptr"] = CpuGpuBuffer(T_dec + 1, publication_group="v4_qo", **i32)
         bufs["v4_draft_qo_indptr"] = torch.arange(T_dec + 1, **i32)
+        bufs["v4_draft_split_indptr"] = torch.zeros(T_dec + 1, **i32)
         self._v4_qo_indptr_np = np.arange(T_dec + 1, dtype=np.int32)
         # Immutable, device-only empty CSR for reusing the H=128 sparse-prefill
         # ASM kernel in decode. Shared read-only across layers and TBO ubatches.
@@ -4874,8 +4870,9 @@ def _alloc_v4_ubatch_metadata_buffers(
             bufs[f"{p}v4_kv_indptr_swa"] = torch.zeros(T_dec + 1, **i32)
             bufs[f"{p}v4_kv_indptr_csa"] = torch.zeros(T_dec + 1, **i32)
             bufs[f"{p}v4_kv_indptr_hca"] = torch.zeros(T_dec + 1, **i32)
-            for kind in ("swa", "csa", "hca"):
-                bufs[f"{p}v4_split_indptr_{kind}"] = torch.zeros(T_dec + 1, **i32)
+            bufs[f"{p}v4_split_indptr_swa"] = torch.zeros(T_dec + 1, **i32)
+            bufs[f"{p}v4_split_indptr_csa"] = torch.zeros(T_dec + 1, **i32)
+            bufs[f"{p}v4_split_indptr_hca"] = torch.zeros(T_dec + 1, **i32)
             bufs[f"{p}v4_csa_n_committed_per_token"] = torch.zeros(T_dec, **i32)
             if not self._indexer_fp4:
                 bufs[f"{p}v4_block_tables_per_token"] = torch.zeros(
diff --git a/atom/model_ops/v4_kernels/__init__.py b/atom/model_ops/v4_kernels/__init__.py
index 753d2c38ae..6b77b185f1 100644
--- a/atom/model_ops/v4_kernels/__init__.py
+++ b/atom/model_ops/v4_kernels/__init__.py
@@ -34,6 +34,7 @@
 from atom.model_ops.v4_kernels.paged_decode import (
     sparse_attn_v4_paged_decode,
     sparse_attn_v4_paged_decode_reference,
+    v4_decode_split_plan,
 )
 from atom.model_ops.v4_kernels.paged_decode_indices import (
     build_v4_paged_decode_indptr,
@@ -89,6 +90,7 @@
     "swa_write",
     "swa_write_2buff_prepacked",
     "update_compressor_states",
+    "v4_decode_split_plan",
     "write_v4_paged_decode_indices",
     "write_v4_paged_decode_indices_reference",
     "write_v4_paged_prefill_indices",
diff --git a/atom/model_ops/v4_kernels/dspark_fp8_indices.py b/atom/model_ops/v4_kernels/dspark_fp8_indices.py
index 1ea4bdcc70..fc77e49e66 100644
--- a/atom/model_ops/v4_kernels/dspark_fp8_indices.py
+++ b/atom/model_ops/v4_kernels/dspark_fp8_indices.py
@@ -41,6 +41,7 @@
 import triton.language as tl
 
 from atom.model_ops.attentions.pool_layout.v4_pool_geometry import WindowParams
+from atom.model_ops.v4_kernels.paged_decode import v4_decode_split_plan
 from atom.model_ops.v4_kernels.pool_index import window_constexprs, window_row
 
 
@@ -142,11 +143,14 @@ class DSparkIndexBuffers:
     kv_indptr: torch.Tensor  # [max_b*T+1] int32
     draft_rows: torch.Tensor  # [max_b*T] int32
     qo_indptr: torch.Tensor  # [max_b*T+1] int32, constant ramp
+    split_indptr: torch.Tensor  # [max_b*T+1] int32, see `split_plan`
     batch_ids: torch.Tensor  # [max_b*T] int32, [0]*T ++ [1]*T ++ ..., -1 on pad
     max_batch: int
     draft_width: int  # T
     draft_window: int  # W
     built_for: int = -1
+    # The fp8 decode ASM split plan for `built_for` requests (None: aiter picks).
+    split_plan: tuple[int, torch.Tensor] | None = None
 
     @classmethod
     def allocate(
@@ -174,6 +178,7 @@ def allocate(
             kv_indptr=torch.empty(n + 1, **i32),
             draft_rows=torch.empty(n, **i32),
             qo_indptr=torch.arange(n + 1, **i32),
+            split_indptr=torch.zeros(n + 1, **i32),
             batch_ids=torch.arange(n, **i32) // draft,
             max_batch=max_batch,
             draft_width=draft,
@@ -199,11 +204,13 @@ def build(
         window: WindowParams,
         slots: torch.Tensor,  # [B] per-request ring slot
         anchors: torch.Tensor,  # [B] per-request anchor position
+        heads: int,  # local attention heads, for the split plan
     ) -> None:
         """Fill this bundle for one block, in one launch. Read it with :meth:`views`.
 
         `T` and `W` come from the bundle, so the slices can never be cut to a shape
-        the buffers were not allocated for.
+        the buffers were not allocated for. The split plan is built alongside,
+        for `B * T` rows of at most `W + T` KV each.
         """
         B = anchors.shape[0]
         T, W = self.draft_width, self.draft_window
@@ -228,6 +235,7 @@ def build(
             BLOCK_B=triton.next_power_of_2(B),
             BLOCK_K=triton.next_power_of_2(W + T),
         )
+        self.split_plan = v4_decode_split_plan(B * T, heads, W + T, self.split_indptr)
         # Last: a launch that raised must not leave the bundle claiming a batch.
         self.built_for = B
 
diff --git a/atom/model_ops/v4_kernels/paged_decode.py b/atom/model_ops/v4_kernels/paged_decode.py
index 7a911af2ed..73ce7dd434 100644
--- a/atom/model_ops/v4_kernels/paged_decode.py
+++ b/atom/model_ops/v4_kernels/paged_decode.py
@@ -84,104 +84,6 @@
 _FP8_DTYPE = torch.float8_e4m3fnuz
 
 
-@functools.cache
-def _device_arch(device_index: int) -> str:
-    """Return the base GCN architecture name without feature suffixes."""
-    return getattr(
-        torch.cuda.get_device_properties(device_index), "gcnArchName", ""
-    ).split(":", 1)[0]
-
-
-def _v4_aiter_fp8_decode_splits(
-    q_packed: torch.Tensor,
-    *,
-    query_group: int,
-    kv_kind: str,
-) -> int | None:
-    """Return the qualified gfx950 H128/q7 AITER CSA split count.
-
-    AITER #5195's V4 NM wrapper maximizes a pure CU-occupancy score when the
-    caller leaves ``num_kv_splits`` unset. For q7 this creates a sawtooth
-    schedule: nearby batches can select 1 versus 15 splits, and every
-    multi-split launch writes FP32 partial outputs before a second reduction
-    kernel. The measured batch-only table below stays CUDA Graph safe because
-    it depends only on the captured Q shape. It is robust across K384/K640/
-    K1152 and avoids reading GPU-resident ``kv_indptr`` on the host.
-
-    HCA remains on AITER's ragged split planner: its heterogeneous long-KV
-    vectors have a different optimum from the CSA table.
-
-    Only used with aiter builds that predate ``get_mla_v4_nm_split_plan`` (see
-    ``aiter_v4_nm_split_plan_fn``).
-    """
-    shape = getattr(q_packed, "shape", None)
-    device = getattr(q_packed, "device", None)
-    if shape is None or len(shape) != 3 or device is None:
-        return None
-
-    device_index = device.index
-    if device_index is None:
-        device_index = torch.cuda.current_device()
-    tokens, heads, _ = shape
-    if (
-        _device_arch(device_index) != "gfx950"
-        or heads != 128
-        or query_group != 7
-        or kv_kind != "csa"
-        or tokens % query_group
-    ):
-        return None
-
-    requests = tokens // query_group
-    if requests <= 0:
-        return None
-    if requests == 1:
-        return 5
-    if requests == 2:
-        return 4
-    if requests == 3:
-        return 3
-    if requests == 4:
-        return 2
-    return 1
-
-
-@functools.cache
-def aiter_v4_nm_split_plan_fn():
-    """aiter's ``get_mla_v4_nm_split_plan``, or None on builds without it.
-
-    Those older builds only have the occupancy-only auto pick, which the CSA
-    table above works around.
-    """
-    import aiter.mla
-
-    return getattr(aiter.mla, "get_mla_v4_nm_split_plan", None)
-
-
-def _v4_aiter_decode_split_kwargs(
-    q_packed: torch.Tensor,
-    *,
-    query_group: int,
-    kv_kind: str,
-    split_plan: tuple[int, torch.Tensor] | None,
-) -> dict:
-    """Split arguments for ``mla_decode_fwd_v4_nm``.
-
-    Uses the metadata builder's plan when it covers this call's rows; the
-    kernel launch then reads the plan's persistent ``split_indptr``. Without a
-    plan, aiter builds that have the planner fall back to their own pick and
-    older ones keep the measured CSA batch table.
-    """
-    if split_plan is not None and split_plan[1].shape[0] >= q_packed.shape[0] + 1:
-        return {"num_kv_splits": split_plan[0], "split_indptr": split_plan[1]}
-    if aiter_v4_nm_split_plan_fn() is not None:
-        return {}
-    splits = _v4_aiter_fp8_decode_splits(
-        q_packed, query_group=query_group, kv_kind=kv_kind
-    )
-    return {} if splits is None else {"num_kv_splits": splits}
-
-
 @functools.lru_cache(maxsize=1)
 def _cu_count() -> int:
     """Compute-unit count of the active GPU, queried once via aiter.
@@ -1083,6 +985,26 @@ def _sparse_attn_v4_paged_decode_prefill_asm(
     )
 
 
+def v4_decode_split_plan(
+    rows: int, heads: int, kv_len: int, split_indptr: torch.Tensor
+) -> tuple[int, torch.Tensor] | None:
+    """aiter's KV split plan for the fp8 decode ASM kernel, or None.
+
+    For a decode call of `rows` query rows (its `qo_indptr` has `rows + 1`
+    entries), `heads` local heads and at most `kv_len` KV per row. The plan's
+    `split_indptr` is written in place into the caller's persistent buffer, so
+    whoever owns a decode CSR builds its plan next to it and passes it on as
+    ``sparse_attn_v4_paged_decode(..., split_plan=plan)``. None on aiter builds
+    before ROCm/aiter#5890, which leaves the split pick to aiter.
+    """
+    import aiter.mla
+
+    plan_fn = getattr(aiter.mla, "get_mla_v4_nm_split_plan", None)
+    if plan_fn is None:
+        return None
+    return plan_fn(rows, heads, kv_len, split_indptr=split_indptr)
+
+
 def _sparse_attn_v4_paged_decode_asm(
     unified_kv: torch.Tensor,
     kv_indices: torch.Tensor,
@@ -1245,8 +1167,6 @@ def sparse_attn_v4_paged_decode(
     kv_last_page_lens: torch.Tensor | None = None,
     empty_kv_indptr: torch.Tensor | None = None,
     prefix: str = "",
-    query_group: int = 1,
-    kv_kind: str = "",
     split_plan: tuple[int, torch.Tensor] | None = None,
 ) -> torch.Tensor:
     """V4 decode sparse attention over a unified KV pool with paged indices.
@@ -1257,6 +1177,10 @@ def sparse_attn_v4_paged_decode(
     sparse-prefill ASM kernel with an empty extend stream. Both paths consume
     pre-packed fp8 Q and the fp8 NoPE + bf16 RoPE pools with no requant.
 
+    ``split_plan`` is aiter's ``MlaV4NmSplitPlan`` (``num_kv_splits``,
+    ``split_indptr``) for the decode ASM kernel, built by whoever built this
+    call's ``qo_indptr``; None leaves the split pick to aiter.
+
     Otherwise (bf16): the existing Triton / reference path. When ``kv_scales``
     is provided, ``unified_kv`` must be fp8 (e4m3fnuz) and is dequantized
     in-kernel using 1xGROUP_SIZE (default 64) block scales (legacy 1buff,
@@ -1281,6 +1205,7 @@ def sparse_attn_v4_paged_decode(
                 q_packed_in,
                 q_rope_in,
             )
+        num_kv_splits, split_indptr = split_plan or (None, None)
         return _sparse_attn_v4_paged_decode_asm(
             unified_kv,
             kv_indices,
@@ -1292,12 +1217,8 @@ def sparse_attn_v4_paged_decode(
             q_rope_in,
             qo_indptr=qo_indptr,
             kv_last_page_lens=kv_last_page_lens,
-            **_v4_aiter_decode_split_kwargs(
-                q_packed_in,
-                query_group=query_group,
-                kv_kind=kv_kind,
-                split_plan=split_plan,
-            ),
+            num_kv_splits=num_kv_splits,
+            split_indptr=split_indptr,
         )
     gfx = get_gfx()
     if gfx == "gfx1250" or gfx.startswith("gfx94"):
diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py
index f969879866..053621464d 100644
--- a/atom/models/deepseek_v4.py
+++ b/atom/models/deepseek_v4.py
@@ -3313,12 +3313,15 @@ def _sparse_attention(
             if ratio == 0:
                 kv_indices = attn_md.kv_indices_swa
                 kv_indptr = attn_md.kv_indptr_swa
+                split_plan = attn_md.split_plan_swa
             elif ratio == 4:
                 kv_indices = attn_md.kv_indices_csa
                 kv_indptr = attn_md.kv_indptr_csa
+                split_plan = attn_md.split_plan_csa
             else:  # ratio == 128
                 kv_indices = attn_md.kv_indices_hca
                 kv_indptr = attn_md.kv_indptr_hca
+                split_plan = attn_md.split_plan_hca
             # Dispatch on kv-cache layout inside the wrapper: fp8 2buff
             # (unified_kv_rope set) → aiter ASM with pre-packed fp8 Q + the
             # 2buff fp8/bf16 pools read with no requant. The optional H=128
@@ -3336,18 +3339,7 @@ def _sparse_attention(
                 qo_indptr=attn_md.qo_indptr,
                 empty_kv_indptr=attn_md.empty_kv_indptr,
                 prefix=f"{self.layer_name}.sparse_attn_decode",
-                query_group=(
-                    7 if attn_md.min_seqlen_q == attn_md.max_seqlen_q == 7 else 1
-                ),
-                kv_kind="csa" if ratio == 4 else "hca" if ratio == 128 else "swa",
-                # Bridged metadata (e.g. sglang) may not carry split plans.
-                split_plan=getattr(
-                    attn_md,
-                    {0: "split_plan_swa", 4: "split_plan_csa"}.get(
-                        ratio, "split_plan_hca"
-                    ),
-                    None,
-                ),
+                split_plan=split_plan,
             )  # [S, H, head_dim]
         else:
             # Two-source paged prefill: prefix from `unified_kv` (per-ratio
diff --git a/atom/models/deepseek_v4_dspark.py b/atom/models/deepseek_v4_dspark.py
index a172a018cf..257cfcbae9 100644
--- a/atom/models/deepseek_v4_dspark.py
+++ b/atom/models/deepseek_v4_dspark.py
@@ -844,7 +844,7 @@ def dspark_attention(
             # accessor (lazy); `bufs.views` raises if stage 0 did not fill first.
             bufs = self.index_buffers(T, W, x.device)
             if self.stage_id == 0:
-                bufs.build(a.swa_window, slots, positions)
+                bufs.build(a.swa_window, slots, positions, a.n_local_heads)
             kv_indices, kv_indptr, draft_rows = bufs.views(B)
             batch_ids = bufs.batch_ids[: B * T]
 
@@ -895,6 +895,7 @@ def dspark_attention(
                 q_rope_in=qkn.q_rope,
                 qo_indptr=bufs.qo_indptr[: B * T + 1],
                 prefix=f"{a.layer_name}.dspark_attn_fp8",
+                split_plan=bufs.split_plan,
             )  # [B*T, n_heads, head_dim]
             out = out.view(B, T, a.n_local_heads, a.head_dim)
         else:
diff --git a/tests/test_v4_prefill_asm_decode_dispatch.py b/tests/test_v4_prefill_asm_decode_dispatch.py
index 6bd2c49d82..bd8fb7c218 100644
--- a/tests/test_v4_prefill_asm_decode_dispatch.py
+++ b/tests/test_v4_prefill_asm_decode_dispatch.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: MIT
-"""Dispatch coverage for V4 FP8 ASM decode paths and split tuning."""
+"""Dispatch coverage for V4 FP8 ASM decode paths and their split plans."""
 
 from __future__ import annotations
 
@@ -115,168 +115,120 @@ def test_ineligible_decode_keeps_dedicated_asm(monkeypatch, kwargs):
     assert _dispatch(monkeypatch, **kwargs) == "decode"
 
 
-@pytest.mark.parametrize(
-    "requests,expected",
-    [(1, 5), (2, 4), (3, 3), (4, 2), (5, 1), (32, 1), (64, 1), (256, 1)],
-)
-def test_aiter_csa_uses_graphsafe_batch_split_table(monkeypatch, requests, expected):
-    monkeypatch.setattr(paged_decode, "_device_arch", lambda _index: "gfx950")
-    q_packed = SimpleNamespace(
-        shape=(requests * 7, 128, 512), device=SimpleNamespace(index=0)
-    )
-
-    assert (
-        paged_decode._v4_aiter_fp8_decode_splits(
-            q_packed,
-            query_group=7,
-            kv_kind="csa",
-        )
-        == expected
+def _decode_rows(n: int, heads: int = 128) -> dict:
+    return dict(
+        unified_kv=torch.empty((4, 512)),
+        kv_indices=torch.empty(0, dtype=torch.int32),
+        kv_indptr=torch.zeros(n + 1, dtype=torch.int32),
+        attn_sink=torch.empty(heads),
+        softmax_scale=512**-0.5,
+        unified_kv_rope=torch.empty((4, 64)),
+        q_packed_in=torch.empty((n, heads, 512)),
+        q_rope_in=torch.empty((n, heads, 64)),
     )
 
 
 @pytest.mark.parametrize(
-    "tokens,heads,query_group,kv_kind,arch",
+    "split_plan,expected",
     [
-        (6 * 7, 128, 7, "hca", "gfx950"),
-        (6 * 7, 64, 7, "csa", "gfx950"),
-        (6 * 4, 128, 4, "csa", "gfx950"),
-        (6 * 7, 128, 7, "csa", "gfx1250"),
+        ((4, torch.arange(0, 4 * 22, 4, dtype=torch.int32)), 4),
+        (None, None),
     ],
 )
-def test_aiter_split_table_rejects_unqualified_shapes(
-    monkeypatch, tokens, heads, query_group, kv_kind, arch
-):
-    monkeypatch.setattr(paged_decode, "_device_arch", lambda _index: arch)
-    q_packed = SimpleNamespace(
-        shape=(tokens, heads, 512), device=SimpleNamespace(index=0)
-    )
-
-    assert (
-        paged_decode._v4_aiter_fp8_decode_splits(
-            q_packed,
-            query_group=query_group,
-            kv_kind=kv_kind,
-        )
-        is None
-    )
-
-
-def _decode_split_kwargs(monkeypatch, *, planner, n=21, **kwargs):
+def test_split_plan_reaches_decode_asm(monkeypatch, split_plan, expected):
     captured = {}
-
     monkeypatch.setattr(paged_decode.envs, "ATOM_USE_V4_PREFILL_ASM_FOR_DECODE", False)
     monkeypatch.setattr(
         paged_decode,
-        "aiter_v4_nm_split_plan_fn",
-        lambda: (lambda *a, **k: None) if planner else None,
-    )
-    monkeypatch.setattr(
-        paged_decode,
-        "_v4_aiter_fp8_decode_splits",
-        lambda q_packed, *, query_group, kv_kind: 3,
+        "_sparse_attn_v4_paged_decode_asm",
+        lambda *args, **kwargs: captured.update(kwargs) or "decode",
     )
-
-    def fake_decode(*args, **kw):
-        captured.update(kw)
-        return "decode"
-
-    monkeypatch.setattr(paged_decode, "_sparse_attn_v4_paged_decode_asm", fake_decode)
-    heads = 128
+    n = 21
     result = paged_decode.sparse_attn_v4_paged_decode(
         q=None,
-        unified_kv=torch.empty((4, 512)),
-        kv_indices=torch.empty(0, dtype=torch.int32),
-        kv_indptr=torch.zeros(n + 1, dtype=torch.int32),
-        attn_sink=torch.empty(heads),
-        softmax_scale=512**-0.5,
-        unified_kv_rope=torch.empty((4, 64)),
-        q_packed_in=torch.empty((n, heads, 512)),
-        q_rope_in=torch.empty((n, heads, 64)),
         qo_indptr=torch.arange(n + 1, dtype=torch.int32),
-        query_group=7,
-        **kwargs,
+        split_plan=split_plan,
+        **_decode_rows(n),
     )
-    assert result == "decode"
-    return {k: v for k, v in captured.items() if k in _SPLIT_KEYS}
-
-
-_SPLIT_KEYS = ("num_kv_splits", "split_indptr")
-
-
-def test_legacy_aiter_receives_csa_split_table(monkeypatch):
-    got = _decode_split_kwargs(monkeypatch, planner=False, kv_kind="csa")
-    assert got == {"num_kv_splits": 3}
 
+    assert result == "decode"
+    assert captured["num_kv_splits"] == expected
+    assert captured["split_indptr"] is (split_plan[1] if split_plan else None)
 
-@pytest.mark.parametrize("planner", [True, False])
-def test_metadata_split_plan_is_forwarded(monkeypatch, planner):
-    n = 21
-    plan = (4, torch.arange(0, 4 * (n + 1), 4, dtype=torch.int32))
-    got = _decode_split_kwargs(
-        monkeypatch, planner=planner, n=n, kv_kind="hca", split_plan=plan
-    )
-    assert got["num_kv_splits"] == 4
-    assert got["split_indptr"] is plan[1]
 
+def test_decode_asm_trims_split_indptr_to_real_rows(monkeypatch):
+    # An eager forward can run fewer rows than the padded grid its plan was
+    # built for; the uniform plan's prefix must go along with qo_indptr's.
+    import aiter.mla
 
-def test_plan_covering_more_rows_is_forwarded(monkeypatch):
-    # Eager forwards can run fewer rows than the padded grid the plan was built
-    # for; the ASM wrapper trims the uniform plan like qo_indptr.
-    plan = (2, torch.arange(0, 2 * 29, 2, dtype=torch.int32))
-    got = _decode_split_kwargs(
-        monkeypatch, planner=True, n=21, kv_kind="csa", split_plan=plan
-    )
-    assert got["num_kv_splits"] == 2
+    captured = {}
+    monkeypatch.setattr(
+        aiter.mla,
+        "mla_decode_fwd_v4_nm",
+        lambda *args, **kwargs: captured.update(kwargs),
+    )
+    n, padded = 3, 8
+    rows = _decode_rows(n)
+    paged_decode._sparse_attn_v4_paged_decode_asm(
+        rows["unified_kv"],
+        rows["kv_indices"],
+        torch.zeros(padded + 1, dtype=torch.int32),
+        rows["attn_sink"],
+        rows["softmax_scale"],
+        rows["unified_kv_rope"],
+        rows["q_packed_in"],
+        rows["q_rope_in"],
+        qo_indptr=torch.arange(padded + 1, dtype=torch.int32),
+        num_kv_splits=2,
+        split_indptr=torch.arange(0, 2 * (padded + 1), 2, dtype=torch.int32),
+    )
+
+    assert captured["num_kv_splits"] == 2
+    assert captured["split_indptr"].tolist() == [0, 2, 4, 6]
+
+
+@pytest.mark.parametrize("has_planner", [True, False])
+def test_v4_decode_split_plan_follows_aiter(monkeypatch, has_planner):
+    import aiter.mla
 
+    calls = []
 
-def test_plan_for_fewer_rows_falls_back(monkeypatch):
-    plan = (2, torch.arange(0, 2 * 8, 2, dtype=torch.int32))
-    got = _decode_split_kwargs(
-        monkeypatch, planner=True, n=21, kv_kind="csa", split_plan=plan
-    )
-    assert got == {}
+    def planner(rows, heads, kv_len, *, split_indptr):
+        calls.append((rows, heads, kv_len))
+        return 2, split_indptr[: rows + 1]
 
+    if has_planner:
+        monkeypatch.setattr(
+            aiter.mla, "get_mla_v4_nm_split_plan", planner, raising=False
+        )
+    else:
+        monkeypatch.delattr(aiter.mla, "get_mla_v4_nm_split_plan", raising=False)
+    buf = torch.zeros(65, dtype=torch.int32)
+    plan = paged_decode.v4_decode_split_plan(28, 128, 1152, buf)
 
-def test_planner_aiter_without_plan_uses_auto(monkeypatch):
-    got = _decode_split_kwargs(monkeypatch, planner=True, kv_kind="swa")
-    assert got == {}
+    if has_planner:
+        assert calls == [(28, 128, 1152)]
+        assert plan[0] == 2 and plan[1].shape[0] == 29
+    else:
+        assert plan is None
 
 
-@pytest.mark.parametrize("kv_fp8,planner", [(True, True), (False, True), (True, False)])
-def test_builder_attaches_split_plans(monkeypatch, kv_fp8, planner):
+@pytest.mark.parametrize("kv_fp8", [True, False])
+def test_builder_decode_split_plan(monkeypatch, kv_fp8):
     from atom.model_ops.attentions import deepseek_v4_attn
 
     calls = []
-
-    def fake_plan(num_seqs, num_heads, kv_len, *, split_indptr):
-        calls.append((num_seqs, num_heads, kv_len))
-        return 2, split_indptr[: num_seqs + 1]
-
     monkeypatch.setattr(
-        paged_decode,
-        "aiter_v4_nm_split_plan_fn",
-        lambda: fake_plan if planner else None,
-    )
-    bufs = {
-        f"v4_split_indptr_{k}": torch.zeros(65, dtype=torch.int32)
-        for k in ("swa", "csa", "hca")
-    }
-    builder = SimpleNamespace(
-        _kv_fp8=kv_fp8,
-        _decode_local_heads=128,
-        window_size=128,
-        index_topk=1024,
-        max_committed_hca=8192,
-        model_runner=SimpleNamespace(forward_vars=bufs),
+        deepseek_v4_attn,
+        "v4_decode_split_plan",
+        lambda rows, heads, kv_len, buf: calls.append((rows, heads, kv_len)) or "plan",
     )
-    md = SimpleNamespace(split_plan_swa=None, split_plan_csa=None, split_plan_hca=None)
-    deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._attach_v4_decode_split_plans(
-        builder, md, 28, ""
+    builder = SimpleNamespace(_kv_fp8=kv_fp8, _local_heads=128)
+    plan = deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._decode_split_plan(
+        builder, 28, 1152, torch.zeros(65, dtype=torch.int32)
     )
 
-    if kv_fp8 and planner:
-        assert calls == [(28, 128, 128), (28, 128, 1152), (28, 128, 8320)]
-        assert md.split_plan_csa[0] == 2 and md.split_plan_csa[1].shape[0] == 29
+    if kv_fp8:
+        assert plan == "plan" and calls == [(28, 128, 1152)]
     else:
-        assert calls == [] and md.split_plan_csa is None
+        assert plan is None and calls == []