diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 2e757f236f..79d3abc3a1 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,7 +120,10 @@ FP4_MQA_PARALLEL_UNIT_NUM, build_v4_paged_decode_indptr, fp4_indexer_enabled, + hca_persist, plan_context_lens, + v4_decode_split_plan, + v4_uniform_split_table, write_v4_paged_decode_indices, write_v4_paged_prefill_indices, ) @@ -297,6 +301,13 @@ class AttentionMetaData_DSV4(AttentionMetaData): """[padded_T+1] int32 GPU — packed cumsum of per-token HCA kv_len (= `min(positions[t]+1, win) + (positions[t]+1)//128`). Padded tail = last value.""" + split_plan_swa: tuple[int, torch.Tensor] | None = None + """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.""" + split_plan_hca: tuple[int, torch.Tensor] | None = None + """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 @@ -711,6 +722,12 @@ def __init__(self, model_runner): # Worst-case HCA per-token committed compress count # (= max_model_len // 128 for V4-Pro = 8192 at 1M context). self.max_committed_hca = model_runner.config.max_model_len // 128 + # Heads per rank the fp8 decode ASM kernel runs; sizes its split plans. + self._local_heads = ( + model_runner.config.hf_config.num_attention_heads + // get_tensor_model_parallel_world_size() + ) + 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 @@ -2589,6 +2606,9 @@ def prepare_mtp_decode( attn_metadata.max_seqlen_q = 1 attn_metadata.kv_indices_swa = swa_indices_buf attn_metadata.kv_indptr_swa = swa_indptr + attn_metadata.split_plan_swa = self._decode_split_plan( + running_bs, self.window_size + ) 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- @@ -3659,6 +3679,31 @@ 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 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 + ) -> tuple[int, torch.Tensor] | None: + """Split plan for the fp8 decode ASM kernel (see `v4_decode_split_plan`). + + 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, self._split_table) + def _attach_v4_paged_decode_meta( self, attn_metadata, @@ -3781,6 +3826,13 @@ def _attach_v4_paged_decode_meta( win=win, index_topk=index_topk, ) + # aiter decode split plans: sized by the rows this forward runs and the + # longest kv_len a token of each layer type can see. + attn_metadata.split_plan_swa = self._decode_split_plan(T_pad, win) + 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 + ) # Expand block tables per query row so the unchanged aiter paged-MQA # kernels can run once with shape `[decode_rows, 1, ...]`. Source and @@ -4604,6 +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) + # 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. diff --git a/atom/model_ops/v4_kernels/__init__.py b/atom/model_ops/v4_kernels/__init__.py index 753d2c38ae..7c2cab8f95 100644 --- a/atom/model_ops/v4_kernels/__init__.py +++ b/atom/model_ops/v4_kernels/__init__.py @@ -34,6 +34,8 @@ from atom.model_ops.v4_kernels.paged_decode import ( 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, @@ -89,6 +91,8 @@ "swa_write", "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 1ea4bdcc70..b1ad30f7bd 100644 --- a/atom/model_ops/v4_kernels/dspark_fp8_indices.py +++ b/atom/model_ops/v4_kernels/dspark_fp8_indices.py @@ -41,6 +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, + v4_uniform_split_table, +) from atom.model_ops.v4_kernels.pool_index import window_constexprs, window_row @@ -142,11 +146,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_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 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 +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_table=v4_uniform_split_table(n, device), batch_ids=torch.arange(n, **i32) // draft, max_batch=max_batch, draft_width=draft, @@ -199,11 +207,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 +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_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 new file mode 100644 index 0000000000..1073aa5e5f --- /dev/null +++ b/atom/model_ops/v4_kernels/hca_persist.py @@ -0,0 +1,217 @@ +# 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. +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 +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 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 ( + 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 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 + 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 6cb5a10a4b..161f8c181c 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, ) @@ -985,6 +986,46 @@ 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_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. 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 + + num_kv_splits = getattr(aiter.mla, "get_mla_v4_nm_num_kv_splits", None) + if num_kv_splits is None: + return None + 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( unified_kv: torch.Tensor, kv_indices: torch.Tensor, @@ -997,6 +1038,7 @@ def _sparse_attn_v4_paged_decode_asm( qo_indptr: torch.Tensor, kv_last_page_lens: torch.Tensor | None = None, # unused on v4 nm (page_size=1) num_kv_splits: int | None = None, + split_indptr: torch.Tensor | None = None, ) -> torch.Tensor: """Native 2buff fp8 V4 decode via the aiter assembly kernel ``mla_decode_fwd_v4_nm`` (ROCm/aiter#3112, mi350/gfx950). @@ -1106,6 +1148,10 @@ def _sparse_attn_v4_paged_decode_asm( (N, H, V4_DIM_NOPE + V4_DIM_ROPE), dtype=torch.bfloat16, device=device ) + # A uniform split plan's prefix is itself a valid plan, so it trims like + # qo_indptr. + if split_indptr is not None: + split_indptr = split_indptr[: N + 1] aiter.mla.mla_decode_fwd_v4_nm( q_packed, q_rope, @@ -1119,6 +1165,7 @@ def _sparse_attn_v4_paged_decode_asm( sink=attn_sink, sm_scale=softmax_scale, num_kv_splits=num_kv_splits, + split_indptr=split_indptr, ) # Drop padded heads. The slice is a non-contiguous view, so .contiguous() # gives downstream a dense tensor; no-op (and no copy) when H was unpadded. @@ -1141,6 +1188,8 @@ def sparse_attn_v4_paged_decode( kv_last_page_lens: torch.Tensor | None = None, empty_kv_indptr: torch.Tensor | None = None, prefix: str = "", + split_plan: tuple[int, torch.Tensor] | None = None, + compress_ratio: int | None = None, ) -> torch.Tensor: """V4 decode sparse attention over a unified KV pool with paged indices. @@ -1150,6 +1199,15 @@ 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 ``(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). + 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, @@ -1174,6 +1232,27 @@ 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) + and hca_persist.workspace_ready(q_packed_in.device) + ): + 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, kv_indices, @@ -1185,6 +1264,8 @@ def sparse_attn_v4_paged_decode( q_rope_in, qo_indptr=qo_indptr, kv_last_page_lens=kv_last_page_lens, + num_kv_splits=num_kv_splits, + split_indptr=split_indptr, ) gfx = get_gfx() if gfx == "gfx1250" or gfx.startswith("gfx94"): diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 26cca1c87b..76386baf8d 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -112,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, @@ -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=self.kv_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. @@ -3313,12 +3322,15 @@ def _sparse_attention( if ratio == 0: kv_indices = attn_md.kv_indices_swa kv_indptr = attn_md.kv_indptr_swa + split_plan = attn_md.split_plan_swa elif ratio == 4: kv_indices = attn_md.kv_indices_csa kv_indptr = attn_md.kv_indptr_csa + split_plan = attn_md.split_plan_csa else: # ratio == 128 kv_indices = attn_md.kv_indices_hca kv_indptr = attn_md.kv_indptr_hca + split_plan = attn_md.split_plan_hca # Dispatch on kv-cache layout inside the wrapper: fp8 2buff # (unified_kv_rope set) → aiter ASM with pre-packed fp8 Q + the # 2buff fp8/bf16 pools read with no requant. The optional H=128 @@ -3336,6 +3348,8 @@ def _sparse_attention( qo_indptr=attn_md.qo_indptr, empty_kv_indptr=attn_md.empty_kv_indptr, prefix=f"{self.layer_name}.sparse_attn_decode", + split_plan=split_plan, + compress_ratio=ratio, ) # [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/atom/utils/envs.py b/atom/utils/envs.py index 989cce9be5..a4dae33e4c 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -642,6 +642,15 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> 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 d8e7f91911..332eef3284 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -110,6 +110,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 (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 new file mode 100644 index 0000000000..ab0d88480e --- /dev/null +++ b/tests/model_ops/test_v4_hca_persist_gpu.py @@ -0,0 +1,307 @@ +# 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 + +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, +) + +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 { + "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, paged_decode.v4_uniform_split_table(MAX_ROWS_BUF, 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) # 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 = { + "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 + + +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()), + paged_decode.v4_uniform_split_table(MAX_ROWS_BUF, 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_deepseek_v4_transfer_regions.py b/tests/test_deepseek_v4_transfer_regions.py index 76b97fa4d6..59aa578913 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,25 @@ def _stub_v4_runtime_imports(): "fp4_indexer_enabled", "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", ): 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_hca_persist_dispatch.py b/tests/test_v4_hca_persist_dispatch.py new file mode 100644 index 0000000000..1c76ff519d --- /dev/null +++ b/tests/test_v4_hca_persist_dispatch.py @@ -0,0 +1,325 @@ +# 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()) + monkeypatch.setattr(hca_persist, "workspace_ready", lambda device: True) + 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) + + +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 + + +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() + ) diff --git a/tests/test_v4_prefill_asm_decode_dispatch.py b/tests/test_v4_prefill_asm_decode_dispatch.py index 85ad316742..5190c52fbf 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 their split plans.""" from __future__ import annotations +from types import SimpleNamespace + import pytest import torch @@ -111,3 +113,150 @@ def fake_prefill_asm(**kwargs): ) def test_ineligible_decode_keeps_dedicated_asm(monkeypatch, kwargs): assert _dispatch(monkeypatch, **kwargs) == "decode" + + +def _decode_rows(n: int, heads: int = 128) -> dict: + 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( + "split_plan,expected", + [ + ((4, torch.arange(0, 4 * 22, 4, dtype=torch.int32)), 4), + (None, None), + ], +) +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, + "_sparse_attn_v4_paged_decode_asm", + lambda *args, **kwargs: captured.update(kwargs) or "decode", + ) + n = 21 + result = paged_decode.sparse_attn_v4_paged_decode( + q=None, + qo_indptr=torch.arange(n + 1, dtype=torch.int32), + split_plan=split_plan, + **_decode_rows(n), + ) + + assert result == "decode" + assert captured["num_kv_splits"] == expected + assert captured["split_indptr"] is (split_plan[1] if split_plan else None) + + +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 + + 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] + + +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 count(rows, heads, kv_len): + calls.append((rows, heads, kv_len)) + return 3 + + if has_count: + monkeypatch.setattr( + aiter.mla, "get_mla_v4_nm_num_kv_splits", count, raising=False + ) + else: + 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_count: + # A row of the constant table, not a copy: nothing is written per call. + assert calls == [(28, 128, 1152)] + 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 + + calls = [] + monkeypatch.setattr( + deepseek_v4_attn, + "v4_decode_split_plan", + lambda rows, heads, kv_len, table: calls.append((rows, heads, kv_len, table)) + or "plan", + ) + 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 + ) + + if kv_fp8: + assert plan == "plan" and calls == [(28, 128, 1152, table)] + else: + assert plan is None and calls == []