[dsv4] kv cache fp8

atom/model_ops/v4_kernels/__init__.py

@@ -42,12 +42,18 @@
 from atom.model_ops.v4_kernels.qk_norm_rope_maybe_quant import (
     qk_norm_rope_maybe_quant,
     qk_norm_rope_maybe_quant_reference,
+    qk_norm_rope_maybe_quant_fp8_2buff,
+)
+from atom.model_ops.v4_kernels.state_writes import (
+    update_compressor_states,
+    swa_write,
+    swa_write_2buff_prepacked,
 )
-from atom.model_ops.v4_kernels.state_writes import update_compressor_states, swa_write
 
 __all__ = [
     "update_compressor_states",
     "swa_write",
+    "swa_write_2buff_prepacked",
     "fused_compress_attn",
     "fused_compress_attn_reference",
     "sparse_attn_v4_paged_decode",
@@ -66,4 +72,5 @@
     "write_v4_paged_prefill_indices_reference",
     "qk_norm_rope_maybe_quant",
     "qk_norm_rope_maybe_quant_reference",
+    "qk_norm_rope_maybe_quant_fp8_2buff",
 ]

atom/model_ops/v4_kernels/fused_compress.py

@@ -403,7 +403,15 @@ def fused_compress_attn(
     ] = None,  # fp32 [NB, k_per_block]; required when quant=True
     use_ue8m0: bool = True,  # round scale to power-of-2 (UE8M0); only when quant=True
     preshuffle: bool = True,  # MFMA 16x16 preshuffled FP8 layout; only when quant=True
