diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index f37da9b23..81e8ffcce 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -35,6 +35,7 @@ Total = ~26.5 MB / slot """ +import logging import os from dataclasses import dataclass from typing import Any, Dict, Optional, Type, cast @@ -42,6 +43,7 @@ import numpy as np import torch from aiter import dtypes +from aiter.jit.utils.chip_info import get_gfx from atom.distributed.pcp_utils import ( get_pcp_world_size, pcp_is_enabled, @@ -69,6 +71,8 @@ get_forward_context, ) +logger = logging.getLogger("atom") + # --------------------------------------------------------------------------- # Typed metadata surface for V4. The base AttentionMetaData class is shared # across all backends; carrying V4-specific dynamic attributes there would @@ -180,6 +184,21 @@ class AttentionMetaData_DSV4(AttentionMetaData): for the independent SWA pool (parallel to `block_tables`, which addresses the compressed pool). -1 entries are window-freed blocks (never indexed).""" + # ----- Native 2buff fp8 per-token paged-decode index tensors ----- + # Feed the aiter asm decode kernel `mla_decode_fwd_v4_nm` (op5), which treats + # each decode token as a 1-token page (page_size=1). Both depend ONLY on the + # padded decode token count N (the captured kernel grid), never on batch + # content — the values are always arange(N+1) / ones(N). Staged every fwd via + # the SAME forward_vars path as `kv_indptr_*` (CpuGpuBuffer H2D), which is + # what makes them CUDAGraph-safe. Only populated on the fp8 path. + qo_indptr: Optional[torch.Tensor] = None + """[padded_T+1] int32 GPU — per-token q indptr `arange(N+1)` (page_size=1, + max_seqlen_q=1). NOT `cu_seqlens_q` (which is per-seq and differs under + MTP); this is the per-token indptr the decode kernel consumes.""" + kv_last_page_lens: Optional[torch.Tensor] = None + """[padded_T] int32 GPU — per-token last-page length `ones(N)` (page_size=1 + → every page is full).""" + # ----- Indexer / sparse-layout side metadata ----- indexer_meta: Optional[Dict[str, Any]] = None """dict — `Indexer.forward_batched` per-fwd GPU tensors. Notable keys: @@ -307,6 +326,7 @@ def __init__(self, model_runner): self.index_head_dim = getattr(hf, "index_head_dim", 128) self.window_size = getattr(hf, "sliding_window", 128) self.index_topk = getattr(hf, "index_topk", 1024) + self.rope_head_dim = getattr(hf, "qk_rope_head_dim", 64) # MTP-portion of compress_ratios. `prepare_mtp_decode`'s direct-kernel # fast path only handles SWA (ratio=0) draft layers; non-zero ratios # would also need n_committed_{csa,hca} + HCA compress tail rebuilt. @@ -327,8 +347,39 @@ def __init__(self, model_runner): self.k2_hca = self.block_size // 128 # = 1 self._state_dtype = torch.float32 # fp32 required for softmax-pool - self._swa_dtype = torch.bfloat16 # SWA window matches KV dtype - self._classical_dtype = torch.bfloat16 # CSA Main / HCA Main KV is BF16 + # KV cache dtype gate. fp8 → 2buff native layout (nope fp8 in a 512B + # entry with inline e8m0 scale; parallel bf16 rope pool). bf16 → + # unchanged. SWA and classical (CSA/HCA Main) share the nope dtype; the + # rope pool is always bf16. + self._kv_fp8 = model_runner.kv_cache_dtype == "fp8" + # aiter prefill (op4) / decode (op5) implement the fp8 (2buff) path only + # on gfx950 / gfx1250. On any other arch, transparently fall back to a + # bf16 KV cache instead of hard-failing. Flipping self._kv_fp8 here (before + # the *_dtype attrs are read) keeps the whole V4 path consistent: pool + # sizing (compute_block_bytes / swa_block_bytes_per_layer), write_mode, + # and module.kv_fp8 (build_kv_cache_tensor) all key off self._kv_fp8 / + # these dtype attrs. Sync model_runner.kv_cache_dtype (and the shared + # config) so any generic reader / log line agrees. + if self._kv_fp8 and get_gfx() not in ("gfx950", "gfx1250"): + logger.warning( + "DeepSeek-V4 --kv_cache_dtype fp8 (2buff) is only supported on " + "gfx950 / gfx1250 (aiter op4/op5); got %r. Falling back to a " + "bf16 KV cache.", + get_gfx(), + ) + self._kv_fp8 = False + model_runner.kv_cache_dtype = "bf16" + cfg = getattr(model_runner, "config", None) + if cfg is not None and getattr(cfg, "kv_cache_dtype", None) == "fp8": + cfg.kv_cache_dtype = "bf16" + if self._kv_fp8: + self._swa_dtype = dtypes.fp8 + self._classical_dtype = dtypes.fp8 + self._rope_dtype = torch.bfloat16 # rope pool is always bf16 + else: + self._swa_dtype = torch.bfloat16 # SWA window matches KV dtype + self._classical_dtype = torch.bfloat16 # CSA / HCA Main KV is BF16 + self._rope_dtype = torch.bfloat16 # unused in bf16 path (symmetry) # CSA Indexer cache is FP8 + 4-byte fp32 scale per row, aligned to 16 # bytes (matches V3.2 sparse MLA pattern; avoids torch inductor # unaligned-access slowdowns). Written by `indexer_k_quant_and_cache`, @@ -437,7 +488,11 @@ def swa_block_bytes_per_layer(self) -> int: (full-resolution, ratio-1 = block_size tokens x head_dim x classical elem). Single source for the per-layer SWA block size, reused by both `swa_pool_block_bytes` and the KV-transfer region stride.""" - return self.block_size * self.head_dim * self._classical_dtype.itemsize + b = self.block_size * self.head_dim * self._swa_dtype.itemsize + if self._kv_fp8: + # 2buff: parallel bf16 rope pool [block_size, rope_head_dim]. + b += self.block_size * self.rope_head_dim * self._rope_dtype.itemsize + return b def swa_pool_block_bytes(self) -> int: """paged-SWA: bytes of ONE SWA physical block across all layers @@ -505,16 +560,22 @@ def compute_block_bytes(self) -> int: read by `cp_gather_indexer_k_quant_cache`). - HCA Main: k2=1 entry × head_dim BF16 """ - elem_classical = self._classical_dtype.itemsize + elem_classical = self._classical_dtype.itemsize # fp8 = 1 or bf16 = 2 csa_main_per_block = self.k1_csa * self.head_dim * elem_classical csa_idx_per_block = self.k1_csa * self._aligned_index_dim # fp8 = 1B hca_main_per_block = self.k2_hca * self.head_dim * elem_classical # paged-SWA: the sliding-window KV is content-addressed, one full- # resolution (ratio-1) entry per original token in EVERY layer, so each - # block carries `block_size * head_dim` BF16 of SWA per layer. This term + # block carries `block_size * head_dim` of SWA per layer. This term # is charged here but the budget (model_runner.get_num_blocks) strips it # back out into the separate window-freed num_swa_blocks pool. swa_per_block = self.block_size * self.head_dim * elem_classical + if self._kv_fp8: + # 2buff: parallel bf16 rope pool per compress entry AND per SWA token. + elem_rope = self._rope_dtype.itemsize + csa_main_per_block += self.k1_csa * self.rope_head_dim * elem_rope + hca_main_per_block += self.k2_hca * self.rope_head_dim * elem_rope + swa_per_block += self.block_size * self.rope_head_dim * elem_rope return ( len(self.csa_layers) * (csa_main_per_block + csa_idx_per_block) + len(self.hca_layers) * hca_main_per_block @@ -577,7 +638,9 @@ def allocate_per_req_cache(self, num_slots: int) -> dict[str, object]: """ assert self._swa_dtype == self._classical_dtype, ( "unified_kv requires SWA dtype == classical KV dtype " - f"(got SWA={self._swa_dtype}, classical={self._classical_dtype})" + f"(got SWA={self._swa_dtype}, classical={self._classical_dtype}). " + "fp8 path must set both to dtypes.fp8 (rope lives in a separate " + "bf16 pool); a genuine mismatch corrupts the unified layout." ) device = self.model_runner.device num_blocks = self.model_runner.num_physical_kvcache_blocks @@ -608,6 +671,30 @@ def allocate_per_req_cache(self, num_slots: int) -> dict[str, object]: ) ) + # ---- 2buff fp8: parallel per-layer rope pool (bf16) ------------------ + # Same [swa_pages + compress_pages] page count as unified_kv, but width + # = rope_head_dim (64) and dtype bf16 (rope is never quantized). bf16 + # path: list of None (no rope pool; rope stays inline in unified_kv). + unified_kv_rope: list[Optional[torch.Tensor]] = [] + if self._kv_fp8: + for layer_id in range(self.num_layers): + ratio = ratios[layer_id] + if ratio == 4: + compress_pages = num_blocks * self.k1_csa + elif ratio == 128: + compress_pages = num_blocks * self.k2_hca + else: + compress_pages = 0 # Dense + unified_kv_rope.append( + torch.zeros( + (swa_pages + compress_pages, self.rope_head_dim), + dtype=self._rope_dtype, + device=device, + ) + ) + else: + unified_kv_rope = [None] * self.num_layers + # ---- Compressor state tensors (compute-contiguous) ------------------ csa_main_kv = self._zero_state( (n_csa, num_slots, *self.csa_main_state_shape), device @@ -652,6 +739,7 @@ def allocate_per_req_cache(self, num_slots: int) -> dict[str, object]: return { "v4_unified_kv": unified_kv, + "v4_unified_kv_rope": unified_kv_rope, "v4_csa_main_kv_state": csa_main_kv, "v4_csa_main_score_state": csa_main_score, "v4_csa_idx_kv_state": csa_idx_kv, @@ -698,6 +786,17 @@ def build_kv_cache_tensor(self, layer_id: int, module): module.unified_kv = unified module.swa_kv = unified[:swa_pages] module.swa_block_size = self.block_size + module.kv_fp8 = self._kv_fp8 + if self._kv_fp8: + # 2buff: parallel bf16 rope pool, same paged layout. swa_kv_rope + # is the flat [swa_pages, rope_head_dim] SWA region; unified_kv_rope + # the full pool (asm decode op5 reads it). + rope = runner.v4_unified_kv_rope[module.layer_id] + module.unified_kv_rope = rope + module.swa_kv_rope = rope[:swa_pages] + else: + module.unified_kv_rope = None + module.swa_kv_rope = None return None if isinstance(module, _V4Indexer): @@ -763,6 +862,10 @@ def build_kv_cache_tensor(self, layer_id: int, module): stride=(block_fp32_stride, 1), storage_offset=idx_kv_f32.storage_offset() + scale_fp32_offset, ) + # Indexer-inner cache is always fp8 (independent of + # kv_cache_dtype); it has no separate rope pool. + module.write_mode = "indexer_fp8" + module.kv_cache_rope = None elif ratio == 4: pos = self.layer_id_to_csa_pos[layer_id_from_prefix] module.kv_state = runner.v4_csa_main_kv_state[pos] @@ -776,6 +879,15 @@ def build_kv_cache_tensor(self, layer_id: int, module): module.kv_cache = unified[swa_pages:].view( num_blocks, self.k1_csa, self.head_dim ) + if self._kv_fp8: + rope = runner.v4_unified_kv_rope[layer_id_from_prefix] + module.kv_cache_rope = rope[swa_pages:].view( + num_blocks, self.k1_csa, self.rope_head_dim + ) + module.write_mode = "main_2buff_fp8" + else: + module.kv_cache_rope = None + module.write_mode = "bf16" elif ratio == 128: pos = self.layer_id_to_hca_pos[layer_id_from_prefix] module.kv_state = runner.v4_hca_main_kv_state[pos] @@ -785,6 +897,15 @@ def build_kv_cache_tensor(self, layer_id: int, module): module.kv_cache = unified[swa_pages:].view( num_blocks, self.k2_hca, self.head_dim ) + if self._kv_fp8: + rope = runner.v4_unified_kv_rope[layer_id_from_prefix] + module.kv_cache_rope = rope[swa_pages:].view( + num_blocks, self.k2_hca, self.rope_head_dim + ) + module.write_mode = "main_2buff_fp8" + else: + module.kv_cache_rope = None + module.write_mode = "bf16" else: raise ValueError( f"Unknown V4 compress_ratio={ratio} on Compressor at " @@ -803,6 +924,11 @@ def get_kv_transfer_tensors(self): runner = self.model_runner if not hasattr(runner, "v4_unified_kv"): return None + if self._kv_fp8: + # PD disaggregation with 2buff fp8 KV cache is not yet supported: + # the byte-region math below assumes a single unified pool and + # ignores the parallel rope pool. Disable KV transfer for fp8. + return None num_slots = runner.max_per_req_cache_slots # paged-SWA: SWA lives in a SEPARATE num_swa_blocks pool at the head @@ -1157,6 +1283,17 @@ def prepare_mtp_decode( attn_metadata.kv_indptr_swa = swa_indptr attn_metadata.batch_id_per_token = batch_id_per_token + # fp8 asm decode per-token index tensors. MTP draft step is 1-token-per- + # seq → the asm kernel sees N = bs. Stage the constant per-token tensors + # to that length via the same builder-staged path as the verify fwd. + if self._kv_fp8: + attn_metadata.qo_indptr = self._stage( + "v4_qo_indptr", self._v4_qo_indptr_np[: bs + 1] + ) + attn_metadata.kv_last_page_lens = self._stage( + "v4_kv_last_page_lens", self._v4_kv_last_page_lens_np[:bs] + ) + # NOT rebuilt (unused by SWA-only MTP layer; would block a future # CSA/HCA MTP layer — assert at top guards): # - n_committed_{csa,hca}_per_seq{,_cpu} (compressor/HCA tail math) @@ -2206,6 +2343,26 @@ def _attach_v4_paged_decode_meta( attn_metadata.kv_indptr_hca = hca_indptr_gpu attn_metadata.swa_pages = swa_pages + # Per-token paged-decode index tensors for the fp8 asm decode kernel. The + # kernel sees N = q_packed.shape[0] = T_pad (padded decode grid). Both + # are re-staged every fwd (like kv_indptr_*) so the captured graph sees a + # freshly-copied backing store at replay. + # qo_indptr: per-token q indptr (page_size=1, max_seqlen_q=1). The REAL + # region [0..T] is arange(T+1) — one 1-length query per real decode token. + # The CG-padded tail [T+1..T_pad] must NOT keep counting up: repeating + # the last real value makes each padded slot a 0-length query + # (qo_indptr[t+1]-qo_indptr[t]==0) that the asm kernel bails on, exactly + # like the kv_indptr pad tail. Per-token, so correct for MTP too. + if self._kv_fp8: + qo_indptr_np = np.empty(T_pad + 1, dtype=np.int32) + qo_indptr_np[: T + 1] = np.arange(T + 1, dtype=np.int32) + if T_pad > T: + qo_indptr_np[T + 1 :] = T + attn_metadata.qo_indptr = self._stage("v4_qo_indptr", qo_indptr_np) + attn_metadata.kv_last_page_lens = self._stage( + "v4_kv_last_page_lens", self._v4_kv_last_page_lens_np[:T_pad] + ) + def _build_paged_prefill_meta( self, attn_metadata: AttentionMetaData_DSV4, @@ -2754,6 +2911,20 @@ def _alloc_v4_metadata_buffers(self) -> None: bufs["v4_kv_indptr_swa"] = CpuGpuBuffer(T_dec + 1, **i32) bufs["v4_kv_indptr_csa"] = CpuGpuBuffer(T_dec + 1, **i32) bufs["v4_kv_indptr_hca"] = CpuGpuBuffer(T_dec + 1, **i32) + + # Per-token paged-decode index tensors for the fp8 asm decode kernel + # (`mla_decode_fwd_v4_nm`, page_size=1). Values are CONSTANT — they + # depend only on the (padded) decode token count N, not the batch: + # qo_indptr = arange(N+1) (per-token q indptr, max_seqlen_q=1) + # kv_last_page_lens = ones(N) (page_size=1 → every page full) + # Built the SAME way as `kv_indptr_*`: a CpuGpuBuffer re-staged via + # `self._stage(...)` EVERY fwd, which is what makes them CUDAGraph-safe + # (re-copied into the captured buffer before graph.replay). The constant + # numpy sources are precomputed once so the per-fwd cost is a slice + H2D. + bufs["v4_qo_indptr"] = CpuGpuBuffer(T_dec + 1, **i32) + bufs["v4_kv_last_page_lens"] = CpuGpuBuffer(T_dec, **i32) + self._v4_qo_indptr_np = np.arange(T_dec + 1, dtype=np.int32) + self._v4_kv_last_page_lens_np = np.ones(T_dec, dtype=np.int32) # Per-seq `ctx_len // 4` (raw, no clamp). Consumed by csa_translate_pack # (kernel masks `(k < n_committed) & (k < index_topk)`) AND by the # indexer (cast to int64 inline). Built unconditionally in diff --git a/atom/model_ops/v4_kernels/__init__.py b/atom/model_ops/v4_kernels/__init__.py index 0b86fff07..d6ac0d3b1 100644 --- a/atom/model_ops/v4_kernels/__init__.py +++ b/atom/model_ops/v4_kernels/__init__.py @@ -40,14 +40,21 @@ write_v4_paged_prefill_indices_reference, ) from atom.model_ops.v4_kernels.qk_norm_rope_maybe_quant import ( + QKNormRopeOut, 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", @@ -64,6 +71,8 @@ "write_v4_paged_decode_indices_reference", "write_v4_paged_prefill_indices", "write_v4_paged_prefill_indices_reference", + "QKNormRopeOut", "qk_norm_rope_maybe_quant", "qk_norm_rope_maybe_quant_reference", + "qk_norm_rope_maybe_quant_fp8_2buff", ] diff --git a/atom/model_ops/v4_kernels/fused_compress.py b/atom/model_ops/v4_kernels/fused_compress.py index 38a68f19f..d3cab55ef 100644 --- a/atom/model_ops/v4_kernels/fused_compress.py +++ b/atom/model_ops/v4_kernels/fused_compress.py @@ -406,7 +406,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] prefix: str = "", ) -> None: """Batched fused per-source-position pool + RMSNorm + RoPE + cache scatter, @@ -466,6 +474,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, @@ -484,9 +494,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, @@ -506,10 +522,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 diff --git a/atom/model_ops/v4_kernels/paged_decode.py b/atom/model_ops/v4_kernels/paged_decode.py index 50d9a03dd..dd5efc3fe 100644 --- a/atom/model_ops/v4_kernels/paged_decode.py +++ b/atom/model_ops/v4_kernels/paged_decode.py @@ -906,6 +906,135 @@ 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). + # These are all staged as int32 by the metadata builder, so no dtype cast is + # needed — only the trim to the real query count N (and a contiguity fixup on + # the possibly-viewed kv_indices). + qo_indptr = qo_indptr[: N + 1] + kv_last_page_lens = kv_last_page_lens[:N] + kv_indptr_i32 = kv_indptr[: N + 1] + kv_page_indices_i32 = kv_indices.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 + return logits[:, 0].to(torch.bfloat16) + + @mark_trace def sparse_attn_v4_paged_decode( q: torch.Tensor, @@ -915,13 +1044,37 @@ 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, prefix: str = "", ) -> 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, + ) gfx = get_gfx() if gfx == "gfx1250" or gfx.startswith("gfx94"): return pa_decode_sparse( diff --git a/atom/model_ops/v4_kernels/paged_prefill.py b/atom/model_ops/v4_kernels/paged_prefill.py index 6acd09549..ea688beb4 100644 --- a/atom/model_ops/v4_kernels/paged_prefill.py +++ b/atom/model_ops/v4_kernels/paged_prefill.py @@ -585,30 +585,73 @@ def sparse_attn_v4_paged_prefill( attn_sink: torch.Tensor, softmax_scale: float, out: torch.Tensor | None = None, + *, + unified_kv_rope: torch.Tensor | None = None, + q_packed: torch.Tensor | None = None, + q_rope: torch.Tensor | None = None, + k_packed: torch.Tensor | None = None, + k_rope: torch.Tensor | None = None, prefix: str = "", ) -> torch.Tensor: """V4 prefill sparse attention over two KV sources (paged unified_kv + - flat per-fwd kv). + flat per-fwd kv), dispatching on the kv-cache layout. + + Native 2buff fp8 (``unified_kv_rope`` provided): routes to aiter's + ``pa_sparse_prefill_fp8_opus`` (op4). ``unified_kv`` is the packed fp8 NoPE + prefix pool and ``unified_kv_rope`` the bf16 RoPE prefix pool; the + op-quantized fp8 Q (``q_packed`` / ``q_rope``) and extend K + (``k_packed`` / ``k_rope``) are fed directly, no dequant of the prefix and no + torch quant. gfx950/gfx1250. (fp8 + PCP is out of scope: the extend K is this + fwd's shard and is not PCP all-gathered — the caller must pass the sharded + ``k_packed`` / ``k_rope``.) + + Otherwise (bf16): the existing OPUS / Triton / reference path over ``q`` and + the bf16 extend ``kv``. Args: - q: [T, H, D] BF16/FP16 — query. - unified_kv: [total_pages, D] BF16/FP16 — prefix source (paged). + q: [T, H, D] BF16/FP16 — query (bf16 path). + unified_kv: [total_pages, D] — prefix source (paged). BF16/FP16 on + the bf16 path; packed fp8 NoPE pool on the fp8 2buff path. kv_indices_prefix: [total_prefix] int32 — flat per-token slot lists into unified_kv. -1 sentinels skipped. kv_indptr_prefix: [T+1] int32 — true prefix sum. - kv: [total_tokens, D] BF16/FP16 — extend source (this - fwd's input K, NOT yet in swa_kv ring). + kv: [total_tokens, D] BF16/FP16 — extend source (bf16 path; + this fwd's input K, NOT yet in swa_kv ring). kv_indices_extend: [total_extend] int32 — flat per-token row idx lists into kv. -1 sentinels skipped. kv_indptr_extend: [T+1] int32 — true prefix sum. attn_sink: [H] — per-head softmax-denom bias. softmax_scale: float. out: optional output buffer. Can alias q when the caller no - longer needs q after attention. + longer needs q after attention (bf16 path only). + unified_kv_rope: [total_pages, rd] bf16 RoPE prefix pool — presence + selects the native fp8 2buff path. + q_packed / q_rope: fp8 2buff Q ([T, H, 512] fp8 / [T, H, 64] bf16). + k_packed / k_rope: fp8 2buff extend K ([T, 512] fp8 / [T, 64] bf16, or the + [T, 1, *] views produced by the quant kernel — flattened here). Returns: - out: [T, H, D] same dtype as q. + out: [T, H, D] bf16. """ + if unified_kv_rope is not None: + # Native fp8 prefill (op4): 2buff fp8 prefix pool + op-quantized fp8 Q + # and extend K fed directly. No dequant of the prefix, no torch quant. + from aiter.ops.pa_sparse_prefill_opus import pa_sparse_prefill_fp8_opus + + return pa_sparse_prefill_fp8_opus( + q_packed, + q_rope, + unified_kv, + unified_kv_rope, + kv_indices_prefix, + kv_indptr_prefix, + k_packed.view(k_packed.shape[0], -1), + k_rope.view(k_rope.shape[0], -1), + kv_indices_extend, + kv_indptr_extend, + attn_sink, + softmax_scale, + ) # [S, H, head_dim] bf16 if get_gfx() == "gfx1250": return pa_prefill_sparse( q, diff --git a/atom/model_ops/v4_kernels/qk_norm_rope_maybe_quant.py b/atom/model_ops/v4_kernels/qk_norm_rope_maybe_quant.py index 9579c8821..ec6fbac74 100644 --- a/atom/model_ops/v4_kernels/qk_norm_rope_maybe_quant.py +++ b/atom/model_ops/v4_kernels/qk_norm_rope_maybe_quant.py @@ -29,6 +29,7 @@ ops anyway. """ +from dataclasses import dataclass from typing import Optional, Tuple import torch @@ -38,6 +39,40 @@ from atom.model_ops.v4_kernels.state_writes import swa_write from atom.utils.decorators import mark_trace + +@dataclass +class QKNormRopeOut: + """Unified carrier for :func:`qk_norm_rope_maybe_quant`. + + The bf16 and fp8-2buff paths return genuinely different tensors, so both + are named fields on one struct (the inactive path's fields stay ``None``). + Downstream ``sparse_attn_v4_paged_{decode,prefill}`` / ``swa_write`` read + the fields they need and dispatch on the kv-cache layout themselves — the + model no longer branches on ``kv_fp8`` to pick tensors. + + bf16 path (fp8 fields None): + - ``q_sa`` [T, H, D] bf16 — post norm+rope Q (a.k.a. the old ``q_out``). + - ``kv`` [T, D] bf16 — post norm+rope KV. + - ``q_scale`` [T, H] fp32 — only when per-row ``quant_q``, else None. + - ``kv_scale``[T] fp32 — only when per-row ``quant_k``, else None. + + fp8 2buff path (bf16 fields None): + - ``q_packed`` [T, H, 512] fp8 — NoPE fp8 + inline e8m0 scale + pad. + - ``q_rope`` [T, H, 64] bf16 — rotated Q-PE (not quantized). + - ``k_packed`` [T, 1, 512] fp8 — NoPE fp8 (single MQA KV head). + - ``k_rope`` [T, 1, 64] bf16 — rotated K-PE. + """ + + q_sa: Optional[torch.Tensor] = None + kv: Optional[torch.Tensor] = None + q_scale: Optional[torch.Tensor] = None + kv_scale: Optional[torch.Tensor] = None + q_packed: Optional[torch.Tensor] = None + q_rope: Optional[torch.Tensor] = None + k_packed: Optional[torch.Tensor] = None + k_rope: Optional[torch.Tensor] = None + + # Lazy-imported flydsl path (optional dependency). Set to None when flydsl # is unavailable; the dispatch in ``qk_norm_rope_maybe_quant`` will fall # back to the Triton kernel. @@ -306,8 +341,7 @@ def _qk_norm_rope_maybe_quant_kernel( tl.store(kv_out_base + NOPE + rd_offs[None, :], pe.to(ot), mask=m_mask[:, None]) -@mark_trace -def qk_norm_rope_maybe_quant( +def _qk_norm_rope_maybe_quant_bf16( q: torch.Tensor, kv: torch.Tensor, kv_weight: torch.Tensor, @@ -560,6 +594,125 @@ def qk_norm_rope_maybe_quant( return q_out, kv_out, q_scale, kv_scale +@mark_trace +def qk_norm_rope_maybe_quant( + 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, + quant_q: bool = False, + quant_k: bool = False, + swa_kv: Optional[torch.Tensor] = None, + state_slot_mapping: Optional[torch.Tensor] = None, + batch_id_per_token: Optional[torch.Tensor] = None, + swa_cu_seqlens_q: Optional[torch.Tensor] = None, + swa_cache_size: Optional[int] = None, + swa_write_per_batch: Optional[int] = None, + swa_block_tables: Optional[torch.Tensor] = None, + swa_block_size: Optional[int] = None, + prefix: str = "", + *, + fp8_2buff: bool = False, + swa_nope_scale_buff: Optional[torch.Tensor] = None, + swa_rope_buff: Optional[torch.Tensor] = None, +) -> QKNormRopeOut: + """Per-token RMSNorm + GPT-J RoPE, dispatching on the kv-cache layout. + + This is the single entry the V4 model calls; it picks the kernel from the + kv-cache dtype so the model no longer branches on ``kv_fp8``: + + - ``fp8_2buff=False`` (bf16 kv cache): the existing fused Triton / flydsl + per-token RMSNorm + RoPE (+ optional per-row FP8 quant + fused bf16 SWA + scatter). All the ``swa_*`` / ``quant_*`` args behave exactly as before. + Returns a :class:`QKNormRopeOut` with ``q_sa`` / ``kv`` (/ ``q_scale`` / + ``kv_scale``) populated. + + - ``fp8_2buff=True`` (fp8 kv cache): dispatches to aiter's + ``fused_qk_norm_rope_group_quant`` — per-head weightless Q RMSNorm + + weighted KV RMSNorm + GPT-J RoPE + 1x64 e8m0 fp8 group-quant into the + native 2buff layout (NoPE-fp8 [.,512] + RoPE-bf16 [.,64]) consumed by op4 + (prefill) / op5 (decode) with no requant. The decode path additionally + fuses the paged SWA scatter into the same launch via ``swa_nope_scale_buff`` + / ``swa_rope_buff`` + ``swa_block_tables`` / ``swa_block_size`` / + ``batch_id_per_token`` (pass ``None`` for prefill, which scatters its + window tail post-attention). Returns a :class:`QKNormRopeOut` with + ``q_packed`` / ``q_rope`` / ``k_packed`` / ``k_rope`` populated. + + See :func:`_qk_norm_rope_maybe_quant_bf16` for the bf16 arg contract. + """ + if not fp8_2buff: + q_out, kv_out, q_scale, kv_scale = _qk_norm_rope_maybe_quant_bf16( + q, + kv, + kv_weight, + cos_cache, + sin_cache, + positions, + n_local_heads, + head_dim, + rope_head_dim, + eps, + quant_q=quant_q, + quant_k=quant_k, + swa_kv=swa_kv, + state_slot_mapping=state_slot_mapping, + batch_id_per_token=batch_id_per_token, + swa_cu_seqlens_q=swa_cu_seqlens_q, + swa_cache_size=swa_cache_size, + swa_write_per_batch=swa_write_per_batch, + swa_block_tables=swa_block_tables, + swa_block_size=swa_block_size, + prefix=prefix, + ) + return QKNormRopeOut(q_sa=q_out, kv=kv_out, q_scale=q_scale, kv_scale=kv_scale) + + # ---- fp8 native 2buff path (aiter group-quant, moved here from the model). + # Single fused launch: per-head weightless Q RMSNorm + weighted KV RMSNorm + + # GPT-J RoPE + 1x64 e8m0 fp8 group-quant into the 2buff layout + # (NoPE-fp8 [.,512] + RoPE-bf16 [.,64]) that op4 (prefill) / op5 (decode) + # consume directly. is_neox=False = GPT-J adjacent-pair RoPE; q_weight=None = + # V4-Pro weightless Q. Decode fuses the paged SWA cache-write via the + # swa_* buffers; prefill passes them None and scatters its tail post-attn. + from aiter import dtypes + from aiter.ops.fused_qk_norm_rope_cache_quant import ( + fused_qk_norm_rope_group_quant, + ) + + num_tokens = q.shape[0] + # aiter derives rot_dim from cos_cache.shape[-1]*2, so it needs the 2D + # [max_pos, rd//2] tables; the _V4RoPE caches are 4D [., rd//2, 1, 1]. + cos_2d = cos_cache.squeeze(-2).squeeze(-2) + sin_2d = sin_cache.squeeze(-2).squeeze(-2) + q_packed, q_rope, k_packed, k_rope = fused_qk_norm_rope_group_quant( + q.view(num_tokens, n_local_heads, head_dim), + kv, + kv_weight, + positions, + cos_2d, + sin_2d, + eps, + is_neox=False, + q_out_dtype=dtypes.fp8, + q_weight=None, + quant_group_size=64, + scale_dtype="e8m0", + swa_nope_scale_buff=swa_nope_scale_buff, + swa_rope_buff=swa_rope_buff, + swa_block_tables=swa_block_tables, + swa_block_size=swa_block_size, + batch_id_per_token=batch_id_per_token, + ) + return QKNormRopeOut( + q_packed=q_packed, q_rope=q_rope, k_packed=k_packed, k_rope=k_rope + ) + + def qk_norm_rope_maybe_quant_reference( q: torch.Tensor, kv: torch.Tensor, @@ -644,3 +797,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_bf16( + 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 diff --git a/atom/model_ops/v4_kernels/state_writes.py b/atom/model_ops/v4_kernels/state_writes.py index 45fc9ecbf..ec92bca7c 100644 --- a/atom/model_ops/v4_kernels/state_writes.py +++ b/atom/model_ops/v4_kernels/state_writes.py @@ -127,10 +127,24 @@ def swa_write( swa_region: torch.Tensor, block_size: int, write_per_batch: int, + *, + k_packed: torch.Tensor | None = None, + k_rope: torch.Tensor | None = None, + swa_region_rope: torch.Tensor | None = None, prefix: str = "", ) -> None: - """paged-SWA in-place write. For the last - `min(tok_n_b, write_per_batch)` tokens of every seq `b ∈ [0, bs)` this fwd + """paged-SWA in-place write, dispatching on the kv-cache layout. + + Native 2buff fp8 (``swa_region_rope`` provided): the op-quantized extend K + comes in as ``k_packed`` (fp8 NoPE) + ``k_rope`` (bf16 RoPE tail), in the + ``[T, *]`` or ``[T, 1, *]`` layout produced by the quant kernel; delegates to + :func:`swa_write_2buff_prepacked`, which scatters both into their paged pools + (``swa_region`` = NoPE pool, ``swa_region_rope`` = RoPE pool) — a pure + dtype-agnostic copy, no requant. The bf16 ``kv`` arg is unused on this path + (the caller may pass ``None``). + + Otherwise (bf16): for the last `min(tok_n_b, write_per_batch)` tokens of + every seq `b ∈ [0, bs)` this fwd (`tok_n_b = cu_seqlens_q[b+1] - cu_seqlens_q[b]`, `bs = block_tables.shape[0]`), write `kv[r]` to the content-addressed SWA region: swa_region[block_tables[b, pos//block_size] * block_size @@ -142,7 +156,8 @@ def swa_write( cached block instead of a stale ring (issue #1417). Args: - kv: [T, head_dim] per-fwd KV (BF16). `T = cu_seqlens_q[bs]`. + kv: [T, head_dim] per-fwd KV (BF16). bf16 path only; `T = cu_seqlens_q[bs]`. + May be ``None`` on the fp8 2buff path (``k_packed`` is used instead). positions: [T'] int — full forward_vars["positions"] (`T' >= T`). cu_seqlens_q: [bs+1] int — exact size (`bs == block_tables.shape[0]`). block_tables: [bs, max_blocks_per_seq] int32 — logical→physical block. @@ -152,7 +167,28 @@ def swa_write( block_size: tokens per block (= V4 block_size, 128). write_per_batch: `min(max_q_len, block_size_window)` — max tokens written per seq this fwd (grid y dim, kernel `constexpr`). + k_packed: [T, 512] or [T, 1, 512] fp8 NoPE extend K — fp8 2buff path only. + k_rope: [T, rope_head_dim] or [T, 1, rope_head_dim] bf16 RoPE tail — fp8 + 2buff path only. + swa_region_rope: [num_pages, rope_head_dim] bf16 RoPE pool — presence + selects the fp8 2buff path. """ + if swa_region_rope is not None: + # fp8 2buff: scatter the op-quantized extend K (k_packed/k_rope) into both + # paged SWA pools. Flatten the [T, 1, *] quant-kernel views to [T, *]; the + # bf16 `kv` source is unused here. + swa_write_2buff_prepacked( + k_packed.view(k_packed.shape[0], -1), + k_rope.view(k_rope.shape[0], -1), + positions, + cu_seqlens_q, + block_tables, + swa_region, + swa_region_rope, + block_size, + write_per_batch, + ) + return assert kv.dim() == 2, f"kv must be [T, D], got {kv.shape}" assert positions.dim() == 1 assert ( @@ -226,6 +262,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 diff --git a/atom/model_ops/v4_kernels/v4_quant.py b/atom/model_ops/v4_kernels/v4_quant.py new file mode 100644 index 000000000..49f078620 --- /dev/null +++ b/atom/model_ops/v4_kernels/v4_quant.py @@ -0,0 +1,549 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""V4 MLA on-the-fly quantization helpers for the hipkittens decode kernel. + +The hipkittens kernel ``aiter.mla.mla_v40_decode_fwd`` consumes Q/KV in a +two-buffer layout: + + nope_scale_buff: [..., 512] FP8 (1 byte/elem) + bytes [0 , 448): NOPE FP8 (per-token feature, position-agnostic) + bytes [448 , 462): 14 duplicated E8M0 scales (7 per-64-elt scales x2) + bytes [462 , 512): 50 bytes unused trailing pad + rope_buff: [..., 64] BF16 (per-token RoPE-rotated, kept BF16) + +ATOM currently stores ``unified_kv`` as a single contiguous bf16 tensor of +shape ``[..., 512]`` (NoPE 448 + RoPE 64 concatenated). For the PoC we +convert bf16 -> (fp8+scale+pad, bf16-rope) on every decode call. A future +phase 2 may change ``unified_kv`` to store the packed layout directly so the +runtime quantization is skipped. + +All constants and pack arithmetic mirror +``/mnt/raid0/ruitang3/git_repo/aiter/op_tests/test_mla_v4_persistent.py`` +(``V4_*`` constants and ``pack_v4_nope_scale``/``quantize_v4_nope_bpad8``). +""" + +from __future__ import annotations + +import math +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl +from aiter import dtypes + +V4_DIM_NOPE = 448 +V4_DIM_ROPE = 64 +V4_DIM_QK = V4_DIM_NOPE + V4_DIM_ROPE # 512 +V4_TILE = 64 +V4_NUM_TILES = V4_DIM_NOPE // V4_TILE # 7 +V4_DIM_SCALE_DUP = V4_DIM_NOPE // (V4_TILE // 2) # 14 +V4_DIM_QK_PACKED = 512 +V4_PACK_OFF_NOPE = 0 +V4_PACK_OFF_SCALE = V4_DIM_NOPE # 448 +V4_PACK_OFF_PAD = V4_DIM_NOPE + V4_DIM_SCALE_DUP # 462 + + +def _fp32_pow2_to_e8m0(pow2_fp32: torch.Tensor) -> torch.Tensor: + """Pack a power-of-2 fp32 scale into a 1-byte E8M0 exponent + (byte B encodes 2^(B-127); B=0 -> 0.0, B=255 -> INF).""" + safe = torch.where(pow2_fp32 > 0, pow2_fp32, torch.ones_like(pow2_fp32)) + biased = torch.log2(safe).round().to(torch.int32) + 127 + biased = torch.clamp(biased, 0, 254) + biased = torch.where(pow2_fp32 > 0, biased, torch.zeros_like(biased)) + return biased.to(torch.uint8) + + +def _cast_scale_inv_to_ue8m0_pow2(scales_inv: torch.Tensor) -> torch.Tensor: + """amax/FP8_AMAX -> ceil-log2 -> power-of-2 fp32.""" + return torch.pow(2.0, torch.clamp_min(scales_inv, 1e-4).log2().ceil()).to( + torch.float32 + ) + + +def _duplicate_each_lastdim(x: torch.Tensor) -> torch.Tensor: + """[..., N] -> [..., 2*N] with each element written twice.""" + return x.unsqueeze(-1).expand(*x.shape, 2).reshape(*x.shape[:-1], x.shape[-1] * 2) + + +def quantize_v4_nope_bpad8( + nope_src: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-tile (64-elt) E8M0 quantization of the NOPE segment. + + Returns ``(nope_fp8 [..., 448], scale_e8m0 [..., 7] uint8)``. + """ + fp8_amax = float(torch.finfo(dtypes.fp8).max) + nope_fp32 = nope_src.float() + leading = nope_fp32.shape[:-1] + tiled = nope_fp32.reshape(*leading, V4_NUM_TILES, V4_TILE) + active_scale_pow2 = _cast_scale_inv_to_ue8m0_pow2( + tiled.abs().amax(dim=-1) / fp8_amax + ) + nope_fp8 = ( + (tiled / active_scale_pow2.unsqueeze(-1)) + .to(dtypes.fp8) + .reshape(*leading, V4_DIM_NOPE) + ) + scale_e8m0 = _fp32_pow2_to_e8m0(active_scale_pow2) # [..., 7] uint8 + return nope_fp8, scale_e8m0 + + +def pack_v4_nope_scale( + nope_fp8: torch.Tensor, scale_e8m0: torch.Tensor +) -> torch.Tensor: + """Pack NOPE + duplicated E8M0 scale into a single 512-byte/token FP8 tensor. + + nope_fp8: [..., 448] FP8 + scale_e8m0: [..., 7] uint8 (E8M0 byte per quant tile) + returns: [..., 512] FP8 (NoPE | dup-scale x14 | pad x50) + """ + leading = nope_fp8.shape[:-1] + assert nope_fp8.shape[-1] == V4_DIM_NOPE + assert scale_e8m0.shape[-1] == V4_NUM_TILES + assert scale_e8m0.shape[:-1] == leading + + packed = torch.zeros( + (*leading, V4_DIM_QK_PACKED), dtype=torch.uint8, device=nope_fp8.device + ) + packed[..., V4_PACK_OFF_NOPE : V4_PACK_OFF_NOPE + V4_DIM_NOPE] = nope_fp8.view( + torch.uint8 + ) + packed[..., V4_PACK_OFF_SCALE : V4_PACK_OFF_SCALE + V4_DIM_SCALE_DUP] = ( + _duplicate_each_lastdim(scale_e8m0) + ) + return packed.view(dtypes.fp8) + + +def _e8m0_to_fp32_pow2(scale_e8m0: torch.Tensor) -> torch.Tensor: + """Inverse of ``_fp32_pow2_to_e8m0``: E8M0 byte B -> fp32 2^(B-127). + + B == 0 decodes to 0.0 (the zero-scale sentinel produced by the forward + path for all-zero tiles).""" + biased = scale_e8m0.to(torch.int32) + pow2 = torch.pow(2.0, (biased - 127).float()) + return torch.where(biased > 0, pow2, torch.zeros_like(pow2)) + + +def dequantize_v4_2buff_to_bf16( + packed_fp8: torch.Tensor, + rope_bf16: torch.Tensor, +) -> torch.Tensor: + """Inverse of ``quantize_bf16_to_v4_2buff``. + + Takes the two-buffer layout ``(packed_fp8 [..., 512], rope_bf16 [..., 64])`` + and reconstructs the bf16 ``[..., 512]`` row (NoPE 448 + RoPE 64). + + The NoPE half is dequantized per 64-elt tile: ``fp8_val * 2^(B-127)`` where + ``B`` is the tile's E8M0 scale byte (read from the dup-scale region; we use + the first of each duplicated pair). Round-trips + ``quantize_bf16_to_v4_2buff`` to within fp8 per-tile quantization error. + """ + assert packed_fp8.shape[-1] == V4_DIM_QK_PACKED, ( + f"dequantize_v4_2buff_to_bf16: packed last dim must be " + f"{V4_DIM_QK_PACKED}, got {tuple(packed_fp8.shape)}" + ) + assert rope_bf16.shape[-1] == V4_DIM_ROPE, ( + f"dequantize_v4_2buff_to_bf16: rope last dim must be {V4_DIM_ROPE}, " + f"got {tuple(rope_bf16.shape)}" + ) + leading = packed_fp8.shape[:-1] + packed_u8 = packed_fp8.view(torch.uint8) + + nope_fp8 = packed_u8[..., V4_PACK_OFF_NOPE : V4_PACK_OFF_NOPE + V4_DIM_NOPE].view( + dtypes.fp8 + ) + nope_fp32 = nope_fp8.float().reshape(*leading, V4_NUM_TILES, V4_TILE) + + # Dup-scale region holds each of the 7 tile scales twice; take the even + # entries to recover the 7 per-tile E8M0 bytes. + scale_dup = packed_u8[..., V4_PACK_OFF_SCALE : V4_PACK_OFF_SCALE + V4_DIM_SCALE_DUP] + scale_e8m0 = scale_dup[..., 0::2] # [..., 7] + scale_pow2 = _e8m0_to_fp32_pow2(scale_e8m0) # [..., 7] fp32 + + nope_bf16 = ( + (nope_fp32 * scale_pow2.unsqueeze(-1)) + .reshape(*leading, V4_DIM_NOPE) + .to(torch.bfloat16) + ) + rope = rope_bf16.to(torch.bfloat16) + return torch.cat([nope_bf16, rope], dim=-1) + + +def quantize_bf16_to_v4_2buff( + bf16_src: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """End-to-end helper: bf16 [..., 512] -> (packed_fp8 [..., 512], rope_bf16 [..., 64]). + + Splits the input on the NoPE/RoPE boundary, quantizes the NoPE half via + ``quantize_v4_nope_bpad8`` + ``pack_v4_nope_scale``, and keeps the RoPE + half in bf16 (contiguous). + """ + assert bf16_src.shape[-1] == V4_DIM_QK, ( + f"quantize_bf16_to_v4_2buff: last dim must be {V4_DIM_QK}, " + f"got {tuple(bf16_src.shape)}" + ) + nope_src = bf16_src[..., :V4_DIM_NOPE] + rope_src = bf16_src[..., V4_DIM_NOPE:].to(torch.bfloat16).contiguous() + nope_fp8, scale_e8m0 = quantize_v4_nope_bpad8(nope_src) + packed_fp8 = pack_v4_nope_scale(nope_fp8, scale_e8m0) + return packed_fp8, rope_src + + +# --------------------------------------------------------------------------- +# Triton port of quantize_bf16_to_v4_2buff. +# +# The torch helper above chains ~6 tensor ops (reshape, amax, log2/ceil, +# to(fp8), uint8 bit-views, scatter-pack) — fine as an eager reference, but it +# graph-breaks / is CUDAGraph-hostile inside a `@support_torch_compile` region. +# This single-kernel version is the compile-safe path used to build the fp8 +# 2buff Q/K without the aiter HIP op (op1). Output is bit-compatible with +# ``quantize_bf16_to_v4_2buff`` (same round-to-nearest fp8 cast + same +# ceil-log2 e8m0 scale), so ``dequantize_v4_2buff_to_bf16`` round-trips either. +# --------------------------------------------------------------------------- + + +@triton.jit +def _bf16_to_v4_2buff_kernel( + src_ptr, # [N, 512] bf16 (NoPE 448 | RoPE 64) + packed_fp8_ptr, # [N, 512] fp8 out — NoPE fp8 written here (rest pre-zeroed) + packed_u8_ptr, # same buffer, uint8 view — e8m0 dup-scale bytes written here + rope_ptr, # [N, 64] bf16 out + src_row_stride, + N, + FP8_AMAX: tl.constexpr, + NOPE: tl.constexpr, + ROPE: tl.constexpr, + TILE: tl.constexpr, + NUM_TILES: tl.constexpr, + PACK_OFF_SCALE: tl.constexpr, + PACKED_ROW: tl.constexpr, # 512 (packed row stride) + BLOCK_N: tl.constexpr, +): + """One program per ``BLOCK_N``-row tile. For each of the 7 NoPE tiles: + per-64-elt amax -> ceil-log2 e8m0 power-of-2 scale -> fp8 quant, storing the + fp8 bytes into the NoPE region and the E8M0 byte (duplicated) into the + dup-scale region. The RoPE tail is copied through in bf16. The 50-byte pad + is left as the caller's pre-zeroed value.""" + pid = tl.program_id(0) + row = pid * BLOCK_N + tl.arange(0, BLOCK_N) + row_mask = row < N + + d_tile = tl.arange(0, TILE) + # Loop is unrolled at compile time (NUM_TILES=7 constexpr). + for t in tl.static_range(NUM_TILES): + cols = t * TILE + d_tile + x = tl.load( + src_ptr + row[:, None] * src_row_stride + cols[None, :], + mask=row_mask[:, None], + other=0.0, + ).to(tl.float32) + + amax = tl.max(tl.abs(x), axis=1) # [BLOCK_N] + scale_inv = amax / FP8_AMAX + # ceil-log2 power-of-2 scale (matches _cast_scale_inv_to_ue8m0_pow2: + # clamp_min(1e-4) guards log2 of an all-zero tile). + e = tl.ceil(tl.log2(tl.maximum(scale_inv, 1e-4))) # integer-valued fp32 + scale_pow2 = tl.exp2(e) + + xq = (x / scale_pow2[:, None]).to(packed_fp8_ptr.dtype.element_ty) + tl.store( + packed_fp8_ptr + row[:, None] * PACKED_ROW + cols[None, :], + xq, + mask=row_mask[:, None], + ) + + # E8M0 byte = clamp(round(log2(scale_pow2)) + 127, 0, 254). scale_pow2 is + # an exact power of two so round(log2)==e. Written twice (dup layout). + byte_f = tl.minimum(tl.maximum(e + 127.0, 0.0), 254.0) + byte = byte_f.to(tl.uint8) + off = PACK_OFF_SCALE + 2 * t + tl.store(packed_u8_ptr + row * PACKED_ROW + off, byte, mask=row_mask) + tl.store(packed_u8_ptr + row * PACKED_ROW + off + 1, byte, mask=row_mask) + + # RoPE tail passthrough (bf16, unquantized). + r_cols = tl.arange(0, ROPE) + r = tl.load( + src_ptr + row[:, None] * src_row_stride + (NOPE + r_cols)[None, :], + mask=row_mask[:, None], + ) + tl.store( + rope_ptr + row[:, None] * ROPE + r_cols[None, :], + r, + mask=row_mask[:, None], + ) + + +def quantize_bf16_to_v4_2buff_triton( + bf16_src: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Triton port of :func:`quantize_bf16_to_v4_2buff` (compile-safe, single + launch). ``bf16_src`` may have any leading shape; the last dim must be 512. + + Returns ``(packed_fp8 [..., 512], rope_bf16 [..., 64])`` — bit-compatible + with the torch helper (same fp8 rounding + e8m0 ceil-log2 scale). + """ + assert bf16_src.shape[-1] == V4_DIM_QK, ( + f"quantize_bf16_to_v4_2buff_triton: last dim must be {V4_DIM_QK}, " + f"got {tuple(bf16_src.shape)}" + ) + assert bf16_src.dtype == torch.bfloat16, ( + f"quantize_bf16_to_v4_2buff_triton: input must be bf16, " + f"got {bf16_src.dtype}" + ) + leading = bf16_src.shape[:-1] + src = bf16_src.reshape(-1, V4_DIM_QK) + src = src.contiguous() if src.stride(-1) != 1 else src + N = src.shape[0] + + # Zeroed so the 50-byte trailing pad (and any untouched byte) is 0, matching + # pack_v4_nope_scale's torch.zeros allocation. + packed = torch.zeros((N, V4_DIM_QK_PACKED), dtype=dtypes.fp8, device=src.device) + rope = torch.empty((N, V4_DIM_ROPE), dtype=torch.bfloat16, device=src.device) + + if N > 0: + block_n = 8 + while block_n > 1 and block_n > N: + block_n //= 2 + grid = (triton.cdiv(N, block_n),) + _bf16_to_v4_2buff_kernel[grid]( + src, + packed, + packed.view(torch.uint8), + rope, + src.stride(0), + N, + FP8_AMAX=float(torch.finfo(dtypes.fp8).max), + NOPE=V4_DIM_NOPE, + ROPE=V4_DIM_ROPE, + TILE=V4_TILE, + NUM_TILES=V4_NUM_TILES, + PACK_OFF_SCALE=V4_PACK_OFF_SCALE, + PACKED_ROW=V4_DIM_QK_PACKED, + BLOCK_N=block_n, + ) + + return ( + packed.view(*leading, V4_DIM_QK_PACKED), + rope.view(*leading, V4_DIM_ROPE), + ) + + +# --------------------------------------------------------------------------- +# Torch reference for aiter.mla.mla_decode_fwd_v4_nm +# +# Mirrors op_tests/test_mla_v4_nm.py::_torch_attn_decode_fp8_dequant_ref + +# _torch_attn_decode_bf16_golden: dequantize the exact 2-buffer FP8 tensors the +# asm kernel consumes, then run plain-torch scaled-dot-product attention with a +# per-batch loop, GQA broadcast, and an attention sink. Reusing this module's +# `dequantize_v4_2buff_to_bf16` (the inverse of the packing above) keeps the +# reference self-consistent with ATOM's own quant path. +# +# This isolates "kernel math bug" from "FP8 quant noise": feed the SAME packed +# bytes to the kernel and to this reference, and any diff is kernel math. +# --------------------------------------------------------------------------- + + +def _attn_decode_bf16_golden( + q_bf16: torch.Tensor, # [total_q, num_heads, D] + kv_bf16: torch.Tensor, # [num_page, page_size, num_kv_heads, D] + qo_indptr: torch.Tensor, # [num_seqs+1] q rows per seq (cumulative) + kv_indptr: torch.Tensor, # [num_seqs+1] pages per seq (cumulative) + kv_page_indices: torch.Tensor, # [total_pages_used] + kv_last_page_lens: torch.Tensor, # [num_seqs] + sm_scale: float, + v_head_dim: int, + attn_sink: Optional[torch.Tensor] = None, # [num_heads] or None +) -> Tuple[torch.Tensor, torch.Tensor]: + """Pure-torch BF16 attention reference (FP32 accum). + + Per-batch loop, scaled-dot-product attention over the full D=NoPE+RoPE + query/key dim, with GQA broadcast (each KV head serves gqa_ratio Q heads) + and an optional per-head attention sink (virtual K-column of logit + ``sink[h]`` shared by every query token). V is the first ``v_head_dim`` + columns of the (dequantized) KV row. + + Returns: + out [total_q, num_heads, v_head_dim] FP32 + lse [total_q, num_heads] FP32 + """ + total_q, num_heads, d = q_bf16.shape + num_kv_heads = kv_bf16.size(2) + gqa_ratio = num_heads // num_kv_heads + page_size = kv_bf16.size(1) + device = q_bf16.device + + out = torch.zeros( + (total_q, num_heads, v_head_dim), dtype=torch.float32, device=device + ) + lse_full = torch.full( + (total_q, num_heads), float("inf"), dtype=torch.float32, device=device + ) + + batch = qo_indptr.numel() - 1 + qo_cpu = qo_indptr.cpu().tolist() + kv_cpu = kv_indptr.cpu().tolist() + last_cpu = kv_last_page_lens.cpu().tolist() + + for b in range(batch): + qs, qe = qo_cpu[b], qo_cpu[b + 1] + ps, pe = kv_cpu[b], kv_cpu[b + 1] + num_pages_b = pe - ps + if num_pages_b == 0: + continue + + page_ids = kv_page_indices[ps:pe] + kv_pages = kv_bf16[page_ids] # [num_pages_b, page_size, num_kv_heads, D] + # Flatten (page, slot) -> token, then trim the partial last page. + kv_flat = kv_pages.reshape(-1, num_kv_heads, d) + total_tokens = (num_pages_b - 1) * page_size + last_cpu[b] + kv_b = kv_flat[:total_tokens].float() # [seq_k, num_kv_heads, D] + + # GQA broadcast: replicate each KV head across its gqa_ratio Q heads. + kv_heads = kv_b.repeat_interleave(gqa_ratio, dim=1) # [seq_k, num_heads, D] + + q_b = q_bf16[qs:qe].float() # [s_q, num_heads, D] + scores = torch.einsum("shd,khd->shk", q_b, kv_heads) * sm_scale # [s_q,H,seq_k] + + lse = scores.logsumexp(dim=-1) # [s_q, H] + if attn_sink is not None: + sink_b = attn_sink.view(1, num_heads).float() # [1, H] + m = torch.maximum(lse, sink_b) + denom = torch.exp(lse - m) + torch.exp(sink_b - m) + lse_final = m + torch.log(denom) + else: + lse_final = lse + probs = torch.exp(scores - lse_final.unsqueeze(-1)) # [s_q, H, seq_k] + + v_heads = kv_heads[..., :v_head_dim] # MLA: V == first v_head_dim of K + out_b = torch.einsum("shk,khv->shv", probs, v_heads) # [s_q, H, v_head_dim] + out[qs:qe] = out_b + lse_full[qs:qe] = lse_final + + return out, lse_full + + +def mla_decode_fwd_v4_nm_ref( + q, # [total_q, num_heads, 512] FP8 packed (NoPE 448 | dup-scale 14 | pad 50) + qrope, # [total_q, num_heads, 64] BF16 + kv_buffer, # [num_page, page_size, num_kv_heads, 512] FP8 packed + kvrope, # [num_page, page_size, num_kv_heads, 64] BF16 + output, # [total_q, num_heads, v_head_dim] BF16 (written in-place) + qo_indptr, # [num_seqs+1] + kv_indptr, # [num_seqs+1] + kv_page_indices, # [num_page_used] + kv_last_page_lens, # [num_seqs] + max_seqlen_q, + *, + sink=None, # [num_heads] FP32, or None for "no sink" math + split_indptr=None, # accepted for signature parity; math is split-invariant + sm_scale=None, # None -> 1/sqrt(D) matching the kernel's hardcoded scale + out_16_nosplit=0, # accepted for parity; reference always writes both buffers + num_kv_splits=1, # accepted for parity; KV splitting does not change the math + logits=None, + attn_lse=None, +): + """Torch reference for ``aiter.mla.mla_decode_fwd_v4_nm``. + + Dequantizes the same 2-buffer FP8 tensors the asm kernel reads (via + :func:`dequantize_v4_2buff_to_bf16`) and computes full MLA decode attention + in plain torch. Matches the kernel's public contract: + + - Returns ``(logits, attn_lse)`` in kernel-native layout + ``logits [total_q, num_kv_splits, num_heads, v_head_dim]`` FP32 and + ``attn_lse [total_q, num_kv_splits, num_heads, 1]`` FP32. The final + attention result lands in the ``[:, 0]`` split slot (single-pass + contract; any remaining split slots are left zero). + - Also writes the BF16 result into ``output`` in-place, so callers that + read the merged/``out_16_nosplit`` buffer see the same answer. + + ``num_kv_splits`` / ``split_indptr`` / ``out_16_nosplit`` are accepted only + for signature parity: the KV split is a kernel-side perf optimization whose + logsumexp-merged result is mathematically identical to the un-split full + attention this reference computes. ``sm_scale`` defaults to ``1/sqrt(D)`` + (D = NoPE + RoPE = 512) to mirror the kernel, which ignores the passed + value and hardcodes that scale. + """ + num_seqs = qo_indptr.numel() - 1 + num_heads = q.size(1) + num_kv_heads = kv_buffer.size(2) + gqa_ratio = num_heads // num_kv_heads + v_head_dim = output.size(2) + total_q = num_seqs * max_seqlen_q + device = q.device + + # ---- sink validation (mirror the kernel wrapper's contract) ---- + if sink is not None: + if sink.dtype != dtypes.fp32: + raise ValueError( + f"mla_decode_fwd_v4_nm_ref: `sink` must be FP32, got {sink.dtype}." + ) + if not sink.is_contiguous(): + raise ValueError( + "mla_decode_fwd_v4_nm_ref: `sink` must be contiguous " + f"(got strides={sink.stride()})." + ) + if sink.numel() != num_heads: + raise ValueError( + f"mla_decode_fwd_v4_nm_ref: `sink` numel {sink.numel()} != " + f"num_heads {num_heads} (= num_kv_heads({num_kv_heads}) * " + f"gqa_ratio({gqa_ratio}))." + ) + if sink.device != device: + raise ValueError( + "mla_decode_fwd_v4_nm_ref: `sink` must be on the same device " + f"as `q` (got sink={sink.device}, q={device})." + ) + # An all -inf sink is the documented "no sink" no-op; treat it as None + # so exp(-inf - m) = 0 contributes nothing (and avoids NaN when m=-inf). + attn_sink = None if bool(torch.all(torch.isneginf(sink))) else sink + else: + attn_sink = None + + # ---- dequantize the exact packed bytes the kernel consumes ---- + q_bf16 = dequantize_v4_2buff_to_bf16(q, qrope) # [total_q, num_heads, 512] + kv_bf16 = dequantize_v4_2buff_to_bf16( + kv_buffer, kvrope + ) # [num_page, page_size, num_kv_heads, 512] + + d = q_bf16.size(-1) + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(d) + + out_fp32, lse_fp32 = _attn_decode_bf16_golden( + q_bf16, + kv_bf16, + qo_indptr, + kv_indptr, + kv_page_indices, + kv_last_page_lens, + float(sm_scale), + v_head_dim, + attn_sink=attn_sink, + ) + + # ---- write outputs in the kernel's public layout ---- + output.copy_(out_fp32.to(output.dtype)) + + if logits is None: + logits = torch.zeros( + (total_q, num_kv_splits, num_heads, v_head_dim), + dtype=dtypes.fp32, + device=device, + ) + else: + logits.zero_() + if attn_lse is None: + attn_lse = torch.zeros( + (total_q, num_kv_splits, num_heads, 1), + dtype=dtypes.fp32, + device=device, + ) + else: + attn_lse.zero_() + + logits[:, 0] = out_fp32 + attn_lse[:, 0, :, 0] = lse_fp32 + + return logits, attn_lse diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 4900a9f24..3388c3840 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -958,6 +958,13 @@ def __init__( # of `self.kv_cache`. Bound by the V4 builder when `kv_cache.dtype` is # FP8 (Indexer-inner Compressor); None for BF16 cache (Main path). self.cache_scale: Optional[torch.Tensor] = None + # Native 2buff fp8 Main path (CSA/HCA Main under --kv_cache_dtype fp8): + # parallel bf16 rope pool [NB, k_per_block, rope_head_dim] and a write + # mode tag. Bound by DeepseekV4AttentionMetadataBuilder; "main_2buff_fp8" + # routes the compress scatter to the flydsl group_fp8 path, "bf16" / + # "indexer_fp8" keep the existing behavior. + self.kv_cache_rope: Optional[torch.Tensor] = None + self.write_mode: str = "bf16" # State cache (per paper §3.6.1 "uncompressed tail + B-side overlap # window" portion). Indexed as a single ring buffer of size @@ -1096,19 +1103,27 @@ def forward( # fwd's data for the NEXT fwd's overlap — must run AFTER the fused # kernel. cos_cache, sin_cache = self.rotary_emb.cos_cache, self.rotary_emb.sin_cache - # Quant path triggers when the bound cache is FP8 (Indexer-inner). - # `self.cache_scale` is bound alongside `self.kv_cache` by the V4 - # builder when the cache is FP8 (strided fp32 view of the per-block - # scale region). - is_quant = self.kv_cache is not None and self.kv_cache.dtype != torch.bfloat16 + # Three scatter modes, discriminated by the bound `write_mode`: + # - "indexer_fp8": per-row fp8 + preshuffle (Indexer-inner Compressor). + # - "main_2buff_fp8": CSA/HCA Main under --kv_cache_dtype fp8 — native + # group_fp8 2buff (nope-fp8 + inline e8m0 into `kv_cache`, bf16 rope + # into `kv_cache_rope`). Routed to the flydsl group_fp8 path. + # - "bf16": plain bf16 Main scatter. + # `self.cache_scale` is bound alongside an fp8 `kv_cache` by the V4 + # builder (indexer-inner only; group_fp8 carries scale inline so + # cache_scale stays None). + main_2buff_fp8 = self.write_mode == "main_2buff_fp8" + is_quant = self.write_mode == "indexer_fp8" # Skip the kernel's cache scatter during warmup (kv_cache/block_tables # not yet bound). if block_tables is None or self.kv_cache is None: scatter_kv_cache = None scatter_block_tables = None + scatter_kv_cache_rope = None else: scatter_kv_cache = self.kv_cache scatter_block_tables = block_tables + scatter_kv_cache_rope = self.kv_cache_rope if main_2buff_fp8 else None fused_compress_attn( kv_in=kv, score_in=score, @@ -1133,6 +1148,8 @@ def forward( use_ue8m0=(self.scale_fmt == "ue8m0"), preshuffle=True, fp8_max=(torch.finfo(self.kv_cache.dtype).max if is_quant else None), + main_2buff_fp8=main_2buff_fp8, + kv_cache_rope=scatter_kv_cache_rope, prefix=f"{self.prefix}.fused_compress_attn", ) update_compressor_states( @@ -1788,6 +1805,12 @@ def __init__( self.layer_name = prefix atom_config = get_current_atom_config() atom_config.compilation_config.static_forward_context[self.layer_name] = self + # Frozen bool: when KV cache dtype is fp8, route writes/attention to the + # native 2buff fp8 path (compute-only 2buff quant, op4 fp8 prefill, op5 + # asm decode). Dynamo specializes on this constant so the bf16 path + # traces unchanged. The rope buffers (swa_kv_rope / unified_kv_rope) are + # bound onto the module by DeepseekV4AttentionMetadataBuilder. + self.kv_fp8 = atom_config.kv_cache_dtype == "fp8" def process_weights_after_loading(self) -> None: """Dequant wo_a (FP8 + e8m0 block scale) → BF16 in place. @@ -2161,7 +2184,20 @@ def _attn_core( # prefix reads see prior-chunk contents), so writing here (before the # decode attention reads the window) is safe. Prefill → swa_kv=None, # separate post-attn swa_write below. - q_sa, kv, q_scale, kv_scale = qk_norm_rope_maybe_quant( + # Fused per-head weightless Q RMSNorm + weighted KV RMSNorm + GPT-J RoPE, + # dispatching on the kv-cache layout inside the wrapper (fp8 2buff ↔ bf16) + # so this call site is branch-free. Decode FUSES the paged + # (content-addressed) SWA cache-write into the launch via swa_block_tables + # (bf16: flydsl swa_kv; fp8: aiter swa_nope/rope buffers — each K row + # scattered into its SWA pool, batch_id<0 skips CG-pad), running before the + # decode attention reads the window (no ordering hazard). Prefill passes + # swa_*=None and scatters its window tail post-attn below. + # + # bf16 → qkn.q_sa / qkn.kv populated; fp8 2buff → qkn.q_packed / qkn.q_rope + # / qkn.k_packed / qkn.k_rope populated (the 2buff layout nope-fp8 [.,512] + + # rope-bf16 [.,64] that op4 (prefill) / op5 (decode) consume with no + # requant). The inactive path's fields stay None. + qkn = qk_norm_rope_maybe_quant( q, kv_pre, self.kv_norm.weight, @@ -2174,16 +2210,21 @@ def _attn_core( self.eps, quant_q=False, quant_k=False, - swa_kv=self.swa_kv if is_decode else None, + fp8_2buff=self.kv_fp8, batch_id_per_token=attn_md.batch_id_per_token if is_decode else None, swa_block_tables=swa_block_tables_gpu if is_decode else None, swa_block_size=swa_block_size if is_decode else None, + # bf16 SWA fusion (flydsl kernel / Triton fallback): + swa_kv=self.swa_kv if (is_decode and not self.kv_fp8) else None, swa_cu_seqlens_q=attn_md.cu_seqlens_q if is_decode else None, swa_write_per_batch=attn_md.max_seqlen_q if is_decode else None, + # fp8 2buff SWA fusion (aiter group-quant launch): + swa_nope_scale_buff=self.swa_kv if (is_decode and self.kv_fp8) else None, + swa_rope_buff=self.swa_kv_rope if (is_decode and self.kv_fp8) else None, prefix=f"{self.layer_name}.qk_norm_rope_maybe_quant", ) - if _V4_USE_REF_QUANT: - act_quant_inplace(kv[..., :-rd], 64, self.scale_fmt) + if _V4_USE_REF_QUANT and not self.kv_fp8: + act_quant_inplace(qkn.kv[..., :-rd], 64, self.scale_fmt) # HCA if use_async_compress: @@ -2227,13 +2268,23 @@ def _attn_core( else: # ratio == 128 kv_indices = attn_md.kv_indices_hca kv_indptr = attn_md.kv_indptr_hca + # Dispatch on kv-cache layout inside the wrapper: fp8 2buff + # (unified_kv_rope set) → aiter asm op5 with pre-packed fp8 Q + the + # 2buff fp8/bf16 pools read with no requant; bf16 (unified_kv_rope + # None) → Triton. qo_indptr / kv_last_page_lens come from the builder + # (None on bf16, ignored by the Triton path). o = sparse_attn_v4_paged_decode( - q_sa, + qkn.q_sa, self.unified_kv, kv_indices, kv_indptr, self.attn_sink, self.softmax_scale, + unified_kv_rope=self.unified_kv_rope, + q_packed_in=qkn.q_packed, + q_rope_in=qkn.q_rope, + qo_indptr=attn_md.qo_indptr, + kv_last_page_lens=attn_md.kv_last_page_lens, prefix=f"{self.layer_name}.sparse_attn_decode", ) # [S, H, head_dim] else: @@ -2259,7 +2310,7 @@ def _attn_core( pcp_on = _pcp_active() if pcp_on: pcp_ws = get_pcp_world_size() - kv_full = pcp_allgather_rerange(kv, pcp_ws) + kv_full = pcp_allgather_rerange(qkn.kv, pcp_ws) # positions must match kv_full's full-sequence coords for the # swa_write ring addressing (`positions[src] % cache_size`). # `positions` here is this rank's 1/W shard (split in @@ -2268,7 +2319,7 @@ def _attn_core( # the builder reindexed to 1/W). positions_full = pcp_allgather_rerange(positions, pcp_ws) else: - kv_full = kv + kv_full = qkn.kv positions_full = positions if ratio == 0: @@ -2282,8 +2333,17 @@ def _attn_core( kv_indptr_prefix = attn_md.kv_indptr_prefix_hca else: raise ValueError(f"Unsupported compress_ratio {ratio}") + # Dispatch on kv-cache layout inside the wrapper: fp8 2buff + # (unified_kv_rope set) → aiter op4 with the 2buff fp8 prefix pool + # (nope-fp8 + rope-bf16) + op-quantized fp8 Q and extend K, no dequant + # of the prefix and no torch quant; bf16 → OPUS / Triton over q_sa and + # the bf16 extend kv_full. (fp8 + PCP is out of scope: k_packed/k_rope + # are this fwd's shard, NOT PCP all-gathered like kv_full.) On bf16 the + # wrapper reuses out=qkn.q_sa as the attention output buffer (q_sa is + # not needed after this call → avoids an extra empty_like); fp8 ignores + # both q_sa and out. o = sparse_attn_v4_paged_prefill( - q_sa, + qkn.q_sa, self.unified_kv, kv_indices_prefix, kv_indptr_prefix, @@ -2294,9 +2354,14 @@ def _attn_core( self.softmax_scale, # Reuse q_sa as the attention output buffer; q_sa is not needed # after this call and this avoids an extra empty_like allocation. - out=q_sa, + out=qkn.q_sa, + unified_kv_rope=self.unified_kv_rope, + q_packed=qkn.q_packed, + q_rope=qkn.q_rope, + k_packed=qkn.k_packed, + k_rope=qkn.k_rope, prefix=f"{self.layer_name}.sparse_attn_prefill", - ) # [S, H, head_dim] + ) # [S, H, head_dim] bf16 # swa_write AFTER attn so chunked-prefill prefix SWA reads see the # prior chunk's contents (not this chunk's just-computed tail). # OPT (window-only prefill write): only write each seq's trailing @@ -2309,6 +2374,12 @@ def _attn_core( # window(128) <= chunk size, so any token's window reaches back at # most 127 tokens = within the prior chunk's written last-128; # free_after_prefill_chunk keeps that trailing block until read. + # Dispatch on kv-cache layout inside the wrapper (swa_region_rope set + # → fp8 2buff path, which scatters the op-quantized extend K + # (k_packed/k_rope) into both paged SWA pools; else bf16 single-pool + # write over kv_full). Both are window-only (same trailing-window + # semantics). fp8 uses this fwd's shard, no PCP — positions_full == + # positions when PCP off, and kv_full is None on fp8. swa_write( kv_full, positions_full, @@ -2317,6 +2388,9 @@ def _attn_core( self.swa_kv, swa_block_size, min(self.window_size, attn_md.max_seqlen_q), + k_packed=qkn.k_packed, + k_rope=qkn.k_rope, + swa_region_rope=self.swa_kv_rope, prefix=f"{self.layer_name}.swa_write", )