perf(dsv4): adapt tuned AITER ASM FP8 decode

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,
 )
@@ -296,6 +298,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."""
+    split_plan_swa: tuple[int, torch.Tensor] | None = None
+    """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
+    """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
@@ -710,6 +719,11 @@ 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 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()
+        )
 
         # Sparse-attn + per-fwd metadata buffers (CG-A: pre-allocate for fixed
         # GPU pointers, prerequisite for CUDAGraph capture). All H2D copies in
@@ -2607,6 +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
+        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-
@@ -3677,6 +3694,20 @@ def _attach_v4_per_fwd_meta(
             buf_prefix_ubatch=buf_prefix_ubatch,
         )
 
+    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`).
+
+        `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,
         attn_metadata,
@@ -3799,6 +3830,19 @@ def _attach_v4_paged_decode_meta(
             win=win,
             index_topk=index_topk,
         )
+        # 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
@@ -4622,6 +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)
+        # 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.
@@ -4649,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.
@@ -4821,6 +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)
+            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(

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",

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
 

atom/model_ops/v4_kernels/paged_decode.py

@@ -985,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,
@@ -997,6 +1017,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,
+    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).
@@ -1106,6 +1127,10 @@ def _sparse_attn_v4_paged_decode_asm(
         (N, H, V4_DIM_NOPE + V4_DIM_ROPE), dtype=torch.bfloat16, device=device
     )
 
+    # 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,
@@ -1119,6 +1144,7 @@ def _sparse_attn_v4_paged_decode_asm(
         sink=attn_sink,
         sm_scale=softmax_scale,
         num_kv_splits=num_kv_splits,
+        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.
@@ -1141,6 +1167,7 @@ def sparse_attn_v4_paged_decode(
     kv_last_page_lens: torch.Tensor | None = None,
     empty_kv_indptr: torch.Tensor | None = None,
     prefix: str = "",
+    split_plan: tuple[int, torch.Tensor] | None = None,
 ) -> torch.Tensor:
     """V4 decode sparse attention over a unified KV pool with paged indices.
 
@@ -1150,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,
@@ -1174,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,
@@ -1185,6 +1217,8 @@ def sparse_attn_v4_paged_decode(
             q_rope_in,
             qo_indptr=qo_indptr,
             kv_last_page_lens=kv_last_page_lens,
+            num_kv_splits=num_kv_splits,
+            split_indptr=split_indptr,
         )
     gfx = get_gfx()
     if gfx == "gfx1250" or gfx.startswith("gfx94"):

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,6 +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",
+                split_plan=split_plan,
             )  # [S, H, head_dim]
         else:
             # Two-source paged prefill: prefix from `unified_kv` (per-ratio

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: