From b9ba7bbaee0d15bd77ef8d193a9829067d77f099 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Fri, 25 Sep 2026 22:48:35 -0500 Subject: [PATCH 01/12] 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 Date: Fri, 25 Sep 2026 22:51:23 -0500 Subject: [PATCH 02/12] 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 Date: Sun, 27 Sep 2026 08:20:51 -0500 Subject: [PATCH 03/12] 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) Signed-off-by: yihonglie --- 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 Date: Mon, 28 Sep 2026 01:05:33 -0500 Subject: [PATCH 04/12] 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) Signed-off-by: yihonglie --- 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 Date: Mon, 28 Sep 2026 02:20:07 -0500 Subject: [PATCH 05/12] 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) Signed-off-by: yihonglie --- 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 == [] From ac8e09da674a8addeb3d37d1d62f561831bcf56c Mon Sep 17 00:00:00 2001 From: yihonglie Date: Tue, 29 Sep 2026 11:19:49 -0500 Subject: [PATCH 06/12] [DSV4] Run HCA fp8 decode through aiter's persistent V4-NM kernel HCA (compress_ratio 128) decode calls now use aiter.mla.mla_decode_fwd_v4_nm_ps: one launch that plans the KV split on the GPU from kv_indptr and merges the partials in-kernel, so no host-side split plan and no stage-2 merge. - atom/model_ops/v4_kernels/hca_persist.py: gates, per-device workspace (get_mla_v4_nm_ps_workspace) and the call. Taken only for fp8 KV, 128 local heads, gfx950, row-dense KV pools and ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= 32768; everything else stays on the decode ASM + split plan. An aiter without the kernel stays on the ASM path and logs once. - The V4 metadata builder allocates the workspace at init, before warmup, KV sizing and graph capture. All HCA decode calls run on the forward's compute stream, so one workspace per device is enough. - sparse_attn_v4_paged_decode takes the caller's compress_ratio (passed from DeepseekV4Attention's decode call site). - ATOM_V4_HCA_PERSIST (default 1) and ATOM_V4_HCA_PERSIST_MIN_ROWS (default 15), documented in docs/environment_variables.md. - Tests: mocked dispatch rules, and a gfx950 test against the ASM path and the torch reference, including padded rows, an offset kv_indptr and CUDA-graph replay with changing KV lengths. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: yihonglie --- atom/model_ops/attentions/deepseek_v4_attn.py | 22 ++ atom/model_ops/v4_kernels/hca_persist.py | 163 +++++++++++ atom/model_ops/v4_kernels/paged_decode.py | 26 ++ atom/models/deepseek_v4.py | 1 + atom/utils/envs.py | 9 + docs/environment_variables.md | 2 + tests/model_ops/test_v4_hca_persist_gpu.py | 269 ++++++++++++++++++ tests/test_v4_hca_persist_dispatch.py | 250 ++++++++++++++++ 8 files changed, 742 insertions(+) create mode 100644 atom/model_ops/v4_kernels/hca_persist.py create mode 100644 tests/model_ops/test_v4_hca_persist_gpu.py create mode 100644 tests/test_v4_hca_persist_dispatch.py diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index d3538db3d4..9e75e7d6d0 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -119,6 +119,7 @@ FP4_MQA_PARALLEL_UNIT_NUM, build_v4_paged_decode_indptr, fp4_indexer_enabled, + hca_persist, plan_context_lens, v4_decode_split_plan, write_v4_paged_decode_indices, @@ -724,6 +725,7 @@ def __init__(self, model_runner): model_runner.config.hf_config.num_attention_heads // get_tensor_model_parallel_world_size() ) + self._prepare_hca_persist() # Sparse-attn + per-fwd metadata buffers (CG-A: pre-allocate for fixed # GPU pointers, prerequisite for CUDAGraph capture). All H2D copies in @@ -3694,6 +3696,26 @@ def _attach_v4_per_fwd_meta( buf_prefix_ubatch=buf_prefix_ubatch, ) + def _prepare_hca_persist(self) -> None: + """Allocate the persistent HCA decode workspace here, at init: before + warmup, before KV sizing (so it is accounted for) and before any + CUDA-graph capture, which may be the first forward to reach it.""" + if not envs.ATOM_V4_HCA_PERSIST: + return + why = None + if not self._kv_fp8: + why = "kv cache is not fp8" + elif not self.hca_layers: + why = "no HCA (ratio 128) layers" + elif self._local_heads != hca_persist.HEADS: + why = f"{self._local_heads} local heads (needs {hca_persist.HEADS})" + elif get_gfx() != "gfx950": + why = f"arch {get_gfx()} (needs gfx950)" + if why is not None: + logger.info("V4 HCA persistent decode not used: %s", why) + return + hca_persist.prepare(self.device) + def _decode_split_plan( self, rows: int, kv_len: int, split_indptr: torch.Tensor ) -> tuple[int, torch.Tensor] | None: diff --git a/atom/model_ops/v4_kernels/hca_persist.py b/atom/model_ops/v4_kernels/hca_persist.py new file mode 100644 index 0000000000..71d98410b4 --- /dev/null +++ b/atom/model_ops/v4_kernels/hca_persist.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Persistent HCA fp8 decode (``ATOM_V4_HCA_PERSIST``, on by default). + +HCA (compress_ratio 128) decode calls run aiter's persistent V4-NM kernel +(``aiter.mla.mla_decode_fwd_v4_nm_ps``): one launch that plans the KV split +from ``kv_indptr`` on the GPU, runs the attention and merges the partials, so +no host-side split plan and no stage-2 merge. Same tensors and layouts as the +decode ASM path (``_sparse_attn_v4_paged_decode_asm``); rows with K = 0 (CUDA +graph padding) are left unwritten, as the ASM does. Limits: 128 local heads, +``N <= 32768`` rows, gfx950, row-dense KV pools. Everything else, and an aiter +without the kernel, stays on the ASM + split plan. + +Workspace: one per device, allocated by :func:`prepare` outside graph capture +(the V4 metadata builder calls it at init, before warmup and KV sizing). It +holds the split partials and the merge counters, which return to zero at the +end of every call, so calls sharing it must be ordered on the GPU. In ATOM +every HCA decode call runs on the forward's compute stream (layers are +sequential; TBO micro-batches share one compute stream; eager forwards and +graph replays do not overlap), so one workspace per device is enough. +""" + +from __future__ import annotations + +import logging + +import torch + +from atom.utils import envs + +logger = logging.getLogger("atom") + +try: + from aiter.mla import get_mla_v4_nm_ps_workspace, mla_decode_fwd_v4_nm_ps +except ImportError: # aiter without the persistent V4-NM kernel + get_mla_v4_nm_ps_workspace = mla_decode_fwd_v4_nm_ps = None + +HEADS = 128 +MAX_ROWS = 32768 # byte-offset wrap inherited from the ASM +HCA_RATIO = 128 +PARTITIONS = 128 +_DIM = 512 +_ROPE = 64 + +# Host-side call counts: eager calls count per call, a CUDA graph once at +# capture (replays do not run Python). +stats = {"persist": 0, "persist_captured": 0} +_logged: set = set() +_workspaces: dict[int, object] = {} + + +def _log_once(key, msg: str, *args) -> None: + if key not in _logged: + _logged.add(key) + logger.info(msg, *args) + + +def available() -> bool: + if mla_decode_fwd_v4_nm_ps is None: + _log_once( + "no_api", + "V4 HCA persistent decode unavailable: aiter lacks " + "mla_decode_fwd_v4_nm_ps; HCA decode stays on the ASM path", + ) + return False + return True + + +def _dev_index(device) -> int: + device = torch.device(device) + return device.index if device.index is not None else torch.cuda.current_device() + + +def prepare(device) -> None: + """Allocate this device's workspace (idempotent). Must run outside CUDA + graph capture; a no-op when aiter lacks the kernel.""" + idx = _dev_index(device) + if idx in _workspaces or not available(): + return + _workspaces[idx] = get_mla_v4_nm_ps_workspace( + torch.device("cuda", idx), num_partitions=PARTITIONS + ) + logger.info( + "V4 HCA persistent decode: device %d ready (P=%d, min rows %d)", + idx, + PARTITIONS, + envs.ATOM_V4_HCA_PERSIST_MIN_ROWS, + ) + + +def wanted(*, compress_ratio: int | None, heads: int, rows: int, gfx: str) -> bool: + """Shape/config gate; the layout gate is :func:`layout_ok`.""" + return ( + envs.ATOM_V4_HCA_PERSIST + and compress_ratio == HCA_RATIO + and heads == HEADS + and gfx == "gfx950" + and envs.ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= MAX_ROWS + and available() + ) + + +def layout_ok(unified_kv: torch.Tensor, unified_kv_rope: torch.Tensor) -> bool: + """The kernel addresses pool row r at base + r * row_bytes.""" + return ( + unified_kv.dim() == 2 + and unified_kv.shape[-1] == _DIM + and unified_kv.element_size() == 1 + and unified_kv.stride() == (_DIM, 1) + and unified_kv_rope.dim() == 2 + and unified_kv_rope.shape[-1] == _ROPE + and unified_kv_rope.dtype == torch.bfloat16 + and unified_kv_rope.stride() == (_ROPE, 1) + ) + + +def hca_persist_decode( + unified_kv: torch.Tensor, + kv_indices: torch.Tensor, + kv_indptr: torch.Tensor, + attn_sink: torch.Tensor, + unified_kv_rope: torch.Tensor, + q_packed_in: torch.Tensor, + q_rope_in: torch.Tensor, +) -> torch.Tensor: + """HCA decode of the ``N = q_packed_in.shape[0]`` rows the ASM path would + run (``kv_indptr[:N + 1]``; an eager forward may have N < T_pad).""" + n, h, _ = q_packed_in.shape + if h != HEADS or n > MAX_ROWS or kv_indptr.numel() < n + 1: + raise ValueError( + f"hca_persist: needs {HEADS} heads, <= {MAX_ROWS} rows and " + f">= N+1 kv_indptr entries; got heads={h}, rows={n}, " + f"kv_indptr={kv_indptr.numel()}" + ) + idx = _dev_index(q_packed_in.device) + if idx not in _workspaces: + prepare(q_packed_in.device) # aiter refuses during capture + out = torch.empty((n, HEADS, _DIM), dtype=torch.bfloat16, device=q_packed_in.device) + if n == 0: + return out + mla_decode_fwd_v4_nm_ps( + q_packed_in.contiguous(), + q_rope_in.contiguous(), + unified_kv, + unified_kv_rope, + kv_indptr[: n + 1], + kv_indices.contiguous(), + attn_sink.contiguous(), + _workspaces[idx], + out=out, + ) + capturing = torch.cuda.is_current_stream_capturing() + stats["persist"] += 1 + stats["persist_captured"] += int(capturing) + _log_once( + ("hit", capturing), + "V4 HCA persistent decode: first %s call on device %d (rows=%d)", + "CUDA-graph captured" if capturing else "eager", + idx, + n, + ) + return out diff --git a/atom/model_ops/v4_kernels/paged_decode.py b/atom/model_ops/v4_kernels/paged_decode.py index 73ce7dd434..8da10e30cb 100644 --- a/atom/model_ops/v4_kernels/paged_decode.py +++ b/atom/model_ops/v4_kernels/paged_decode.py @@ -59,6 +59,7 @@ from aiter.ops.triton.utils.device_info import get_num_sms from atom.model_ops.sparse_attn_v4 import _sparse_attn_ragged_torch +from atom.model_ops.v4_kernels import hca_persist from atom.model_ops.v4_kernels.paged_decode_gluon import ( _paged_decode_fused_gluon_kernel, ) @@ -1168,6 +1169,7 @@ def sparse_attn_v4_paged_decode( empty_kv_indptr: torch.Tensor | None = None, prefix: str = "", split_plan: tuple[int, torch.Tensor] | None = None, + compress_ratio: int | None = None, ) -> torch.Tensor: """V4 decode sparse attention over a unified KV pool with paged indices. @@ -1181,6 +1183,11 @@ def sparse_attn_v4_paged_decode( ``split_indptr``) for the decode ASM kernel, built by whoever built this call's ``qo_indptr``; None leaves the split pick to aiter. + ``compress_ratio`` is the calling layer's ratio (0 SWA, 4 CSA, 128 HCA). + With ``ATOM_V4_HCA_PERSIST`` (default on), fp8 HCA calls with 128 heads on + gfx950 and ``ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= 32768`` run aiter's + persistent kernel (``hca_persist``) instead; ``split_plan`` is then unused. + 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, @@ -1205,6 +1212,25 @@ def sparse_attn_v4_paged_decode( q_packed_in, q_rope_in, ) + if ( + q_packed_in is not None + and hca_persist.wanted( + compress_ratio=compress_ratio, + heads=q_packed_in.shape[1], + rows=q_packed_in.shape[0], + gfx=get_gfx(), + ) + and hca_persist.layout_ok(unified_kv, unified_kv_rope) + ): + return hca_persist.hca_persist_decode( + unified_kv, + kv_indices, + kv_indptr, + attn_sink, + unified_kv_rope, + 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, diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 053621464d..85c3f75487 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -3340,6 +3340,7 @@ def _sparse_attention( empty_kv_indptr=attn_md.empty_kv_indptr, prefix=f"{self.layer_name}.sparse_attn_decode", split_plan=split_plan, + compress_ratio=ratio, ) # [S, H, head_dim] else: # Two-source paged prefill: prefix from `unified_kv` (per-ratio diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 68d98aec82..e2c9fafe4e 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -525,6 +525,15 @@ def _positive_float_env(name: str, default: str) -> float: "ATOM_USE_V4_PREFILL_ASM_FOR_DECODE": lambda: ( os.getenv("ATOM_USE_V4_PREFILL_ASM_FOR_DECODE", "0") == "1" ), + # DeepSeek-V4 HCA (compress_ratio 128) fp8 decode through aiter's + # persistent V4-NM kernel (one launch, in-kernel split + merge) instead of + # the decode ASM + split plan. gfx950, 128 local heads only; quietly off + # when aiter lacks mla_decode_fwd_v4_nm_ps. + "ATOM_V4_HCA_PERSIST": lambda: os.getenv("ATOM_V4_HCA_PERSIST", "1") == "1", + # Calls with fewer q rows stay on the ASM path, which is faster there. + "ATOM_V4_HCA_PERSIST_MIN_ROWS": lambda: int( + os.getenv("ATOM_V4_HCA_PERSIST_MIN_ROWS", "15") + ), # Route the paged decode to aiter's FlyDSL kernel (#4332) instead of gluon. "ATOM_PA_FLYDSL": lambda: (os.getenv("ATOM_PA_FLYDSL", "0") == "1"), # FlyDSL GPU work planner, built once per forward in the metadata diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 62963b04da..9eb4458a94 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -108,6 +108,8 @@ make duplicate prefill useful. Pure-attention models do not use checkpoint waits | **ATOM_USE_FLYDSL_FP8_PREFILL_ATTN** | bool | 0 (false) | Enable FlyDSL FP8 MLA prefill on gfx950, including fused K concatenation with QKV quantization and direct FP8 output from FlyDSL gather where supported. AITER capability checks are cached per layer. Missing FP8 attention support or unsupported device/model configurations select BF16 attention before quantization and emit a warning once per process. FP16 activations and synthetic RoPE-padding configurations also use the normal attention path. Cached K/V reuse descales `max(new_token_descale, 1e-6) * 2`; larger outliers saturate. Unsupported FP8 gather configurations use BF16 gather followed by fused K/V dynamic quantization when FP8 attention is supported. Kernel execution errors propagate without retry. Added 2026-09-10. | | **ATOM_PA_FLYDSL** | bool | 0 (false) | Route the MHA paged decode to aiter's FlyDSL kernel (aiter #4332) instead of the gluon one, where FlyDSL's domain covers the call. Off by default: the kernel is not fully tested on every shape ATOM ships. The env is the first of two gates; the second mirrors aiter's own validation so an unsupported shape falls back to gluon rather than raising inside aiter, and the `and` short-circuits so nothing is evaluated when this is off. Global, not per call site: the vLLM bridge is rerouted too. Needs aiter `94dca7bc6` or later. | | **ATOM_PA_FLYDSL_PLAN** | bool | 1 (true) | Use aiter #5546's GPU work planner on the dense decode, built once per forward in the metadata builder. Requires `ATOM_PA_FLYDSL=1` — with FlyDSL off no plan is built at all. Sizes each request's partition count from its real context length instead of splitting the batch uniformly, which is worth p90 interactivity `+20.5%` at concurrency 20 on the MiniMax-M3 agentic trace and nothing at concurrency 1, where there is nothing to rebalance. Kept off the two MiniMax-M3 sparse call sites, whose contexts are a fixed topk window and where it measures a net loss. Plans are only minted during cudagraph capture — aiter's planner takes the batch as a `tl.constexpr`, so a new value is a cold kernel specialization and a plan no captured graph can ever release — which leaves the planner inert under `--enforce-eager`. | +| **ATOM_V4_HCA_PERSIST** | bool | 1 (true) | DeepSeek-V4 HCA (`compress_ratio` 128) fp8 decode through aiter's persistent V4-NM kernel (`aiter.mla.mla_decode_fwd_v4_nm_ps`): one launch that splits the KV on the GPU and merges in-kernel, instead of the decode ASM plus a host-built split plan. Taken only with fp8 KV, 128 local heads, gfx950, row-dense KV pools and `ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= 32768`; everything else stays on the ASM path. An aiter without the kernel keeps the ASM path (logged once). One workspace per device is allocated by the V4 metadata builder at init. Set `0` to force the ASM path. | +| **ATOM_V4_HCA_PERSIST_MIN_ROWS** | int | 15 | Smallest decode call (q rows) that uses `ATOM_V4_HCA_PERSIST`; smaller calls stay on the ASM path. | ### GLM-5.3 diff --git a/tests/model_ops/test_v4_hca_persist_gpu.py b/tests/model_ops/test_v4_hca_persist_gpu.py new file mode 100644 index 0000000000..61309e97af --- /dev/null +++ b/tests/model_ops/test_v4_hca_persist_gpu.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: MIT +"""GPU check of the persistent HCA decode kernel through +``sparse_attn_v4_paged_decode``: persistent path vs the aiter decode ASM path +vs the torch reference (``mla_decode_fwd_v4_nm_ref``), on HCA-shaped inputs. + +Needs gfx950 and an aiter with ``mla_decode_fwd_v4_nm_ps``; skipped otherwise. +""" + +from __future__ import annotations + +import random + +import pytest +import torch + +pytest.importorskip("aiter") +if not torch.cuda.is_available(): + pytest.skip("needs a GPU", allow_module_level=True) + +from aiter.jit.utils.chip_info import get_gfx # noqa: E402 + +from atom.model_ops.v4_kernels import hca_persist, paged_decode # noqa: E402 +from atom.model_ops.v4_kernels.v4_quant import ( # noqa: E402 + mla_decode_fwd_v4_nm_ref, + quantize_bf16_to_v4_2buff, +) + +if get_gfx() != "gfx950": + pytest.skip("persistent HCA kernel is gfx950-only", allow_module_level=True) +if hca_persist.mla_decode_fwd_v4_nm_ps is None: + pytest.skip("aiter lacks mla_decode_fwd_v4_nm_ps", allow_module_level=True) + +DEV = "cuda" +H = 128 +MAX_ROWS_BUF = 1024 + + +def _hca_k(pos: int) -> int: + """HCA kv length at absolute position ``pos``: 128-token window + one + compressed entry per 128 tokens.""" + return min(pos + 1, 128) + (pos + 1) // 128 + + +def _q7_lens(rows, seed): + """kv lens of `rows` q-len-7 verify rows of requests with random contexts.""" + rng = random.Random(seed) + lens = [] + while len(lens) < rows: + c = int(2 ** rng.uniform(13, 17)) # 8k .. 128k context + lens += [_hca_k(c + i) for i in range(7)] + return lens[:rows] + + +def _rnd(g, *shape): + """bf16 randn with a per-64-group magnitude 2^U(-2, 2), so every tile + scale differs.""" + x = torch.randn(*shape, 512, device=DEV, generator=g) + mag = torch.exp2(torch.rand(*shape, 8, device=DEV, generator=g) * 4 - 2) + return (x * mag.repeat_interleave(64, -1)).to(torch.bfloat16) + + +def _inputs(kv_lens, seed, pool_extra=257): + g = torch.Generator(device=DEV).manual_seed(seed) + n, tot = len(kv_lens), int(sum(kv_lens)) + pool = tot + pool_extra + qp, qr = quantize_bf16_to_v4_2buff(_rnd(g, n, H)) + kp, kr = quantize_bf16_to_v4_2buff(_rnd(g, pool)) + indptr = torch.zeros(n + 1, dtype=torch.int32, device=DEV) + indptr[1:] = torch.tensor(kv_lens, device=DEV).cumsum(0) + idx = torch.randint(0, pool, (tot,), device=DEV, generator=g) + return dict( + q_packed=qp, + q_rope=qr, + kv_packed=kp, + kv_rope=kr, + kv_indptr=indptr, + kv_page_indices=idx.to(torch.int32), + sink=(torch.randn(H, device=DEV, generator=g) * 2).float(), + ) + + +def _reference(inp, n): + out = torch.empty((n, H, 512), dtype=torch.bfloat16, device=DEV) + mla_decode_fwd_v4_nm_ref( + inp["q_packed"][:n], + inp["q_rope"][:n], + inp["kv_packed"].view(-1, 1, 1, 512), + inp["kv_rope"].view(-1, 1, 1, 64), + out, + torch.arange(n + 1, dtype=torch.int32, device=DEV), + inp["kv_indptr"][: n + 1], + inp["kv_page_indices"], + torch.ones(n, dtype=torch.int32, device=DEV), + 1, + sink=inp["sink"], + ) + return out + + +def _close(a, b): + """Per (row, head): cosine >= 0.9999 and relative L2 <= 1e-2.""" + a, b = a.double().reshape(-1, 512), b.double().reshape(-1, 512) + cos = torch.nn.functional.cosine_similarity(a, b, dim=1).min().item() + rel = ((a - b).norm(dim=1) / b.norm(dim=1).clamp_min(1e-30)).max().item() + assert cos >= 0.9999 and rel <= 1e-2, (cos, rel) + + +@pytest.fixture +def mode(monkeypatch): + def set_mode(persist, min_rows=15): + monkeypatch.setenv("ATOM_V4_HCA_PERSIST", "1" if persist else "0") + monkeypatch.setenv("ATOM_V4_HCA_PERSIST_MIN_ROWS", str(min_rows)) + + return set_mode + + +def _run(inp, *, n=None, t_pad=None): + """One ATOM decode call: q has `n` rows (default all), the CSR / qo_indptr + have `t_pad` + 1 entries (default n), like the metadata builder stages them.""" + n = inp["q_packed"].shape[0] if n is None else n + t_pad = n if t_pad is None else t_pad + kv_indptr = inp["kv_indptr"][: t_pad + 1] + kv_len = int((kv_indptr[1:] - kv_indptr[:-1]).max()) + plan = paged_decode.v4_decode_split_plan( + t_pad, H, kv_len, torch.zeros(MAX_ROWS_BUF + 1, dtype=torch.int32, device=DEV) + ) + return paged_decode.sparse_attn_v4_paged_decode( + None, + inp["kv_packed"], + inp["kv_page_indices"], + kv_indptr, + inp["sink"], + 512**-0.5, + unified_kv_rope=inp["kv_rope"], + q_packed_in=inp["q_packed"][:n], + q_rope_in=inp["q_rope"][:n], + qo_indptr=torch.arange(t_pad + 1, dtype=torch.int32, device=DEV), + split_plan=plan, + compress_ratio=128, + ) + + +def _both(mode, inp, **kw): + mode(False) + asm = _run(inp, **kw) + mode(True, min_rows=1) + before = hca_persist.stats["persist"] + per = _run(inp, **kw) + assert hca_persist.stats["persist"] == before + 1, "persistent path not taken" + torch.cuda.synchronize() + return asm, per + + +@pytest.mark.parametrize("rows", [7, 15, 21, 112, 448, 896]) +def test_rows_vs_asm_and_reference(mode, rows): + inp = _inputs(_q7_lens(rows, rows), seed=rows) + asm, per = _both(mode, inp) + ref = _reference(inp, rows) + _close(asm, ref) + _close(per, ref) + _close(per, asm) + + +@pytest.mark.parametrize("rows", [7, 14, 15, 112]) +def test_default_min_rows_dispatch(mode, rows): + inp = _inputs(_q7_lens(rows, 1), seed=1) + mode(True) # default MIN_ROWS 15 + before = hca_persist.stats["persist"] + _run(inp) + assert hca_persist.stats["persist"] - before == (1 if rows >= 15 else 0) + + +def test_graph_padded_rows(mode): + """Captured-grid shape: T_pad rows, the tail has K = 0 (indptr repeats).""" + real = _q7_lens(35, 3) + inp = _inputs(real + [0] * 13, seed=3) + asm, per = _both(mode, inp) + r = len(real) + _close(per[:r], _reference(inp, r)) + _close(per[:r], asm[:r]) + + +def test_eager_fewer_rows_than_staged(mode): + """Eager forward: q has N real rows, the CSR was staged for T_pad > N.""" + inp = _inputs(_q7_lens(21, 4) + [0] * 11, seed=4) + asm, per = _both(mode, inp, n=21, t_pad=32) + assert per.shape[0] == 21 and asm.shape[0] == 21 + _close(per, _reference(inp, 21)) + _close(per, asm) + + +def test_offset_kv_indptr(mode): + inp = _inputs(_q7_lens(28, 5), seed=5) + ref = _reference(inp, 28) + off = 777 + shifted = dict(inp) + shifted["kv_indptr"] = inp["kv_indptr"] + off + junk = torch.randint( + 0, inp["kv_packed"].shape[0], (off,), dtype=torch.int32, device=DEV + ) + shifted["kv_page_indices"] = torch.cat([junk, inp["kv_page_indices"]]) + asm, per = _both(mode, shifted) + _close(per, ref) + _close(asm, ref) + + +def test_cuda_graph_capture_replay_changing_csr(mode): + """Capture once at T_pad rows, replay with new kv_indptr / indices / q.""" + t_pad, cap_k = 112, 1400 + mode(True) + hca_persist.prepare(DEV) # the builder does this at init, before capture + first = _inputs(_q7_lens(100, 6) + [0] * 12, seed=6, pool_extra=t_pad * cap_k) + pool_rows = first["kv_packed"].shape[0] + s = dict( + q_packed=first["q_packed"].clone(), + q_rope=first["q_rope"].clone(), + kv_packed=first["kv_packed"], + kv_rope=first["kv_rope"], + kv_indptr=torch.zeros(t_pad + 1, dtype=torch.int32, device=DEV), + kv_page_indices=torch.zeros(t_pad * cap_k, dtype=torch.int32, device=DEV), + sink=first["sink"], + ) + + def stage(lens, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + ip = torch.zeros(t_pad + 1, dtype=torch.int32, device=DEV) + ip[1:] = torch.tensor(lens, device=DEV).cumsum(0) + idx = torch.randint(0, pool_rows, (int(ip[-1]),), generator=g, device=DEV) + s["kv_indptr"].copy_(ip) + s["kv_page_indices"][: idx.numel()].copy_(idx.to(torch.int32)) + other = _inputs([1] * t_pad, seed=seed) + s["q_packed"].copy_(other["q_packed"]) + s["q_rope"].copy_(other["q_rope"]) + + stage(_q7_lens(100, 6) + [0] * 12, 60) + before = hca_persist.stats["persist"] + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = paged_decode.sparse_attn_v4_paged_decode( + None, + s["kv_packed"], + s["kv_page_indices"], + s["kv_indptr"], + s["sink"], + 512**-0.5, + unified_kv_rope=s["kv_rope"], + q_packed_in=s["q_packed"], + q_rope_in=s["q_rope"], + qo_indptr=torch.arange(t_pad + 1, dtype=torch.int32, device=DEV), + compress_ratio=128, + ) + assert hca_persist.stats["persist"] == before + 1 + for step, lens in enumerate( + [ + _q7_lens(100, 6) + [0] * 12, + _q7_lens(112, 7), + _q7_lens(42, 8) + [0] * 70, + [random.Random(9).randint(1, cap_k) for _ in range(t_pad)], + ] + ): + stage(lens, 100 + step) + graph.replay() + torch.cuda.synchronize() + r = sum(1 for k in lens if k > 0) + _close(out[:r], _reference(s, r)) + # the merge counters are back to zero after every completed call + ws = hca_persist._workspaces[torch.cuda.current_device()] + assert int(ws.cnt.abs().sum()) == 0 diff --git a/tests/test_v4_hca_persist_dispatch.py b/tests/test_v4_hca_persist_dispatch.py new file mode 100644 index 0000000000..6131408de8 --- /dev/null +++ b/tests/test_v4_hca_persist_dispatch.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: MIT +"""Dispatch rules of the persistent HCA decode kernel (``ATOM_V4_HCA_PERSIST``). +Mocked: no GPU and no aiter kernel needed.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("triton", reason="paged_decode defines Triton kernels") +pytest.importorskip("aiter", reason="paged_decode imports the AITER runtime") + +from atom.model_ops.v4_kernels import hca_persist, paged_decode + + +@pytest.fixture +def routes(monkeypatch): + """Replace both backends by recorders; return the call log.""" + calls = [] + monkeypatch.setattr(paged_decode.envs, "ATOM_USE_V4_PREFILL_ASM_FOR_DECODE", False) + monkeypatch.setattr( + paged_decode, + "_sparse_attn_v4_paged_decode_asm", + lambda *a, **k: calls.append(("asm", a, k)) or "asm", + ) + monkeypatch.setattr( + hca_persist, + "hca_persist_decode", + lambda *a, **k: calls.append(("persist", a, k)) or "persist", + ) + monkeypatch.setattr(hca_persist, "mla_decode_fwd_v4_nm_ps", object()) + return calls + + +def _env(monkeypatch, enabled=True, min_rows=15, gfx="gfx950"): + monkeypatch.setenv("ATOM_V4_HCA_PERSIST", "1" if enabled else "0") + monkeypatch.setenv("ATOM_V4_HCA_PERSIST_MIN_ROWS", str(min_rows)) + monkeypatch.setattr(paged_decode, "get_gfx", lambda: gfx) + + +def _call(n, *, heads=128, ratio=128, kv_rows=8, kv=None, kv_rope=None, plan=None): + kv = torch.empty((kv_rows, 512), dtype=torch.uint8) if kv is None else kv + kv_rope = ( + torch.empty((kv_rows, 64), dtype=torch.bfloat16) if kv_rope is None else kv_rope + ) + return paged_decode.sparse_attn_v4_paged_decode( + q=None, + unified_kv=kv, + kv_indices=torch.zeros(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=kv_rope, + q_packed_in=torch.empty((n, heads, 512), dtype=torch.uint8), + q_rope_in=torch.empty((n, heads, 64), dtype=torch.bfloat16), + qo_indptr=torch.arange(n + 1, dtype=torch.int32), + split_plan=plan, + compress_ratio=ratio, + ) + + +@pytest.mark.parametrize("n", [15, 21, 112, 448, 896, 32768]) +def test_hca_rows_at_or_above_min_go_persistent(monkeypatch, routes, n): + _env(monkeypatch) + assert _call(n) == "persist" + + +@pytest.mark.parametrize("n", [1, 7, 14]) +def test_small_batches_stay_on_asm(monkeypatch, routes, n): + _env(monkeypatch) + plan = (2, torch.arange(0, 2 * (n + 1), 2, dtype=torch.int32)) + assert _call(n, plan=plan) == "asm" + # the ASM path keeps its split plan unchanged + assert routes[-1][2]["num_kv_splits"] == 2 + assert routes[-1][2]["split_indptr"] is plan[1] + + +def test_min_rows_is_configurable(monkeypatch, routes): + _env(monkeypatch, min_rows=8) + assert _call(7) == "asm" + assert _call(8) == "persist" + + +def test_rows_above_limit_fall_back(monkeypatch, routes): + _env(monkeypatch) + assert _call(32769) == "asm" + + +@pytest.mark.parametrize( + "kw", + [ + {"ratio": 0}, + {"ratio": 4}, + {"ratio": None}, + {"heads": 64}, + {"heads": 16}, + ], +) +def test_non_hca_or_other_heads_stay_on_asm(monkeypatch, routes, kw): + _env(monkeypatch) + assert _call(112, **kw) == "asm" + + +@pytest.mark.parametrize("gfx", ["gfx942", "gfx1250"]) +def test_other_arch_stays_on_asm(monkeypatch, routes, gfx): + _env(monkeypatch, gfx=gfx) + # gfx1250 H=128 could take the prefill-ASM route; it is off in `routes` + assert _call(112) == "asm" + + +def test_switch_on_by_default(monkeypatch, routes): + monkeypatch.delenv("ATOM_V4_HCA_PERSIST", raising=False) + monkeypatch.setattr(paged_decode, "get_gfx", lambda: "gfx950") + assert _call(896) == "persist" + + +def test_switch_off(monkeypatch, routes): + _env(monkeypatch, enabled=False) + assert _call(896) == "asm" + + +def test_aiter_without_kernel_stays_on_asm(monkeypatch, routes): + _env(monkeypatch) + monkeypatch.setattr(hca_persist, "mla_decode_fwd_v4_nm_ps", None) + assert _call(896) == "asm" + monkeypatch.setattr(hca_persist, "get_mla_v4_nm_ps_workspace", None) + monkeypatch.setattr(hca_persist, "_workspaces", {}) + hca_persist.prepare("cuda:0") # quietly nothing to prepare + assert hca_persist._workspaces == {} + + +def test_strided_pool_falls_back(monkeypatch, routes): + _env(monkeypatch) + wide = torch.empty((8, 1024), dtype=torch.uint8)[:, :512] + assert _call(112, kv=wide) == "asm" + wide_rope = torch.empty((8, 128), dtype=torch.bfloat16)[:, :64] + assert _call(112, kv_rope=wide_rope) == "asm" + + +def test_persistent_gets_the_asm_path_tensors(monkeypatch, routes): + _env(monkeypatch) + _call(21) + kind, args, _ = routes[-1] + assert kind == "persist" + kv, kv_indices, kv_indptr, sink, kv_rope, q_packed, q_rope = args + assert kv.shape == (8, 512) and kv_rope.shape == (8, 64) + assert kv_indptr.numel() == 22 and q_packed.shape == (21, 128, 512) + assert sink.numel() == 128 and q_rope.shape == (21, 128, 64) + + +# ---------------------------------------------------------------- the entry + + +def test_decode_trims_indptr_to_real_rows(monkeypatch): + """Eager forward with N < T_pad: pass exactly the N rows the ASM would.""" + calls = [] + ws = object() + monkeypatch.setitem(hca_persist._workspaces, 0, ws) + monkeypatch.setattr(hca_persist, "_dev_index", lambda d: 0) + monkeypatch.setattr( + hca_persist, + "mla_decode_fwd_v4_nm_ps", + lambda *a, **k: calls.append((a, k)), + ) + n, t_pad = 5, 16 + indptr = torch.arange(3, 3 + t_pad + 1, dtype=torch.int32) # offset CSR + out = hca_persist.hca_persist_decode( + torch.empty((4, 512), dtype=torch.uint8), + torch.zeros(t_pad, dtype=torch.int32), + indptr, + torch.zeros(128), + torch.empty((4, 64), dtype=torch.bfloat16), + torch.empty((n, 128, 512), dtype=torch.uint8), + torch.empty((n, 128, 64), dtype=torch.bfloat16), + ) + assert out.shape == (n, 128, 512) and out.dtype == torch.bfloat16 + ((args, kwargs),) = calls + assert args[4].tolist() == list(range(3, 3 + n + 1)) + assert args[7] is ws and kwargs["out"] is out + + +@pytest.mark.parametrize("n", [32769]) +def test_decode_enforces_row_limit(n): + with pytest.raises(ValueError, match="32768"): + hca_persist.hca_persist_decode( + None, + None, + torch.zeros(n + 1, dtype=torch.int32), + None, + None, + torch.empty((n, 128, 1), dtype=torch.uint8), + None, + ) + + +def test_prepare_allocates_once_per_device(monkeypatch): + made = [] + monkeypatch.setattr(hca_persist, "_workspaces", {}) + monkeypatch.setattr(hca_persist, "mla_decode_fwd_v4_nm_ps", object()) + monkeypatch.setattr( + hca_persist, + "get_mla_v4_nm_ps_workspace", + lambda dev, num_partitions: made.append((dev, num_partitions)) or dev, + ) + hca_persist.prepare("cuda:3") + hca_persist.prepare("cuda:3") + assert made == [(torch.device("cuda", 3), hca_persist.PARTITIONS)] + + +# ---------------------------------------------------------------- the builder + + +@pytest.mark.parametrize( + "kv_fp8,hca_layers,heads,gfx,expect", + [ + (True, [2], 128, "gfx950", True), + (False, [2], 128, "gfx950", False), + (True, [], 128, "gfx950", False), + (True, [2], 16, "gfx950", False), + (True, [2], 128, "gfx1250", False), + ], +) +def test_builder_prepares_only_when_usable( + monkeypatch, kv_fp8, hca_layers, heads, gfx, expect +): + from atom.model_ops.attentions import deepseek_v4_attn + + monkeypatch.setenv("ATOM_V4_HCA_PERSIST", "1") + monkeypatch.setattr(deepseek_v4_attn, "get_gfx", lambda: gfx) + prepared = [] + monkeypatch.setattr(hca_persist, "prepare", prepared.append) + builder = SimpleNamespace( + _kv_fp8=kv_fp8, hca_layers=hca_layers, _local_heads=heads, device="cuda:3" + ) + deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._prepare_hca_persist(builder) + assert prepared == (["cuda:3"] if expect else []) + + +def test_builder_does_nothing_when_off(monkeypatch): + from atom.model_ops.attentions import deepseek_v4_attn + + monkeypatch.setenv("ATOM_V4_HCA_PERSIST", "0") + monkeypatch.setattr( + hca_persist, "prepare", lambda d: pytest.fail("prepared while off") + ) + deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._prepare_hca_persist( + SimpleNamespace() + ) From 7d39187a203f6abd512a402fed63ef756504dddf Mon Sep 17 00:00:00 2001 From: yihonglie Date: Tue, 29 Sep 2026 11:58:00 -0500 Subject: [PATCH 07/12] [DSV4] Allocate the persistent HCA workspace at model load on every path Only the native V4 metadata builder called hca_persist.prepare(). The vLLM and SGLang plugins reach the same persistent HCA decode branch without that builder, so their workspace was allocated lazily on the first qualifying call: after KV sizing, and as a hard failure if that call was being captured into a CUDA graph. - Every HCA DeepseekV4Attention layer now calls hca_persist.prepare_if_usable() at construction (switch on, fp8 KV, 128 local heads, gfx950), so model load allocates the workspace in native ATOM and in both plugins, before KV sizing and graph capture. The native builder call stays (idempotent) and shares the eligibility check. - The dispatch gate asks hca_persist.workspace_ready(): a call that finds no workspace allocates it when eager, and stays on the ASM path under capture instead of raising. - Tests: CPU dispatch tests for the capture fallback and the layer-init hook; GPU test capturing a graph with no workspace (ASM result checked against the reference). Signed-off-by: yihonglie Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/model_ops/attentions/deepseek_v4_attn.py | 13 +--- atom/model_ops/v4_kernels/hca_persist.py | 62 ++++++++++++++- atom/model_ops/v4_kernels/paged_decode.py | 1 + atom/models/deepseek_v4.py | 9 +++ tests/model_ops/test_v4_hca_persist_gpu.py | 40 +++++++++- tests/test_v4_hca_persist_dispatch.py | 75 +++++++++++++++++++ 6 files changed, 186 insertions(+), 14 deletions(-) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 9e75e7d6d0..cbf1655813 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -3702,15 +3702,10 @@ def _prepare_hca_persist(self) -> None: CUDA-graph capture, which may be the first forward to reach it.""" if not envs.ATOM_V4_HCA_PERSIST: return - why = None - if not self._kv_fp8: - why = "kv cache is not fp8" - elif not self.hca_layers: - why = "no HCA (ratio 128) layers" - elif self._local_heads != hca_persist.HEADS: - why = f"{self._local_heads} local heads (needs {hca_persist.HEADS})" - elif get_gfx() != "gfx950": - why = f"arch {get_gfx()} (needs gfx950)" + why = "no HCA (ratio 128) layers" if not self.hca_layers else None + why = why or hca_persist.unusable_reason( + kv_fp8=self._kv_fp8, heads=self._local_heads, gfx=get_gfx() + ) if why is not None: logger.info("V4 HCA persistent decode not used: %s", why) return diff --git a/atom/model_ops/v4_kernels/hca_persist.py b/atom/model_ops/v4_kernels/hca_persist.py index 71d98410b4..1073aa5e5f 100644 --- a/atom/model_ops/v4_kernels/hca_persist.py +++ b/atom/model_ops/v4_kernels/hca_persist.py @@ -12,8 +12,13 @@ ``N <= 32768`` rows, gfx950, row-dense KV pools. Everything else, and an aiter without the kernel, stays on the ASM + split plan. -Workspace: one per device, allocated by :func:`prepare` outside graph capture -(the V4 metadata builder calls it at init, before warmup and KV sizing). It +Workspace: one per device, allocated by :func:`prepare` outside graph capture. +Every HCA attention layer asks for it at construction +(:func:`prepare_if_usable`, so model load allocates it in native ATOM and in +the vLLM / SGLang plugins alike, before KV sizing and warmup), and the native +V4 metadata builder asks again at init (idempotent). A call that finds no +workspace allocates it when eager; under CUDA-graph capture it stays on the +ASM path instead (:func:`workspace_ready`). It holds the split partials and the merge counters, which return to zero at the end of every call, so calls sharing it must be ordered on the GPU. In ATOM every HCA decode call runs on the forward's compute stream (layers are @@ -89,6 +94,52 @@ def prepare(device) -> None: ) +def unusable_reason(*, kv_fp8: bool, heads: int, gfx: str) -> str | None: + """Why the persistent kernel cannot serve this model/rank, or None.""" + if not kv_fp8: + return "kv cache is not fp8" + if heads != HEADS: + return f"{heads} local heads (needs {HEADS})" + if gfx != "gfx950": + return f"arch {gfx} (needs gfx950)" + return None + + +def prepare_if_usable(*, kv_fp8: bool, heads: int, gfx: str, device=None) -> bool: + """Allocate the workspace when this rank can use the kernel (switch on, + fp8 KV, 128 local heads, gfx950). Called at model construction, outside + graph capture; ``device`` defaults to the current CUDA device.""" + if not envs.ATOM_V4_HCA_PERSIST: + return False + why = unusable_reason(kv_fp8=kv_fp8, heads=heads, gfx=gfx) + if why is not None: + _log_once(("unused", why), "V4 HCA persistent decode not used: %s", why) + return False + device = torch.device("cuda") if device is None else device + prepare(device) + return _dev_index(device) in _workspaces + + +def workspace_ready(device) -> bool: + """True when this device has a workspace, allocating it if the stream is + not capturing. Under capture with no workspace (nothing prepared it before + capture) the call stays on the ASM path: aiter refuses to allocate there.""" + idx = _dev_index(device) + if idx in _workspaces: + return True + if torch.cuda.is_current_stream_capturing(): + _log_once( + ("no_ws_capture", idx), + "V4 HCA persistent decode: no workspace on device %d at CUDA-graph " + "capture (prepare() was not called before capture); this graph " + "keeps HCA decode on the ASM path", + idx, + ) + return False + prepare(device) + return idx in _workspaces + + def wanted(*, compress_ratio: int | None, heads: int, rows: int, gfx: str) -> bool: """Shape/config gate; the layout gate is :func:`layout_ok`.""" return ( @@ -134,8 +185,11 @@ def hca_persist_decode( f"kv_indptr={kv_indptr.numel()}" ) idx = _dev_index(q_packed_in.device) - if idx not in _workspaces: - prepare(q_packed_in.device) # aiter refuses during capture + if not workspace_ready(q_packed_in.device): + raise RuntimeError( + "hca_persist: no workspace on this device and the stream is " + "capturing; call prepare() before CUDA-graph capture" + ) out = torch.empty((n, HEADS, _DIM), dtype=torch.bfloat16, device=q_packed_in.device) if n == 0: return out diff --git a/atom/model_ops/v4_kernels/paged_decode.py b/atom/model_ops/v4_kernels/paged_decode.py index 8da10e30cb..6522e4ebfc 100644 --- a/atom/model_ops/v4_kernels/paged_decode.py +++ b/atom/model_ops/v4_kernels/paged_decode.py @@ -1221,6 +1221,7 @@ def sparse_attn_v4_paged_decode( gfx=get_gfx(), ) and hca_persist.layout_ok(unified_kv, unified_kv_rope) + and hca_persist.workspace_ready(q_packed_in.device) ): return hca_persist.hca_persist_decode( unified_kv, diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 85c3f75487..ce682927cb 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -104,6 +104,7 @@ from atom.model_ops.triton_hash_topk import hash_topk_triton from atom.model_ops.triton_rmsnorm_nw import rmsnorm_nw from atom.model_ops.utils import atom_parameter, shuffle_weights +from atom.model_ops.v4_kernels import hca_persist from atom.model_ops.v4_kernels import ( FP4_MQA_BLOCK_K, FP4_MQA_PARALLEL_UNIT_NUM, @@ -2631,6 +2632,14 @@ def __init__( # traces unchanged. The rope planes (swa_plane_rope / unified_kv_rope) are # bound onto the module by DeepseekV4AttentionMetadataBuilder. self.kv_fp8 = atom_config.kv_cache_dtype == "fp8" + # Persistent HCA decode workspace: allocate at model load (before KV + # sizing and graph capture) on every serving path, native or plugin. + if self.compress_ratio == hca_persist.HCA_RATIO: + hca_persist.prepare_if_usable( + kv_fp8=str(atom_config.kv_cache_dtype).startswith("fp8"), + heads=self.n_local_heads, + gfx=arch, + ) def process_weights_after_loading(self) -> None: """Prepare wo_a (FP8 + e8m0 block scale) for the grouped output LoRA. diff --git a/tests/model_ops/test_v4_hca_persist_gpu.py b/tests/model_ops/test_v4_hca_persist_gpu.py index 61309e97af..2196d5296e 100644 --- a/tests/model_ops/test_v4_hca_persist_gpu.py +++ b/tests/model_ops/test_v4_hca_persist_gpu.py @@ -208,7 +208,7 @@ def test_cuda_graph_capture_replay_changing_csr(mode): """Capture once at T_pad rows, replay with new kv_indptr / indices / q.""" t_pad, cap_k = 112, 1400 mode(True) - hca_persist.prepare(DEV) # the builder does this at init, before capture + hca_persist.prepare(DEV) # model load does this, before capture first = _inputs(_q7_lens(100, 6) + [0] * 12, seed=6, pool_extra=t_pad * cap_k) pool_rows = first["kv_packed"].shape[0] s = dict( @@ -267,3 +267,41 @@ def stage(lens, seed): # the merge counters are back to zero after every completed call ws = hca_persist._workspaces[torch.cuda.current_device()] assert int(ws.cnt.abs().sum()) == 0 + + +def test_capture_without_workspace_stays_on_asm(mode, monkeypatch): + """Nothing prepared the workspace before capture: the captured call keeps + HCA decode on the ASM path instead of failing, and allocates nothing.""" + mode(True) + monkeypatch.setattr(hca_persist, "_workspaces", {}) + t_pad = 112 + inp = _inputs(_q7_lens(100, 10) + [0] * 12, seed=10) + plan = paged_decode.v4_decode_split_plan( + t_pad, + H, + int((inp["kv_indptr"][1:] - inp["kv_indptr"][:-1]).max()), + torch.zeros(MAX_ROWS_BUF + 1, dtype=torch.int32, device=DEV), + ) + before = hca_persist.stats["persist"] + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = paged_decode.sparse_attn_v4_paged_decode( + None, + inp["kv_packed"], + inp["kv_page_indices"], + inp["kv_indptr"], + inp["sink"], + 512**-0.5, + unified_kv_rope=inp["kv_rope"], + q_packed_in=inp["q_packed"], + q_rope_in=inp["q_rope"], + qo_indptr=torch.arange(t_pad + 1, dtype=torch.int32, device=DEV), + split_plan=plan, + compress_ratio=128, + ) + assert hca_persist.stats["persist"] == before + assert hca_persist._workspaces == {} + graph.replay() + torch.cuda.synchronize() + _close(out[:100], _reference(inp, 100)) diff --git a/tests/test_v4_hca_persist_dispatch.py b/tests/test_v4_hca_persist_dispatch.py index 6131408de8..870558d9df 100644 --- a/tests/test_v4_hca_persist_dispatch.py +++ b/tests/test_v4_hca_persist_dispatch.py @@ -31,6 +31,7 @@ def routes(monkeypatch): lambda *a, **k: calls.append(("persist", a, k)) or "persist", ) monkeypatch.setattr(hca_persist, "mla_decode_fwd_v4_nm_ps", object()) + monkeypatch.setattr(hca_persist, "workspace_ready", lambda device: True) return calls @@ -150,6 +151,80 @@ def test_persistent_gets_the_asm_path_tensors(monkeypatch, routes): assert sink.numel() == 128 and q_rope.shape == (21, 128, 64) +def test_no_workspace_under_capture_stays_on_asm(monkeypatch, routes): + _env(monkeypatch) + monkeypatch.setattr(hca_persist, "workspace_ready", lambda device: False) + assert _call(896) == "asm" + + +# ------------------------------------------------------------ the workspace + + +def _fake_alloc(monkeypatch, capturing): + made = [] + monkeypatch.setattr(hca_persist, "_workspaces", {}) + monkeypatch.setattr(hca_persist, "mla_decode_fwd_v4_nm_ps", object()) + monkeypatch.setattr( + hca_persist, + "get_mla_v4_nm_ps_workspace", + lambda dev, num_partitions: made.append(dev) or dev, + ) + monkeypatch.setattr( + hca_persist.torch.cuda, "is_current_stream_capturing", lambda: capturing + ) + return made + + +def test_workspace_ready_allocates_when_eager(monkeypatch): + made = _fake_alloc(monkeypatch, capturing=False) + assert hca_persist.workspace_ready("cuda:2") + assert hca_persist.workspace_ready("cuda:2") + assert made == [torch.device("cuda", 2)] + + +def test_workspace_ready_never_allocates_under_capture(monkeypatch): + made = _fake_alloc(monkeypatch, capturing=True) + assert not hca_persist.workspace_ready("cuda:2") + assert made == [] + hca_persist._workspaces[2] = object() # prepared before capture + assert hca_persist.workspace_ready("cuda:2") + + +@pytest.mark.parametrize( + "enabled,kv_fp8,heads,gfx,expect", + [ + (True, True, 128, "gfx950", True), + (False, True, 128, "gfx950", False), + (True, False, 128, "gfx950", False), + (True, True, 64, "gfx950", False), + (True, True, 128, "gfx942", False), + ], +) +def test_layer_init_prepares_only_when_usable( + monkeypatch, enabled, kv_fp8, heads, gfx, expect +): + """Model-construction hook shared by native ATOM and the vLLM / SGLang + plugins (their builders never call prepare()).""" + monkeypatch.setenv("ATOM_V4_HCA_PERSIST", "1" if enabled else "0") + made = _fake_alloc(monkeypatch, capturing=False) + got = hca_persist.prepare_if_usable( + kv_fp8=kv_fp8, heads=heads, gfx=gfx, device="cuda:1" + ) + assert got is expect + assert made == ([torch.device("cuda", 1)] if expect else []) + + +def test_attention_layer_calls_prepare_for_hca_only(): + """DeepseekV4Attention.__init__ (every serving path) gates on the ratio.""" + import inspect + + from atom.models import deepseek_v4 + + src = inspect.getsource(deepseek_v4.DeepseekV4Attention.__init__) + assert "hca_persist.prepare_if_usable(" in src + assert "self.compress_ratio == hca_persist.HCA_RATIO" in src + + # ---------------------------------------------------------------- the entry From c8757bd5d7ad310e83e7b5b7de460398f76a56cc Mon Sep 17 00:00:00 2001 From: yihonglie Date: Sat, 3 Oct 2026 08:50:25 -0500 Subject: [PATCH 08/12] [DSV4] Satisfy ruff in the persistent HCA changes Fold hca_persist into the existing v4_kernels import, drop the E402 noqa directives ruff no longer enables, use dict literals in the GPU test, and underscore an unused unpacked tensor. No behaviour change. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: yihonglie --- atom/models/deepseek_v4.py | 2 +- tests/model_ops/test_v4_hca_persist_gpu.py | 42 +++++++++++----------- tests/test_v4_hca_persist_dispatch.py | 2 +- 3 files changed, 23 insertions(+), 23 deletions(-) diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index ce682927cb..fa5ae32139 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -104,7 +104,6 @@ from atom.model_ops.triton_hash_topk import hash_topk_triton from atom.model_ops.triton_rmsnorm_nw import rmsnorm_nw from atom.model_ops.utils import atom_parameter, shuffle_weights -from atom.model_ops.v4_kernels import hca_persist from atom.model_ops.v4_kernels import ( FP4_MQA_BLOCK_K, FP4_MQA_PARALLEL_UNIT_NUM, @@ -113,6 +112,7 @@ csa_translate_pack, fp4_indexer_enabled, fused_compress_attn, + hca_persist, inverse_rope_inplace, qk_norm_rope_maybe_quant, scale_indexer_weights, diff --git a/tests/model_ops/test_v4_hca_persist_gpu.py b/tests/model_ops/test_v4_hca_persist_gpu.py index 2196d5296e..beb277aa87 100644 --- a/tests/model_ops/test_v4_hca_persist_gpu.py +++ b/tests/model_ops/test_v4_hca_persist_gpu.py @@ -17,10 +17,10 @@ if not torch.cuda.is_available(): pytest.skip("needs a GPU", allow_module_level=True) -from aiter.jit.utils.chip_info import get_gfx # noqa: E402 +from aiter.jit.utils.chip_info import get_gfx -from atom.model_ops.v4_kernels import hca_persist, paged_decode # noqa: E402 -from atom.model_ops.v4_kernels.v4_quant import ( # noqa: E402 +from atom.model_ops.v4_kernels import hca_persist, paged_decode +from atom.model_ops.v4_kernels.v4_quant import ( mla_decode_fwd_v4_nm_ref, quantize_bf16_to_v4_2buff, ) @@ -68,15 +68,15 @@ def _inputs(kv_lens, seed, pool_extra=257): indptr = torch.zeros(n + 1, dtype=torch.int32, device=DEV) indptr[1:] = torch.tensor(kv_lens, device=DEV).cumsum(0) idx = torch.randint(0, pool, (tot,), device=DEV, generator=g) - return dict( - q_packed=qp, - q_rope=qr, - kv_packed=kp, - kv_rope=kr, - kv_indptr=indptr, - kv_page_indices=idx.to(torch.int32), - sink=(torch.randn(H, device=DEV, generator=g) * 2).float(), - ) + return { + "q_packed": qp, + "q_rope": qr, + "kv_packed": kp, + "kv_rope": kr, + "kv_indptr": indptr, + "kv_page_indices": idx.to(torch.int32), + "sink": (torch.randn(H, device=DEV, generator=g) * 2).float(), + } def _reference(inp, n): @@ -211,15 +211,15 @@ def test_cuda_graph_capture_replay_changing_csr(mode): hca_persist.prepare(DEV) # model load does this, before capture first = _inputs(_q7_lens(100, 6) + [0] * 12, seed=6, pool_extra=t_pad * cap_k) pool_rows = first["kv_packed"].shape[0] - s = dict( - q_packed=first["q_packed"].clone(), - q_rope=first["q_rope"].clone(), - kv_packed=first["kv_packed"], - kv_rope=first["kv_rope"], - kv_indptr=torch.zeros(t_pad + 1, dtype=torch.int32, device=DEV), - kv_page_indices=torch.zeros(t_pad * cap_k, dtype=torch.int32, device=DEV), - sink=first["sink"], - ) + s = { + "q_packed": first["q_packed"].clone(), + "q_rope": first["q_rope"].clone(), + "kv_packed": first["kv_packed"], + "kv_rope": first["kv_rope"], + "kv_indptr": torch.zeros(t_pad + 1, dtype=torch.int32, device=DEV), + "kv_page_indices": torch.zeros(t_pad * cap_k, dtype=torch.int32, device=DEV), + "sink": first["sink"], + } def stage(lens, seed): g = torch.Generator(device=DEV).manual_seed(seed) diff --git a/tests/test_v4_hca_persist_dispatch.py b/tests/test_v4_hca_persist_dispatch.py index 870558d9df..1c76ff519d 100644 --- a/tests/test_v4_hca_persist_dispatch.py +++ b/tests/test_v4_hca_persist_dispatch.py @@ -145,7 +145,7 @@ def test_persistent_gets_the_asm_path_tensors(monkeypatch, routes): _call(21) kind, args, _ = routes[-1] assert kind == "persist" - kv, kv_indices, kv_indptr, sink, kv_rope, q_packed, q_rope = args + kv, _kv_indices, kv_indptr, sink, kv_rope, q_packed, q_rope = args assert kv.shape == (8, 512) and kv_rope.shape == (8, 64) assert kv_indptr.numel() == 22 and q_packed.shape == (21, 128, 512) assert sink.numel() == 128 and q_rope.shape == (21, 128, 64) From 425f8852dd896e2864e30edde507218515e4aa2f Mon Sep 17 00:00:00 2001 From: yihonglie Date: Sat, 3 Oct 2026 22:42:39 -0500 Subject: [PATCH 09/12] [DSV4] Fix the non-GPU unit tests for the split-plan and HCA imports test_deepseek_v4_transfer_regions loads the real deepseek_v4_attn against a fixed set of import stubs, which did not cover what the decode changes added: aiter.dist.parallel_state.get_tensor_model_parallel_world_size, v4_kernels.v4_decode_split_plan and v4_kernels.hca_persist. Under the CPU runner the module then failed to import (41 errors); with test_pp collected first, its narrower parallel_state stub was picked up instead. Supply the three in the test's own stub set (hca_persist reports itself unusable, as it is without aiter), and use a dict literal in the prefill dispatch test. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: yihonglie --- tests/test_deepseek_v4_transfer_regions.py | 14 ++++++++++++++ tests/test_v4_prefill_asm_decode_dispatch.py | 20 ++++++++++---------- 2 files changed, 24 insertions(+), 10 deletions(-) diff --git a/tests/test_deepseek_v4_transfer_regions.py b/tests/test_deepseek_v4_transfer_regions.py index 76b97fa4d6..abba648f96 100644 --- a/tests/test_deepseek_v4_transfer_regions.py +++ b/tests/test_deepseek_v4_transfer_regions.py @@ -126,6 +126,12 @@ def _stub_v4_runtime_imports(): aiter_jit_utils.__path__ = [] chip_info = types.ModuleType("aiter.jit.utils.chip_info") chip_info.get_gfx = lambda: "gfx950" + # Replaced, not reused: other test modules leave their own, narrower + # parallel_state stub in sys.modules. + aiter_dist = types.ModuleType("aiter.dist") + aiter_dist.__path__ = [] + parallel_state = types.ModuleType("aiter.dist.parallel_state") + parallel_state.get_tensor_model_parallel_world_size = lambda: 1 pcp_utils = types.ModuleType("atom.distributed.pcp_utils") for name in ( @@ -155,16 +161,24 @@ def _stub_v4_runtime_imports(): "fp4_indexer_enabled", "hca_compress_paged_offsets", "plan_context_lens", + "v4_decode_split_plan", "write_v4_paged_decode_indices", "write_v4_paged_prefill_indices", ): setattr(kernels, name, lambda *args, **kwargs: None) + # No aiter here, so the persistent HCA decode is never usable. + kernels.hca_persist = SimpleNamespace( + unusable_reason=lambda **kwargs: "no aiter in this test", + prepare=lambda *args, **kwargs: None, + ) replacements = { "aiter": aiter, "aiter.jit": aiter_jit, "aiter.jit.utils": aiter_jit_utils, "aiter.jit.utils.chip_info": chip_info, + "aiter.dist": aiter_dist, + "aiter.dist.parallel_state": parallel_state, "atom.distributed.pcp_utils": pcp_utils, "atom.model_ops.attentions.backends": backends, "atom.model_ops.v4_kernels": kernels, diff --git a/tests/test_v4_prefill_asm_decode_dispatch.py b/tests/test_v4_prefill_asm_decode_dispatch.py index bd8fb7c218..d550bd7e94 100644 --- a/tests/test_v4_prefill_asm_decode_dispatch.py +++ b/tests/test_v4_prefill_asm_decode_dispatch.py @@ -116,16 +116,16 @@ def test_ineligible_decode_keeps_dedicated_asm(monkeypatch, kwargs): 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)), - ) + return { + "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( From b6f4d7d333edee1630778b8eabdf287f03875747 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Sat, 3 Oct 2026 23:47:51 -0500 Subject: [PATCH 10/12] [DSV4] Skip the HCA split plan where the persistent kernel runs; review fixes - The metadata builder no longer builds the HCA split plan for forwards whose HCA decode the persistent kernel takes (it plans on the GPU from the real kv_indptr). The decision uses only init-time state (workspace allocated before any capture) and the row gate shared with dispatch (hca_persist.rows_ok), so capture and every replay agree; a call the kernel then declines gets aiter's own split pick. - _prepare_hca_persist delegates to hca_persist.prepare_if_usable and returns whether HCA can take the kernel. - DeepseekV4Attention passes its own kv_fp8 flag to prepare_if_usable. - Docs: the workspace is allocated at model load, not by the builder; the split-plan helper reference points at ROCm/aiter#6126. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: yihonglie --- atom/model_ops/attentions/deepseek_v4_attn.py | 49 ++++++++++++------- atom/model_ops/v4_kernels/hca_persist.py | 7 ++- atom/model_ops/v4_kernels/paged_decode.py | 2 +- atom/models/deepseek_v4.py | 2 +- docs/environment_variables.md | 2 +- tests/test_v4_hca_persist_dispatch.py | 36 ++++++++++++++ 6 files changed, 76 insertions(+), 22 deletions(-) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index b7ff3ad5a3..74118dca50 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -306,7 +306,8 @@ class AttentionMetaData_DSV4(AttentionMetaData): 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.""" + """Same for HCA layers; also None when the persistent kernel + (`hca_persist`) takes this forward's HCA decode.""" 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 @@ -447,6 +448,10 @@ class DeepseekV4AttentionMetadataBuilder(CommonAttentionBuilder): `config.kv_cache_block_size` (config.py forces the same value for V4). """ + # Set at init by `_prepare_hca_persist`: HCA decode can take aiter's + # persistent kernel (workspace allocated before any capture). + _hca_persist: bool = False + capture_owns_cu_seqlens_q = True block_size = 256 @@ -726,7 +731,7 @@ def __init__(self, model_runner): model_runner.config.hf_config.num_attention_heads // get_tensor_model_parallel_world_size() ) - self._prepare_hca_persist() + self._hca_persist = self._prepare_hca_persist() # Sparse-attn + per-fwd metadata buffers (CG-A: pre-allocate for fixed # GPU pointers, prerequisite for CUDAGraph capture). All H2D copies in @@ -3678,20 +3683,22 @@ def _attach_v4_per_fwd_meta( buf_prefix_ubatch=buf_prefix_ubatch, ) - def _prepare_hca_persist(self) -> None: + def _prepare_hca_persist(self) -> bool: """Allocate the persistent HCA decode workspace here, at init: before warmup, before KV sizing (so it is accounted for) and before any - CUDA-graph capture, which may be the first forward to reach it.""" - if not envs.ATOM_V4_HCA_PERSIST: - return - why = "no HCA (ratio 128) layers" if not self.hca_layers else None - why = why or hca_persist.unusable_reason( - kv_fp8=self._kv_fp8, heads=self._local_heads, gfx=get_gfx() + CUDA-graph capture, which may be the first forward to reach it. + + True when HCA decode can take the persistent kernel. Fixed for the + builder's lifetime, so capture and every replay see the same answer. + """ + if not envs.ATOM_V4_HCA_PERSIST or not self.hca_layers: + return False + return hca_persist.prepare_if_usable( + kv_fp8=self._kv_fp8, + heads=self._local_heads, + gfx=get_gfx(), + device=self.device, ) - if why is not None: - logger.info("V4 HCA persistent decode not used: %s", why) - return - hca_persist.prepare(self.device) def _decode_split_plan( self, rows: int, kv_len: int, split_indptr: torch.Tensor @@ -3837,11 +3844,17 @@ def _attach_v4_paged_decode_meta( 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"], - ) + # HCA calls the persistent kernel takes plan their split on the GPU from + # the real kv_indptr; a call it then declines (non-dense pool, eager + # N below the row gate) gets aiter's own split pick. + if self._hca_persist and hca_persist.rows_ok(T_pad): + attn_metadata.split_plan_hca = None + else: + 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 diff --git a/atom/model_ops/v4_kernels/hca_persist.py b/atom/model_ops/v4_kernels/hca_persist.py index 1073aa5e5f..48b8d75d10 100644 --- a/atom/model_ops/v4_kernels/hca_persist.py +++ b/atom/model_ops/v4_kernels/hca_persist.py @@ -140,6 +140,11 @@ def workspace_ready(device) -> bool: return idx in _workspaces +def rows_ok(rows: int) -> bool: + """Row-count gate: smaller calls stay on the ASM path, which is faster there.""" + return envs.ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= MAX_ROWS + + def wanted(*, compress_ratio: int | None, heads: int, rows: int, gfx: str) -> bool: """Shape/config gate; the layout gate is :func:`layout_ok`.""" return ( @@ -147,7 +152,7 @@ def wanted(*, compress_ratio: int | None, heads: int, rows: int, gfx: str) -> bo and compress_ratio == HCA_RATIO and heads == HEADS and gfx == "gfx950" - and envs.ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= MAX_ROWS + and rows_ok(rows) and available() ) diff --git a/atom/model_ops/v4_kernels/paged_decode.py b/atom/model_ops/v4_kernels/paged_decode.py index 6522e4ebfc..cc235952d3 100644 --- a/atom/model_ops/v4_kernels/paged_decode.py +++ b/atom/model_ops/v4_kernels/paged_decode.py @@ -996,7 +996,7 @@ def v4_decode_split_plan( `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. + before ROCm/aiter#6126, which leaves the split pick to aiter. """ import aiter.mla diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 3e93382794..76386baf8d 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -2636,7 +2636,7 @@ def __init__( # sizing and graph capture) on every serving path, native or plugin. if self.compress_ratio == hca_persist.HCA_RATIO: hca_persist.prepare_if_usable( - kv_fp8=str(atom_config.kv_cache_dtype).startswith("fp8"), + kv_fp8=self.kv_fp8, heads=self.n_local_heads, gfx=arch, ) diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 469aa66f58..3c06dd1415 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -110,7 +110,7 @@ make duplicate prefill useful. Pure-attention models do not use checkpoint waits | **ATOM_USE_FLYDSL_FP8_PREFILL_ATTN** | bool | 0 (false) | Enable FlyDSL FP8 MLA prefill on gfx950, including fused K concatenation with QKV quantization and direct FP8 output from FlyDSL gather where supported. AITER capability checks are cached per layer. Missing FP8 attention support or unsupported device/model configurations select BF16 attention before quantization and emit a warning once per process. FP16 activations and synthetic RoPE-padding configurations also use the normal attention path. Cached K/V reuse descales `max(new_token_descale, 1e-6) * 2`; larger outliers saturate. Unsupported FP8 gather configurations use BF16 gather followed by fused K/V dynamic quantization when FP8 attention is supported. Kernel execution errors propagate without retry. Added 2026-09-10. | | **ATOM_PA_FLYDSL** | bool | 0 (false) | Route the MHA paged decode to aiter's FlyDSL kernel (aiter #4332) instead of the gluon one, where FlyDSL's domain covers the call. Off by default: the kernel is not fully tested on every shape ATOM ships. The env is the first of two gates; the second mirrors aiter's own validation so an unsupported shape falls back to gluon rather than raising inside aiter, and the `and` short-circuits so nothing is evaluated when this is off. Global, not per call site: the vLLM bridge is rerouted too. Needs aiter `94dca7bc6` or later. | | **ATOM_PA_FLYDSL_PLAN** | bool | 1 (true) | Use aiter #5546's GPU work planner on the dense decode, built once per forward in the metadata builder. Requires `ATOM_PA_FLYDSL=1` — with FlyDSL off no plan is built at all. Sizes each request's partition count from its real context length instead of splitting the batch uniformly, which is worth p90 interactivity `+20.5%` at concurrency 20 on the MiniMax-M3 agentic trace and nothing at concurrency 1, where there is nothing to rebalance. Kept off the two MiniMax-M3 sparse call sites, whose contexts are a fixed topk window and where it measures a net loss. Plans are only minted during cudagraph capture — aiter's planner takes the batch as a `tl.constexpr`, so a new value is a cold kernel specialization and a plan no captured graph can ever release — which leaves the planner inert under `--enforce-eager`. | -| **ATOM_V4_HCA_PERSIST** | bool | 1 (true) | DeepSeek-V4 HCA (`compress_ratio` 128) fp8 decode through aiter's persistent V4-NM kernel (`aiter.mla.mla_decode_fwd_v4_nm_ps`): one launch that splits the KV on the GPU and merges in-kernel, instead of the decode ASM plus a host-built split plan. Taken only with fp8 KV, 128 local heads, gfx950, row-dense KV pools and `ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= 32768`; everything else stays on the ASM path. An aiter without the kernel keeps the ASM path (logged once). One workspace per device is allocated by the V4 metadata builder at init. Set `0` to force the ASM path. | +| **ATOM_V4_HCA_PERSIST** | bool | 1 (true) | DeepSeek-V4 HCA (`compress_ratio` 128) fp8 decode through aiter's persistent V4-NM kernel (`aiter.mla.mla_decode_fwd_v4_nm_ps`): one launch that splits the KV on the GPU and merges in-kernel, instead of the decode ASM plus a host-built split plan. Taken only with fp8 KV, 128 local heads, gfx950, row-dense KV pools and `ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= 32768`; everything else stays on the ASM path. An aiter without the kernel keeps the ASM path (logged once). One workspace per device (about 64 MiB) is allocated at model load, before KV sizing and CUDA-graph capture; the metadata builder then skips the host split plan for the HCA calls this kernel takes. Set `0` to force the ASM path. | | **ATOM_V4_HCA_PERSIST_MIN_ROWS** | int | 15 | Smallest decode call (q rows) that uses `ATOM_V4_HCA_PERSIST`; smaller calls stay on the ASM path. | ### GLM-5.3 diff --git a/tests/test_v4_hca_persist_dispatch.py b/tests/test_v4_hca_persist_dispatch.py index 1c76ff519d..5bb9160b2c 100644 --- a/tests/test_v4_hca_persist_dispatch.py +++ b/tests/test_v4_hca_persist_dispatch.py @@ -323,3 +323,39 @@ def test_builder_does_nothing_when_off(monkeypatch): deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._prepare_hca_persist( SimpleNamespace() ) + + +def test_builder_reports_whether_hca_takes_the_kernel(monkeypatch): + from atom.model_ops.attentions import deepseek_v4_attn + + monkeypatch.setenv("ATOM_V4_HCA_PERSIST", "1") + monkeypatch.setattr(deepseek_v4_attn, "get_gfx", lambda: "gfx950") + monkeypatch.setattr(hca_persist, "prepare_if_usable", lambda **kw: True) + builder = SimpleNamespace( + _kv_fp8=True, hca_layers=[2], _local_heads=128, device="cuda:0" + ) + prepare = deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._prepare_hca_persist + assert prepare(builder) is True + builder.hca_layers = [] + assert prepare(builder) is False + + +@pytest.mark.parametrize( + "rows,ok", [(14, False), (15, True), (32768, True), (32769, False)] +) +def test_rows_gate_shared_by_builder_and_dispatch(monkeypatch, rows, ok): + """The builder skips the HCA split plan exactly for the rows dispatch sends + to the persistent kernel.""" + monkeypatch.setenv("ATOM_V4_HCA_PERSIST_MIN_ROWS", "15") + assert hca_persist.rows_ok(rows) is ok + + +def test_builder_skips_hca_plan_only_for_persistent_rows(): + import inspect + + from atom.model_ops.attentions import deepseek_v4_attn + + src = inspect.getsource( + deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._attach_v4_paged_decode_meta + ) + assert "self._hca_persist and hca_persist.rows_ok(T_pad)" in src From a49fc9d156530fd6c64dec7dc3dbf0ad475b36af Mon Sep 17 00:00:00 2001 From: yihonglie Date: Sun, 4 Oct 2026 00:03:41 -0500 Subject: [PATCH 11/12] [DSV4] Build decode split plans on the host from one constant table Every fp8 decode split plan is now (s, split_table[s - 1]): aiter's host-only split count (get_mla_v4_nm_num_kv_splits) and a row of a [16, T + 1] uniform split_indptr table built once at init (v4_uniform_split_table). No forward writes a split_indptr any more, and CUDA graphs of every size read constant memory, so a graph can no longer see another size's plan. - The builder's per-layer-type, per-ubatch and MTP-draft split_indptr buffers are replaced by the one shared table; DSparkIndexBuffers keeps its own table sized for its draft rows. - With planning free, the HCA plan is built again for every forward (the persistent kernel ignores it), so the builder no longer tracks whether HCA takes that kernel. - A split count beyond the table, or rows beyond it, leaves the pick to aiter. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: yihonglie --- atom/model_ops/attentions/deepseek_v4_attn.py | 77 ++++++------------- atom/model_ops/v4_kernels/__init__.py | 2 + .../v4_kernels/dspark_fp8_indices.py | 11 ++- atom/model_ops/v4_kernels/hca_persist.py | 7 +- atom/model_ops/v4_kernels/paged_decode.py | 34 ++++++-- docs/environment_variables.md | 2 +- tests/model_ops/test_v4_hca_persist_gpu.py | 4 +- tests/test_deepseek_v4_transfer_regions.py | 1 + tests/test_v4_hca_persist_dispatch.py | 36 --------- tests/test_v4_prefill_asm_decode_dispatch.py | 58 ++++++++++---- 10 files changed, 109 insertions(+), 123 deletions(-) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 74118dca50..5b29a163fa 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -123,6 +123,7 @@ hca_persist, plan_context_lens, v4_decode_split_plan, + v4_uniform_split_table, write_v4_paged_decode_indices, write_v4_paged_prefill_indices, ) @@ -306,8 +307,7 @@ class AttentionMetaData_DSV4(AttentionMetaData): 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; also None when the persistent kernel - (`hca_persist`) takes this forward's HCA decode.""" + """Same for HCA layers (unused by calls the persistent kernel takes).""" 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 @@ -448,10 +448,6 @@ class DeepseekV4AttentionMetadataBuilder(CommonAttentionBuilder): `config.kv_cache_block_size` (config.py forces the same value for V4). """ - # Set at init by `_prepare_hca_persist`: HCA decode can take aiter's - # persistent kernel (workspace allocated before any capture). - _hca_persist: bool = False - capture_owns_cu_seqlens_q = True block_size = 256 @@ -731,7 +727,7 @@ def __init__(self, model_runner): model_runner.config.hf_config.num_attention_heads // get_tensor_model_parallel_world_size() ) - self._hca_persist = self._prepare_hca_persist() + self._prepare_hca_persist() # Sparse-attn + per-fwd metadata buffers (CG-A: pre-allocate for fixed # GPU pointers, prerequisite for CUDAGraph capture). All H2D copies in @@ -2611,7 +2607,7 @@ def prepare_mtp_decode( 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"] + running_bs, self.window_size ) attn_metadata.batch_id_per_q_token = batch_id_per_q_token @@ -3683,36 +3679,30 @@ def _attach_v4_per_fwd_meta( buf_prefix_ubatch=buf_prefix_ubatch, ) - def _prepare_hca_persist(self) -> bool: + def _prepare_hca_persist(self) -> None: """Allocate the persistent HCA decode workspace here, at init: before warmup, before KV sizing (so it is accounted for) and before any - CUDA-graph capture, which may be the first forward to reach it. - - True when HCA decode can take the persistent kernel. Fixed for the - builder's lifetime, so capture and every replay see the same answer. - """ - if not envs.ATOM_V4_HCA_PERSIST or not self.hca_layers: - return False - return hca_persist.prepare_if_usable( - kv_fp8=self._kv_fp8, - heads=self._local_heads, - gfx=get_gfx(), - device=self.device, - ) + CUDA-graph capture, which may be the first forward to reach it.""" + if envs.ATOM_V4_HCA_PERSIST and self.hca_layers: + hca_persist.prepare_if_usable( + kv_fp8=self._kv_fp8, + heads=self._local_heads, + gfx=get_gfx(), + device=self.device, + ) def _decode_split_plan( - self, rows: int, kv_len: int, split_indptr: torch.Tensor + self, rows: int, kv_len: int ) -> 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. + Host only: `rows` and `kv_len` are fixed for a captured graph, so capture + and every replay pick the same split count, and the plan's + `split_indptr` is a slice of the constant `_split_table`. """ if not self._kv_fp8: return None - return v4_decode_split_plan(rows, self._local_heads, kv_len, split_indptr) + return v4_decode_split_plan(rows, self._local_heads, kv_len, self._split_table) def _attach_v4_paged_decode_meta( self, @@ -3838,23 +3828,11 @@ def _attach_v4_paged_decode_meta( ) # 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_swa = self._decode_split_plan(T_pad, win) + attn_metadata.split_plan_csa = self._decode_split_plan(T_pad, win + index_topk) + attn_metadata.split_plan_hca = self._decode_split_plan( + T_pad, win + self.max_committed_hca ) - # HCA calls the persistent kernel takes plan their split on the GPU from - # the real kv_indptr; a call it then declines (non-dense pool, eager - # N below the row gate) gets aiter's own split pick. - if self._hca_persist and hca_persist.rows_ok(T_pad): - attn_metadata.split_plan_hca = None - else: - 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 @@ -4678,10 +4656,9 @@ 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) + # Every decode split plan (verify, MTP draft, TBO ubatches) is a slice + # of this constant table; read-only, so all of them share it. + self._split_table = v4_uniform_split_table(T_dec, self.device) 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. @@ -4709,7 +4686,6 @@ 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. @@ -4882,9 +4858,6 @@ 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( diff --git a/atom/model_ops/v4_kernels/__init__.py b/atom/model_ops/v4_kernels/__init__.py index 6b77b185f1..7c2cab8f95 100644 --- a/atom/model_ops/v4_kernels/__init__.py +++ b/atom/model_ops/v4_kernels/__init__.py @@ -35,6 +35,7 @@ sparse_attn_v4_paged_decode, sparse_attn_v4_paged_decode_reference, v4_decode_split_plan, + v4_uniform_split_table, ) from atom.model_ops.v4_kernels.paged_decode_indices import ( build_v4_paged_decode_indptr, @@ -91,6 +92,7 @@ "swa_write_2buff_prepacked", "update_compressor_states", "v4_decode_split_plan", + "v4_uniform_split_table", "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 fc77e49e66..b1ad30f7bd 100644 --- a/atom/model_ops/v4_kernels/dspark_fp8_indices.py +++ b/atom/model_ops/v4_kernels/dspark_fp8_indices.py @@ -41,7 +41,10 @@ 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.paged_decode import ( + v4_decode_split_plan, + v4_uniform_split_table, +) from atom.model_ops.v4_kernels.pool_index import window_constexprs, window_row @@ -143,7 +146,7 @@ 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` + split_table: torch.Tensor # [16, max_b*T+1] int32, constant, 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 @@ -178,7 +181,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), + split_table=v4_uniform_split_table(n, device), batch_ids=torch.arange(n, **i32) // draft, max_batch=max_batch, draft_width=draft, @@ -235,7 +238,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) + self.split_plan = v4_decode_split_plan(B * T, heads, W + T, self.split_table) # 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/hca_persist.py b/atom/model_ops/v4_kernels/hca_persist.py index 48b8d75d10..1073aa5e5f 100644 --- a/atom/model_ops/v4_kernels/hca_persist.py +++ b/atom/model_ops/v4_kernels/hca_persist.py @@ -140,11 +140,6 @@ def workspace_ready(device) -> bool: return idx in _workspaces -def rows_ok(rows: int) -> bool: - """Row-count gate: smaller calls stay on the ASM path, which is faster there.""" - return envs.ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= MAX_ROWS - - def wanted(*, compress_ratio: int | None, heads: int, rows: int, gfx: str) -> bool: """Shape/config gate; the layout gate is :func:`layout_ok`.""" return ( @@ -152,7 +147,7 @@ def wanted(*, compress_ratio: int | None, heads: int, rows: int, gfx: str) -> bo and compress_ratio == HCA_RATIO and heads == HEADS and gfx == "gfx950" - and rows_ok(rows) + and envs.ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= MAX_ROWS and available() ) diff --git a/atom/model_ops/v4_kernels/paged_decode.py b/atom/model_ops/v4_kernels/paged_decode.py index cc235952d3..48296802b4 100644 --- a/atom/model_ops/v4_kernels/paged_decode.py +++ b/atom/model_ops/v4_kernels/paged_decode.py @@ -986,24 +986,44 @@ def _sparse_attn_v4_paged_decode_prefill_asm( ) +# aiter never picks more than 16 KV splits for the v4 nm decode kernel. +V4_DECODE_MAX_SPLITS = 16 + + +def v4_uniform_split_table(rows: int, device) -> torch.Tensor: + """``[V4_DECODE_MAX_SPLITS, rows + 1]`` int32; row ``s - 1`` is the uniform + ``split_indptr`` ``[0, s, 2s, ...]`` of ``s`` splits per row. + + Built once, outside CUDA-graph capture, and never written again: every + split plan is a slice of it, so graphs of any size read constant memory and + no forward rewrites a plan before it runs. + """ + splits = torch.arange(1, V4_DECODE_MAX_SPLITS + 1, dtype=torch.int32, device=device) + return splits[:, None] * torch.arange(rows + 1, dtype=torch.int32, device=device) + + def v4_decode_split_plan( - rows: int, heads: int, kv_len: int, split_indptr: torch.Tensor + rows: int, heads: int, kv_len: int, split_table: 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 + entries), `heads` local heads and at most `kv_len` KV per row. Host only: + aiter picks the split count ``s`` and the plan's ``split_indptr`` is row + ``s - 1`` of `split_table` (:func:`v4_uniform_split_table`), whose prefix + the decode call trims to its rows. Pass it on as ``sparse_attn_v4_paged_decode(..., split_plan=plan)``. None on aiter builds before ROCm/aiter#6126, 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: + num_kv_splits = getattr(aiter.mla, "get_mla_v4_nm_num_kv_splits", None) + if num_kv_splits is None: return None - return plan_fn(rows, heads, kv_len, split_indptr=split_indptr) + s = num_kv_splits(rows, heads, kv_len) + if s > split_table.shape[0] or rows >= split_table.shape[1]: + return None # outside the table: let aiter pick + return s, split_table[s - 1] def _sparse_attn_v4_paged_decode_asm( diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 3c06dd1415..332eef3284 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -110,7 +110,7 @@ make duplicate prefill useful. Pure-attention models do not use checkpoint waits | **ATOM_USE_FLYDSL_FP8_PREFILL_ATTN** | bool | 0 (false) | Enable FlyDSL FP8 MLA prefill on gfx950, including fused K concatenation with QKV quantization and direct FP8 output from FlyDSL gather where supported. AITER capability checks are cached per layer. Missing FP8 attention support or unsupported device/model configurations select BF16 attention before quantization and emit a warning once per process. FP16 activations and synthetic RoPE-padding configurations also use the normal attention path. Cached K/V reuse descales `max(new_token_descale, 1e-6) * 2`; larger outliers saturate. Unsupported FP8 gather configurations use BF16 gather followed by fused K/V dynamic quantization when FP8 attention is supported. Kernel execution errors propagate without retry. Added 2026-09-10. | | **ATOM_PA_FLYDSL** | bool | 0 (false) | Route the MHA paged decode to aiter's FlyDSL kernel (aiter #4332) instead of the gluon one, where FlyDSL's domain covers the call. Off by default: the kernel is not fully tested on every shape ATOM ships. The env is the first of two gates; the second mirrors aiter's own validation so an unsupported shape falls back to gluon rather than raising inside aiter, and the `and` short-circuits so nothing is evaluated when this is off. Global, not per call site: the vLLM bridge is rerouted too. Needs aiter `94dca7bc6` or later. | | **ATOM_PA_FLYDSL_PLAN** | bool | 1 (true) | Use aiter #5546's GPU work planner on the dense decode, built once per forward in the metadata builder. Requires `ATOM_PA_FLYDSL=1` — with FlyDSL off no plan is built at all. Sizes each request's partition count from its real context length instead of splitting the batch uniformly, which is worth p90 interactivity `+20.5%` at concurrency 20 on the MiniMax-M3 agentic trace and nothing at concurrency 1, where there is nothing to rebalance. Kept off the two MiniMax-M3 sparse call sites, whose contexts are a fixed topk window and where it measures a net loss. Plans are only minted during cudagraph capture — aiter's planner takes the batch as a `tl.constexpr`, so a new value is a cold kernel specialization and a plan no captured graph can ever release — which leaves the planner inert under `--enforce-eager`. | -| **ATOM_V4_HCA_PERSIST** | bool | 1 (true) | DeepSeek-V4 HCA (`compress_ratio` 128) fp8 decode through aiter's persistent V4-NM kernel (`aiter.mla.mla_decode_fwd_v4_nm_ps`): one launch that splits the KV on the GPU and merges in-kernel, instead of the decode ASM plus a host-built split plan. Taken only with fp8 KV, 128 local heads, gfx950, row-dense KV pools and `ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= 32768`; everything else stays on the ASM path. An aiter without the kernel keeps the ASM path (logged once). One workspace per device (about 64 MiB) is allocated at model load, before KV sizing and CUDA-graph capture; the metadata builder then skips the host split plan for the HCA calls this kernel takes. Set `0` to force the ASM path. | +| **ATOM_V4_HCA_PERSIST** | bool | 1 (true) | DeepSeek-V4 HCA (`compress_ratio` 128) fp8 decode through aiter's persistent V4-NM kernel (`aiter.mla.mla_decode_fwd_v4_nm_ps`): one launch that splits the KV on the GPU and merges in-kernel, instead of the decode ASM plus a host-built split plan. Taken only with fp8 KV, 128 local heads, gfx950, row-dense KV pools and `ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= 32768`; everything else stays on the ASM path. An aiter without the kernel keeps the ASM path (logged once). One workspace per device (about 64 MiB) is allocated at model load, before KV sizing and CUDA-graph capture. Set `0` to force the ASM path. | | **ATOM_V4_HCA_PERSIST_MIN_ROWS** | int | 15 | Smallest decode call (q rows) that uses `ATOM_V4_HCA_PERSIST`; smaller calls stay on the ASM path. | ### GLM-5.3 diff --git a/tests/model_ops/test_v4_hca_persist_gpu.py b/tests/model_ops/test_v4_hca_persist_gpu.py index beb277aa87..ab0d88480e 100644 --- a/tests/model_ops/test_v4_hca_persist_gpu.py +++ b/tests/model_ops/test_v4_hca_persist_gpu.py @@ -122,7 +122,7 @@ def _run(inp, *, n=None, t_pad=None): kv_indptr = inp["kv_indptr"][: t_pad + 1] kv_len = int((kv_indptr[1:] - kv_indptr[:-1]).max()) plan = paged_decode.v4_decode_split_plan( - t_pad, H, kv_len, torch.zeros(MAX_ROWS_BUF + 1, dtype=torch.int32, device=DEV) + t_pad, H, kv_len, paged_decode.v4_uniform_split_table(MAX_ROWS_BUF, DEV) ) return paged_decode.sparse_attn_v4_paged_decode( None, @@ -280,7 +280,7 @@ def test_capture_without_workspace_stays_on_asm(mode, monkeypatch): t_pad, H, int((inp["kv_indptr"][1:] - inp["kv_indptr"][:-1]).max()), - torch.zeros(MAX_ROWS_BUF + 1, dtype=torch.int32, device=DEV), + paged_decode.v4_uniform_split_table(MAX_ROWS_BUF, DEV), ) before = hca_persist.stats["persist"] torch.cuda.synchronize() diff --git a/tests/test_deepseek_v4_transfer_regions.py b/tests/test_deepseek_v4_transfer_regions.py index abba648f96..59aa578913 100644 --- a/tests/test_deepseek_v4_transfer_regions.py +++ b/tests/test_deepseek_v4_transfer_regions.py @@ -162,6 +162,7 @@ def _stub_v4_runtime_imports(): "hca_compress_paged_offsets", "plan_context_lens", "v4_decode_split_plan", + "v4_uniform_split_table", "write_v4_paged_decode_indices", "write_v4_paged_prefill_indices", ): diff --git a/tests/test_v4_hca_persist_dispatch.py b/tests/test_v4_hca_persist_dispatch.py index 5bb9160b2c..1c76ff519d 100644 --- a/tests/test_v4_hca_persist_dispatch.py +++ b/tests/test_v4_hca_persist_dispatch.py @@ -323,39 +323,3 @@ def test_builder_does_nothing_when_off(monkeypatch): deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._prepare_hca_persist( SimpleNamespace() ) - - -def test_builder_reports_whether_hca_takes_the_kernel(monkeypatch): - from atom.model_ops.attentions import deepseek_v4_attn - - monkeypatch.setenv("ATOM_V4_HCA_PERSIST", "1") - monkeypatch.setattr(deepseek_v4_attn, "get_gfx", lambda: "gfx950") - monkeypatch.setattr(hca_persist, "prepare_if_usable", lambda **kw: True) - builder = SimpleNamespace( - _kv_fp8=True, hca_layers=[2], _local_heads=128, device="cuda:0" - ) - prepare = deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._prepare_hca_persist - assert prepare(builder) is True - builder.hca_layers = [] - assert prepare(builder) is False - - -@pytest.mark.parametrize( - "rows,ok", [(14, False), (15, True), (32768, True), (32769, False)] -) -def test_rows_gate_shared_by_builder_and_dispatch(monkeypatch, rows, ok): - """The builder skips the HCA split plan exactly for the rows dispatch sends - to the persistent kernel.""" - monkeypatch.setenv("ATOM_V4_HCA_PERSIST_MIN_ROWS", "15") - assert hca_persist.rows_ok(rows) is ok - - -def test_builder_skips_hca_plan_only_for_persistent_rows(): - import inspect - - from atom.model_ops.attentions import deepseek_v4_attn - - src = inspect.getsource( - deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._attach_v4_paged_decode_meta - ) - assert "self._hca_persist and hca_persist.rows_ok(T_pad)" in src diff --git a/tests/test_v4_prefill_asm_decode_dispatch.py b/tests/test_v4_prefill_asm_decode_dispatch.py index d550bd7e94..5190c52fbf 100644 --- a/tests/test_v4_prefill_asm_decode_dispatch.py +++ b/tests/test_v4_prefill_asm_decode_dispatch.py @@ -187,32 +187,58 @@ def test_decode_asm_trims_split_indptr_to_real_rows(monkeypatch): 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): +def test_uniform_split_table_rows(): + table = paged_decode.v4_uniform_split_table(4, "cpu") + assert table.shape == (paged_decode.V4_DECODE_MAX_SPLITS, 5) + assert table.dtype == torch.int32 + assert table[0].tolist() == [0, 1, 2, 3, 4] + assert table[2].tolist() == [0, 3, 6, 9, 12] + + +@pytest.mark.parametrize("has_count", [True, False]) +def test_v4_decode_split_plan_follows_aiter(monkeypatch, has_count): import aiter.mla calls = [] - def planner(rows, heads, kv_len, *, split_indptr): + def count(rows, heads, kv_len): calls.append((rows, heads, kv_len)) - return 2, split_indptr[: rows + 1] + return 3 - if has_planner: + if has_count: monkeypatch.setattr( - aiter.mla, "get_mla_v4_nm_split_plan", planner, raising=False + aiter.mla, "get_mla_v4_nm_num_kv_splits", count, 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) + monkeypatch.delattr(aiter.mla, "get_mla_v4_nm_num_kv_splits", raising=False) + table = paged_decode.v4_uniform_split_table(64, "cpu") + plan = paged_decode.v4_decode_split_plan(28, 128, 1152, table) - if has_planner: + if has_count: + # A row of the constant table, not a copy: nothing is written per call. assert calls == [(28, 128, 1152)] - assert plan[0] == 2 and plan[1].shape[0] == 29 + assert plan[0] == 3 and plan[1].data_ptr() == table[2].data_ptr() + assert plan[1][:29].tolist() == list(range(0, 3 * 29, 3)) else: assert plan is None +@pytest.mark.parametrize("splits,rows", [(17, 28), (3, 64)]) +def test_v4_decode_split_plan_outside_table_leaves_pick_to_aiter( + monkeypatch, splits, rows +): + import aiter.mla + + monkeypatch.setattr( + aiter.mla, + "get_mla_v4_nm_num_kv_splits", + lambda *a: splits, + raising=False, + ) + table = paged_decode.v4_uniform_split_table(63, "cpu") + assert paged_decode.v4_decode_split_plan(rows, 128, 1152, table) is None + + @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 @@ -221,14 +247,16 @@ def test_builder_decode_split_plan(monkeypatch, kv_fp8): monkeypatch.setattr( deepseek_v4_attn, "v4_decode_split_plan", - lambda rows, heads, kv_len, buf: calls.append((rows, heads, kv_len)) or "plan", + lambda rows, heads, kv_len, table: calls.append((rows, heads, kv_len, table)) + or "plan", ) - builder = SimpleNamespace(_kv_fp8=kv_fp8, _local_heads=128) + table = paged_decode.v4_uniform_split_table(64, "cpu") + builder = SimpleNamespace(_kv_fp8=kv_fp8, _local_heads=128, _split_table=table) plan = deepseek_v4_attn.DeepseekV4AttentionMetadataBuilder._decode_split_plan( - builder, 28, 1152, torch.zeros(65, dtype=torch.int32) + builder, 28, 1152 ) if kv_fp8: - assert plan == "plan" and calls == [(28, 128, 1152)] + assert plan == "plan" and calls == [(28, 128, 1152, table)] else: assert plan is None and calls == [] From 784eb259109f4bb070eaa4afdb9fcb7d2fc54b06 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Sun, 4 Oct 2026 00:10:24 -0500 Subject: [PATCH 12/12] [DSV4] Describe split plans as plain tuples in the docstrings aiter no longer has MlaV4NmSplitPlan; ATOM's plans were already (num_kv_splits, split_indptr) tuples from v4_decode_split_plan. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: yihonglie --- atom/model_ops/attentions/deepseek_v4_attn.py | 2 +- atom/model_ops/v4_kernels/paged_decode.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 5b29a163fa..79d3abc3a1 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -302,7 +302,7 @@ 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 - """aiter `MlaV4NmSplitPlan` (`num_kv_splits`, `split_indptr[padded_T+1]`) + """Split plan `(num_kv_splits, split_indptr)` (see `v4_decode_split_plan`) 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.""" diff --git a/atom/model_ops/v4_kernels/paged_decode.py b/atom/model_ops/v4_kernels/paged_decode.py index 48296802b4..161f8c181c 100644 --- a/atom/model_ops/v4_kernels/paged_decode.py +++ b/atom/model_ops/v4_kernels/paged_decode.py @@ -1199,8 +1199,8 @@ 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 + ``split_plan`` is ``(num_kv_splits, split_indptr)`` (:func:`v4_decode_split_plan`) + for the decode ASM kernel, built by whoever built this call's ``qo_indptr``; None leaves the split pick to aiter. ``compress_ratio`` is the calling layer's ratio (0 SWA, 4 CSA, 128 HCA).