-    fp8_max: Optional[float] = None,  # E4M3 max; required when quant=True
+    fp8_max: Optional[
+        float
+    ] = None,  # E4M3 max; required only on the quant (indexer_fp8) path
+    # V4-Main native fp8 2buff path (CSA/HCA Main under --kv_cache_dtype fp8).
+    # Distinct from `quant` (Indexer-inner per-row preshuffle): writes per-64-tile
+    # e8m0 nope-fp8 + inline dup-scale into `kv_cache` (fp8 [NB,k,512]) and bf16
+    # rope into `kv_cache_rope` (bf16 [NB,k,64]) via the flydsl group_fp8 scatter.
+    main_2buff_fp8: bool = False,
+    kv_cache_rope: Optional[torch.Tensor] = None,  # bf16 [NB,k_per_block,64]
 ) -> None:
     """Batched fused per-source-position pool + RMSNorm + RoPE + cache scatter,
     dispatched via SGLang-style packed plan.
@@ -462,6 +470,8 @@ def fused_compress_attn(
         and _shape_key == (512, 64, 128, False)
     )
     if _hca_use:
+        # main_2buff_fp8: native group_fp8 2buff scatter (nope-fp8 into
+        # kv_cache, bf16 rope into k_rope_cache). Otherwise plain bf16.
         flydsl_hca_compress_attn(
             kv_in=kv_in,
             score_in=score_in,
@@ -480,9 +490,15 @@ def fused_compress_attn(
             ratio=ratio,
             head_dim=head_dim,
             rope_head_dim=rope_head_dim,
+            quant=main_2buff_fp8,
+            k_rope_cache=kv_cache_rope if main_2buff_fp8 else None,
+            quant_group_size=64,
         )
         return
     if _flydsl_use:
+        # main_2buff_fp8: CSA Main native group_fp8 2buff (nope-fp8 + inline
+        # e8m0 into kv_cache, bf16 rope into k_rope_cache; scale carried inline
+        # so cache_scale stays None). Indexer-inner uses per_row_fp8 preshuffle.
         flydsl_fused_compress_attn(
             kv_in=kv_in,
             score_in=score_in,
@@ -502,10 +518,12 @@ def fused_compress_attn(
             ratio=ratio,
             head_dim=head_dim,
             rope_head_dim=rope_head_dim,
-            quant=quant,
+            quant=quant or main_2buff_fp8,
             cache_scale=cache_scale,
             use_ue8m0=use_ue8m0,
-            preshuffle=preshuffle,
+            preshuffle=preshuffle and not main_2buff_fp8,
+            quant_mode="group_fp8" if main_2buff_fp8 else "per_row_fp8",
+            k_rope_cache=kv_cache_rope if main_2buff_fp8 else None,
         )
         return
 

atom/model_ops/v4_kernels/paged_decode.py

@@ -904,6 +904,132 @@ def sparse_attn_v4_paged_decode_reference(
     return _sparse_attn_ragged_torch(q, unified_kv, attn_sink, topk_idxs, softmax_scale)
 
 
+def _sparse_attn_v4_paged_decode_asm(
+    unified_kv: torch.Tensor,
+    kv_indices: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    attn_sink: torch.Tensor,
+    softmax_scale: float,
+    unified_kv_rope: torch.Tensor,
+    q_packed_in: torch.Tensor,
+    q_rope_in: torch.Tensor,
+    qo_indptr: torch.Tensor,
+    kv_last_page_lens: torch.Tensor,
+    num_kv_splits: 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).
+
+    ``unified_kv`` is the packed 512-byte fp8 NoPE pool ``[P, 512]`` and
+    ``unified_kv_rope`` is the bf16 RoPE pool ``[P, 64]`` — consumed with NO
+    requant (caller stored the 2buff layout natively). Q is supplied pre-packed
+    (``q_packed_in``/``q_rope_in``) by the compute-only 2buff quant; no per-call
+    quantize.
+
+    The flat-CSR ``unified_kv`` is made compatible with the kernel's paged
+    format by treating each token as a 1-token page (``page_size=1``):
+    ``kv_page_indices = kv_indices``, ``kv_indptr`` is the existing CSR indptr,
+    ``kv_last_page_lens = ones(N)``, ``qo_indptr = arange(N+1)``,
+    ``max_seqlen_q = 1``. No translation layer needed.
+
+    softmax_scale is ignored by the kernel (hardcodes 1/sqrt(512)); passed
+    through for API parity.
+    """
+    import aiter
+    import aiter.mla
+
+    from atom.model_ops.v4_kernels.v4_quant import (
+        V4_DIM_NOPE,
+        V4_DIM_QK_PACKED,
+        V4_DIM_ROPE,
+    )
+
+    assert q_packed_in.dim() == 3 and q_packed_in.shape[-1] == V4_DIM_QK_PACKED, (
+        f"asm v4 nm: q_packed_in must be [N, H, {V4_DIM_QK_PACKED}] fp8, "
+        f"got {tuple(q_packed_in.shape)}"
+    )
+    assert q_rope_in.dim() == 3 and q_rope_in.shape[-1] == V4_DIM_ROPE, (
+        f"asm v4 nm: q_rope_in must be [N, H, {V4_DIM_ROPE}] bf16, "
+        f"got {tuple(q_rope_in.shape)}"
+    )
+    q_packed = q_packed_in
+    q_rope = q_rope_in
+    device = q_packed.device
+    N, H, _ = q_packed.shape
+
+    assert unified_kv.dim() == 2 and unified_kv.shape[-1] == V4_DIM_QK_PACKED, (
+        f"asm v4 nm: native fp8 unified_kv must be [P, {V4_DIM_QK_PACKED}] fp8, "
+        f"got {tuple(unified_kv.shape)}"
+    )
+    assert unified_kv_rope.dim() == 2 and unified_kv_rope.shape[-1] == V4_DIM_ROPE, (
+        f"asm v4 nm: native fp8 unified_kv_rope must be [P, {V4_DIM_ROPE}] bf16, "
+        f"got {tuple(unified_kv_rope.shape)}"
+    )
+
+    nhead_kv = 1
+    page_size = 1
+    # Kernel expects KV as [num_page, page_size=1, num_kv_heads=1, dim].
+    kv_packed = unified_kv.view(-1, page_size, nhead_kv, V4_DIM_QK_PACKED)
+    kv_rope = unified_kv_rope.view(-1, page_size, nhead_kv, V4_DIM_ROPE)
+
+    # ---- per-token paged layout (one slot per token, decode=1) ----------
+    # qo_indptr = arange(N+1) and kv_last_page_lens = ones(N) are per-token
+    # (max_seqlen_q=1, page_size=1) and depend ONLY on N. They are built by the
+    # attn_metadata builder (deepseek_v4_attn._attach_v4_paged_decode_meta /
+    # prepare_mtp_decode) via the SAME forward_vars staging path as kv_indptr /
+    # kv_indices — re-staged (H2D into a persistent CpuGpuBuffer) every fwd so
+    # the captured graph reads a freshly-copied backing store at replay — and
+    # threaded in through `sparse_attn_v4_paged_decode`.
+    # The per-seq index tensors (qo_indptr [num_seqs+1], kv_indptr [num_seqs+1],
+    # kv_last_page_lens [num_seqs]) are staged by the attn-metadata builder at the
+    # PADDED decode grid T_pad so one set of buffers serves both a CG-captured
+    # (padded) replay and an eager (real-T) forward. The asm kernel derives
+    # num_seqs from qo_indptr.numel()-1, so on an EAGER forward whose real token
+    # count N is smaller than the staged T_pad — e.g. an uncaptured MTP verify
+    # batch of bs*(k+1) tokens that is not a captured cudagraph size — the kernel
+    # would emit T_pad output rows while q_packed and the `output` buffer hold
+    # only N=q_packed.shape[0] rows. That is an OOB write into `output` on the
+    # split>1 path and a shape mismatch on the logits[:,0] path (o=[T_pad,...] vs
+    # q=[N,...]). Trim every per-seq tensor to the ACTUAL query count N: the real
+    # region is self-consistent (qo_indptr[:N+1] == arange(N+1), kv_indptr[:N+1]
+    # is the real ragged cumsum) and the dropped tail is exactly the 0-length
+    # CG-padded slots. No-op when captured/padded (N == T_pad).
+    qo_indptr = qo_indptr[: N + 1].to(torch.int32)
+    kv_last_page_lens = kv_last_page_lens[:N].to(torch.int32)
+    kv_indptr_i32 = kv_indptr[: N + 1].to(torch.int32)
+    kv_page_indices_i32 = kv_indices.to(torch.int32).contiguous()
+    max_seqlen_q = 1
+
+    output = torch.empty(
+        (N, H, V4_DIM_NOPE + V4_DIM_ROPE), dtype=torch.bfloat16, device=device
+    )
+    out_16_nosplit = 1 if num_kv_splits == 1 else 0
+
+    logits, _ = aiter.mla.mla_decode_fwd_v4_nm(
+        q_packed,
+        q_rope,
+        kv_packed,
+        kv_rope,
+        output,
+        qo_indptr,
+        kv_indptr_i32,
+        kv_page_indices_i32,
+        kv_last_page_lens,
+        max_seqlen_q,
+        sink=attn_sink,
+        sm_scale=softmax_scale,
+        out_16_nosplit=out_16_nosplit,
+        num_kv_splits=num_kv_splits,
+    )
+    # Result lands in `output` for the bf16-direct write (out_16_nosplit=1) or
+    # the stage2 LSE merge (resolved splits > 1). Only the single-pass fp32
+    # path leaves it in logits[:, 0]. logits.shape[1] = resolved split count.
+    resolved_splits = logits.shape[1]
+    if out_16_nosplit != 0 or resolved_splits > 1:
+        return output.to(torch.bfloat16)
+    return logits[:, 0].to(torch.bfloat16)
+
+
 def sparse_attn_v4_paged_decode(
     q: torch.Tensor,
     unified_kv: torch.Tensor,
@@ -912,12 +1038,36 @@ def sparse_attn_v4_paged_decode(
     attn_sink: torch.Tensor,
     softmax_scale: float,
     kv_scales: torch.Tensor | None = None,
+    unified_kv_rope: torch.Tensor | None = None,
+    q_packed_in: torch.Tensor | None = None,
+    q_rope_in: torch.Tensor | None = None,
+    qo_indptr: torch.Tensor | None = None,
+    kv_last_page_lens: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """V4 decode sparse attention over a unified KV pool with paged indices.
 
-    When ``kv_scales`` is provided, ``unified_kv`` must be fp8 (e4m3fnuz) and
-    will be dequantized in-kernel using 1xGROUP_SIZE (default 64) block scales.
+    Native 2buff fp8 (``unified_kv_rope`` provided): routes to the aiter asm
+    kernel (op5) with pre-packed fp8 Q (``q_packed_in``/``q_rope_in``); the
+    fp8 NoPE pool + bf16 RoPE pool are read with no requant.
+
+    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,
+    unreachable from the model).
     """
+    if unified_kv_rope is not None:
+        return _sparse_attn_v4_paged_decode_asm(
+            unified_kv,
+            kv_indices,
+            kv_indptr,
+            attn_sink,
+            softmax_scale,
+            unified_kv_rope,
+            q_packed_in,
+            q_rope_in,
+            qo_indptr=qo_indptr,
+            kv_last_page_lens=kv_last_page_lens,
+        )
     if get_gfx() == "gfx1250":
         return pa_decode_sparse(
             q,

atom/model_ops/v4_kernels/qk_norm_rope_maybe_quant.py

@@ -639,3 +639,60 @@ def _rope_tail(x: torch.Tensor) -> torch.Tensor:
         kv_scale = None
 
     return q_out, kv_out, q_scale, kv_scale
+
+
+def qk_norm_rope_maybe_quant_fp8_2buff(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    kv_weight: torch.Tensor,
+    cos_cache: torch.Tensor,
+    sin_cache: torch.Tensor,
+    positions: torch.Tensor,
+    n_local_heads: int,
+    head_dim: int,
+    rope_head_dim: int,
+    eps: float,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Compile-safe Triton fp8 2buff Q/K: per-head weightless Q RMSNorm +
+    weighted KV RMSNorm + GPT-J RoPE (bf16), then per-64-elt-tile e8m0 fp8 quant
+    of the NoPE half + 2buff pack (RoPE tail kept bf16).
+
+    Composed from two already-tuned Triton pieces (kept separate for reuse and
+    independent testability rather than fused into one launch):
+
+    1. :func:`qk_norm_rope_maybe_quant` (bf16, quant off) — per-head weightless
+       Q RMSNorm + weighted KV RMSNorm + GPT-J RoPE.
+    2. :func:`quantize_bf16_to_v4_2buff_triton` — per-64-elt-tile e8m0 fp8 quant
+       of the NoPE half + 2buff pack, RoPE tail kept bf16.
+
+    Compute-only: it does NOT perform a fused SWA scatter. Callers on the decode
+    path must additionally call ``swa_write_2buff_prepacked`` on the returned
+    ``k_packed``/``k_rope`` (the prefill call site scatters its tail post-attn).
+
+    Returns ``(q_packed [T, H, 512] fp8, q_rope [T, H, 64] bf16,
+    k_packed [T, 1, 512] fp8, k_rope [T, 1, 64] bf16)``.
+    """
+    from atom.model_ops.v4_kernels.v4_quant import quantize_bf16_to_v4_2buff_triton
+
+    # bf16 norm+rope (existing Triton/flydsl kernel, no quant, no SWA scatter).
+    q_bf16, kv_bf16, _, _ = qk_norm_rope_maybe_quant(
+        q,
+        kv,
+        kv_weight,
+        cos_cache,
+        sin_cache,
+        positions,
+        n_local_heads,
+        head_dim,
+        rope_head_dim,
+        eps,
+        quant_q=False,
+        quant_k=False,
+    )
+
+    # q_bf16: [T, H, D]; kv_bf16: [T, D] -> [T, 1, D] (single KV head).
+    T = q_bf16.shape[0]
+    q_packed, q_rope = quantize_bf16_to_v4_2buff_triton(q_bf16)
+    k_packed, k_rope = quantize_bf16_to_v4_2buff_triton(kv_bf16.view(T, 1, head_dim))
+
+    return q_packed, q_rope, k_packed, k_rope

atom/model_ops/v4_kernels/state_writes.py

@@ -222,6 +222,71 @@ def swa_write_reference(
         swa_region[dst_row] = src_kv
 
 
+def swa_write_2buff_prepacked(
+    k_packed: torch.Tensor,
+    k_rope: torch.Tensor,
+    positions: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    swa_block_tables: torch.Tensor,
+    swa_region_nope: torch.Tensor,
+    swa_region_rope: torch.Tensor,
+    block_size: int,
+    write_per_batch: int,
+) -> None:
+    """Native 2buff fp8 paged SWA write: content-addressed scatter of the LAST
+    ``min(tok_n_b, write_per_batch)`` tokens of every seq into the two paged
+    SWA pools (fp8 NoPE + bf16 RoPE). The K is ALREADY in the 2buff layout
+    (nope-fp8 ``[T,512]`` + rope-bf16 ``[T,64]``), produced upstream by the
+    compute-only 2buff quant (:func:`qk_norm_rope_maybe_quant_fp8_2buff`). This
+    is a pure dtype-agnostic scatter (reuses the paged :func:`swa_write` once per
+    pool); NO torch quantization happens here.
+
+    Both pools are the flat content-addressed regions of ``unified_kv`` /
+    ``unified_kv_rope`` (``[num_pages, D]``), addressed by ``swa_block_tables``:
+    ``swa_region[block_tables[b, pos//block_size] * block_size + pos%block_size]``.
+    Replaces the pre-paged per-request ring variant (matches the paged bf16
+    :func:`swa_write` semantics; issue #1417).
+
+    Args:
+        k_packed:        [T, 512] fp8 — quantized K nope + inline e8m0 scale + pad.
+        k_rope:          [T, 64]  bf16 — rotated K-PE (not quantized).
+        swa_block_tables:[bs, max_blocks] int32 — paged-SWA logical→physical map.
+        swa_region_nope: [num_pages, 512] fp8 paged pool (2buff nope).
+        swa_region_rope: [num_pages, 64]  bf16 paged pool (rope).
+        block_size:      paging stride of both pools.
+        (other args as :func:`swa_write`.)
+    """
+    from atom.model_ops.v4_kernels.v4_quant import V4_DIM_QK_PACKED, V4_DIM_ROPE
+
+    assert (
+        k_packed.dim() == 2 and k_packed.shape[1] == V4_DIM_QK_PACKED
+    ), f"k_packed must be [T,{V4_DIM_QK_PACKED}] fp8, got {tuple(k_packed.shape)}"
+    assert (
+        k_rope.dim() == 2 and k_rope.shape[1] == V4_DIM_ROPE
+    ), f"k_rope must be [T,{V4_DIM_ROPE}] bf16, got {tuple(k_rope.shape)}"
+    assert swa_region_nope.dim() == 2 and swa_region_nope.shape[1] == V4_DIM_QK_PACKED
+    assert swa_region_rope.dim() == 2 and swa_region_rope.shape[1] == V4_DIM_ROPE
+
+    swa_write(
+        k_packed.contiguous(),
+        positions,
+        cu_seqlens_q,
+        swa_block_tables,
+        swa_region_nope,
+        block_size,
+        write_per_batch,
+    )
+    swa_write(
+        k_rope.contiguous(),
+        positions,
+        cu_seqlens_q,
+        swa_block_tables,
+        swa_region_rope,
+        block_size,
+        write_per_batch,
+    )
+
+
 # === Unified Compressor state save (plan path) ==========================
 # Paper §3.6.1: per-request fixed-size state cache for "uncompressed tail
 # tokens + previous block as overlap context (B-side, eq 11)". ATOM keeps