v28 attempt: PV MMA (128,64) - cosine 0.004, debugging
This commit is contained in:
651
cutedsl/csa_hca_compressor_PYTORCH_EXAMPLE.py
Normal file
651
cutedsl/csa_hca_compressor_PYTORCH_EXAMPLE.py
Normal file
@@ -0,0 +1,651 @@
|
||||
"""
|
||||
CSA / HCA Token-Level Compressor for DeepSeek-V4.
|
||||
|
||||
Implements Section 2.3 of the DeepSeek-V4 paper exactly:
|
||||
- CSA (m=4): overlapping weighted sum over 2m hidden states per block
|
||||
- HCA (m'=128): non-overlapping weighted sum over m' hidden states per block
|
||||
|
||||
Both produce compressed KV entries C^Comp ∈ R^{n/m × c} where each entry
|
||||
is a weighted sum of hidden states using softmax-normalised gate weights.
|
||||
|
||||
CSA additionally produces compressed indexer keys K^IComp ∈ R^{n/m × c_I}
|
||||
for the Lightning Indexer (top-k sparse selection).
|
||||
|
||||
V4-Pro reference dimensions (Section 4.2.1):
|
||||
d = 7168 hidden dim
|
||||
c = 512 head dim (CSA and HCA both)
|
||||
m = 4 CSA compression ratio
|
||||
m' = 128 HCA compression ratio
|
||||
c_I = 128 indexer head dim
|
||||
n_I_h = 64 num indexer query heads
|
||||
n_win = 128 sliding window size (separate, not handled here)
|
||||
rope_dim = 64 partial RoPE on last 64 dims of each head
|
||||
|
||||
Design notes
|
||||
------------
|
||||
* BF16 matmuls throughout — swap the _proj() calls for your NVFP4 GEMMs.
|
||||
* No batch dimension: one sequence at a time (matching decode latency path).
|
||||
* CompressorState carries the incomplete-block tail and the previous block's
|
||||
raw projections needed for the CSA overlap.
|
||||
* Partial RoPE is applied to the last rope_dim=64 dims of C^Comp before
|
||||
the entry is stored in the compressed KV cache, using the representative
|
||||
position = last token of the block. The inverse-RoPE on attention outputs
|
||||
is handled by your attention kernel (already in blackwell_attention.py).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# State
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class CompressorState:
|
||||
"""
|
||||
Per-sequence mutable state for the compressor.
|
||||
|
||||
tail_hidden — hidden states that have arrived but don't yet fill a
|
||||
complete compression block. Shape (tail_len, d),
|
||||
0 <= tail_len < m.
|
||||
prev_hidden — the m hidden states from the previous complete block.
|
||||
Needed for the CSA overlap (C^b / Z^b projections).
|
||||
None before the first block is committed.
|
||||
Not used by HCA (no overlap).
|
||||
compressed_kv — accumulated C^Comp entries, shape (n_blocks, c).
|
||||
compressed_indexer_kv — accumulated K^IComp entries, shape (n_blocks, c_I).
|
||||
None for HCA layers.
|
||||
"""
|
||||
tail_hidden: Optional[torch.Tensor] = None # (tail_len, d)
|
||||
prev_hidden: Optional[torch.Tensor] = None # (m, d) CSA only
|
||||
compressed_kv: Optional[torch.Tensor] = None # (n_blocks, c)
|
||||
compressed_indexer_kv: Optional[torch.Tensor] = None # (n_blocks, c_I)
|
||||
num_blocks: int = 0
|
||||
|
||||
def reset(self):
|
||||
self.tail_hidden = None
|
||||
self.prev_hidden = None
|
||||
self.compressed_kv = None
|
||||
self.compressed_indexer_kv = None
|
||||
self.num_blocks = 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _apply_partial_rope(
|
||||
x: torch.Tensor, # (..., c)
|
||||
positions: torch.Tensor, # (...,) int64
|
||||
cos_sin_cache: torch.Tensor, # (max_pos, rope_dim)
|
||||
nope_dim: int,
|
||||
rope_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""GPT-J style RoPE on the last rope_dim dimensions only."""
|
||||
if rope_dim == 0:
|
||||
return x
|
||||
half = rope_dim // 2
|
||||
cos = cos_sin_cache[positions, :half].to(x.dtype) # (..., half)
|
||||
sin = cos_sin_cache[positions, half:].to(x.dtype) # (..., half)
|
||||
out = x.clone()
|
||||
rope_part = out[..., nope_dim:] # (..., rope_dim)
|
||||
even = rope_part[..., 0::2]
|
||||
odd = rope_part[..., 1::2]
|
||||
out[..., nope_dim:][..., 0::2] = even * cos - odd * sin
|
||||
out[..., nope_dim:][..., 1::2] = even * sin + odd * cos
|
||||
return out
|
||||
|
||||
|
||||
def _proj(x: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Linear projection: x @ W.
|
||||
x: (..., d) W: (d, out_dim) → (..., out_dim)
|
||||
|
||||
*** SWAP THIS FOR YOUR NVFP4 GEMM. ***
|
||||
The W tensor would become your NVFP4 weight + scale_b + gsb triple,
|
||||
and x would be quantised to NVFP4 activation before the call.
|
||||
"""
|
||||
return x.to(W.dtype) @ W
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CSA compressor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class CSACompressor:
|
||||
"""
|
||||
Compressed Sparse Attention token-level compressor.
|
||||
|
||||
Paper equations (11) and (12), Section 2.3.1:
|
||||
|
||||
C^a = H · W^a_KV Z^a = H · W^a_Z (current block)
|
||||
C^b = H · W^b_KV Z^b = H · W^b_Z (prev-block overlap)
|
||||
|
||||
For compressed block i (0-indexed):
|
||||
[S^a ; S^b] = softmax_row( [Z^a_{cur} + B^a ; Z^b_{prev} + B^b] )
|
||||
C^Comp_i = Σ S^a_j ⊙ C^a_j + Σ S^b_j ⊙ C^b_j
|
||||
|
||||
When i=0: Z^b / C^b are padded with -inf / 0 so only Z^a contributes.
|
||||
|
||||
The same compression is applied independently to produce indexer keys,
|
||||
using separate projections W^I_KV, W^I_Z, B^I_a, B^I_b.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_dim: int = 7168, # d
|
||||
head_dim: int = 512, # c
|
||||
compress_ratio: int = 4, # m
|
||||
indexer_head_dim: int = 128, # c_I
|
||||
num_indexer_heads: int = 64, # n_I_h (not used in compressor itself)
|
||||
nope_dim: int = 448, # c - rope_dim = 512 - 64
|
||||
rope_dim: int = 64,
|
||||
device: str = "cuda",
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
):
|
||||
self.d = hidden_dim
|
||||
self.c = head_dim
|
||||
self.m = compress_ratio
|
||||
self.c_I = indexer_head_dim
|
||||
self.n_I_h = num_indexer_heads
|
||||
self.nope = nope_dim
|
||||
self.rope = rope_dim
|
||||
self.device = device
|
||||
self.dtype = dtype
|
||||
|
||||
# ── Main KV projection weights ──────────────────────────────
|
||||
# W^a_{KV}, W^b_{KV}: (d, c)
|
||||
# W^a_Z, W^b_Z: (d, c)
|
||||
self.W_a_KV = self._param(hidden_dim, head_dim)
|
||||
self.W_b_KV = self._param(hidden_dim, head_dim)
|
||||
self.W_a_Z = self._param(hidden_dim, head_dim)
|
||||
self.W_b_Z = self._param(hidden_dim, head_dim)
|
||||
|
||||
# Positional biases B^a, B^b: (m, c) — learnable per-position offsets
|
||||
# added to the gate logits before softmax.
|
||||
self.B_a = self._param(compress_ratio, head_dim)
|
||||
self.B_b = self._param(compress_ratio, head_dim)
|
||||
|
||||
# ── Indexer key projection weights ──────────────────────────
|
||||
# Same overlap structure, separate projections, output dim c_I.
|
||||
self.W_I_a_KV = self._param(hidden_dim, indexer_head_dim)
|
||||
self.W_I_b_KV = self._param(hidden_dim, indexer_head_dim)
|
||||
self.W_I_a_Z = self._param(hidden_dim, indexer_head_dim)
|
||||
self.W_I_b_Z = self._param(hidden_dim, indexer_head_dim)
|
||||
self.B_I_a = self._param(compress_ratio, indexer_head_dim)
|
||||
self.B_I_b = self._param(compress_ratio, indexer_head_dim)
|
||||
|
||||
def _param(self, *shape) -> torch.Tensor:
|
||||
"""Uninitialised placeholder — replace with checkpoint-loaded tensor."""
|
||||
return torch.empty(*shape, dtype=self.dtype, device=self.device)
|
||||
|
||||
def load_weights(
|
||||
self,
|
||||
W_a_KV, W_b_KV, W_a_Z, W_b_Z, B_a, B_b,
|
||||
W_I_a_KV, W_I_b_KV, W_I_a_Z, W_I_b_Z, B_I_a, B_I_b,
|
||||
):
|
||||
"""Assign weights from checkpoint. All tensors moved to device/dtype."""
|
||||
def _cvt(t): return t.to(device=self.device, dtype=self.dtype)
|
||||
self.W_a_KV = _cvt(W_a_KV); self.W_b_KV = _cvt(W_b_KV)
|
||||
self.W_a_Z = _cvt(W_a_Z); self.W_b_Z = _cvt(W_b_Z)
|
||||
self.B_a = _cvt(B_a); self.B_b = _cvt(B_b)
|
||||
self.W_I_a_KV = _cvt(W_I_a_KV); self.W_I_b_KV = _cvt(W_I_b_KV)
|
||||
self.W_I_a_Z = _cvt(W_I_a_Z); self.W_I_b_Z = _cvt(W_I_b_Z)
|
||||
self.B_I_a = _cvt(B_I_a); self.B_I_b = _cvt(B_I_b)
|
||||
|
||||
# ----------------------------------------------------------------
|
||||
# Core: compress one block of m hidden states
|
||||
# ----------------------------------------------------------------
|
||||
|
||||
def _compress_block(
|
||||
self,
|
||||
cur_hidden: torch.Tensor, # (m, d) current block
|
||||
prev_hidden: Optional[torch.Tensor], # (m, d) or None if block 0
|
||||
cos_sin_cache: Optional[torch.Tensor],
|
||||
block_end_pos: int, # position of last token in block
|
||||
for_indexer: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Compress one block into a single C^Comp entry.
|
||||
|
||||
Returns: (c,) or (c_I,) compressed entry.
|
||||
|
||||
The overlap computation (equations 11-12):
|
||||
Z_cat = [Z^a + B^a ; Z^b + B^b] shape (2m, c) or (m, c) when block 0
|
||||
S = softmax(Z_cat, dim=0) normalise over the 2m position axis
|
||||
C_out = (S[:m] * C^a).sum(0) + (S[m:] * C^b).sum(0)
|
||||
"""
|
||||
m = self.m
|
||||
assert cur_hidden.shape[0] == m
|
||||
|
||||
if for_indexer:
|
||||
W_a_KV, W_b_KV = self.W_I_a_KV, self.W_I_b_KV
|
||||
W_a_Z, W_b_Z = self.W_I_a_Z, self.W_I_b_Z
|
||||
B_a, B_b = self.B_I_a, self.B_I_b
|
||||
else:
|
||||
W_a_KV, W_b_KV = self.W_a_KV, self.W_b_KV
|
||||
W_a_Z, W_b_Z = self.W_a_Z, self.W_b_Z
|
||||
B_a, B_b = self.B_a, self.B_b
|
||||
|
||||
# Current block projections: (m, c)
|
||||
C_a = _proj(cur_hidden, W_a_KV) # KV candidates
|
||||
Z_a = _proj(cur_hidden, W_a_Z) # gate logits
|
||||
|
||||
if prev_hidden is None:
|
||||
# Block 0: no previous block → softmax over m entries only
|
||||
# (paper pads Z^b with -inf, C^b with 0)
|
||||
Z_cat = Z_a + B_a # (m, c)
|
||||
S = F.softmax(Z_cat.float(), dim=0).to(self.dtype) # (m, c)
|
||||
C_out = (S * C_a).sum(dim=0) # (c,)
|
||||
else:
|
||||
# Blocks 1..: overlap with previous block
|
||||
C_b = _proj(prev_hidden, W_b_KV) # (m, c)
|
||||
Z_b = _proj(prev_hidden, W_b_Z) # (m, c)
|
||||
|
||||
# Concatenate along the position axis → (2m, c)
|
||||
Z_cat = torch.cat([Z_a + B_a, Z_b + B_b], dim=0)
|
||||
S = F.softmax(Z_cat.float(), dim=0).to(self.dtype) # (2m, c)
|
||||
|
||||
S_a, S_b = S[:m], S[m:] # each (m, c)
|
||||
C_out = (S_a * C_a).sum(dim=0) + (S_b * C_b).sum(dim=0) # (c,)
|
||||
|
||||
# Partial RoPE on the last rope_dim dims using the block's end position
|
||||
if cos_sin_cache is not None and self.rope > 0 and not for_indexer:
|
||||
pos_t = torch.tensor([block_end_pos], dtype=torch.long,
|
||||
device=cur_hidden.device)
|
||||
C_out = _apply_partial_rope(
|
||||
C_out.unsqueeze(0), pos_t, cos_sin_cache,
|
||||
self.nope, self.rope
|
||||
).squeeze(0)
|
||||
|
||||
return C_out # (c,) or (c_I,)
|
||||
|
||||
# ----------------------------------------------------------------
|
||||
# Prefill: all n tokens at once
|
||||
# ----------------------------------------------------------------
|
||||
|
||||
def prefill(
|
||||
self,
|
||||
hidden: torch.Tensor, # (n, d)
|
||||
cos_sin_cache: Optional[torch.Tensor], # (max_pos, rope_dim)
|
||||
start_pos: int = 0,
|
||||
state: Optional[CompressorState] = None,
|
||||
) -> CompressorState:
|
||||
"""
|
||||
Process all n tokens in one shot (prefill / context ingestion).
|
||||
|
||||
Tokens that don't fill a complete block of m are stored in
|
||||
state.tail_hidden for future incremental decode steps.
|
||||
|
||||
Returns an updated CompressorState.
|
||||
"""
|
||||
if state is None:
|
||||
state = CompressorState()
|
||||
|
||||
n, d = hidden.shape
|
||||
m = self.m
|
||||
|
||||
# If there are tail tokens from a previous call, prepend them
|
||||
if state.tail_hidden is not None and state.tail_hidden.shape[0] > 0:
|
||||
hidden = torch.cat([state.tail_hidden, hidden], dim=0)
|
||||
# The positions need to be adjusted accordingly
|
||||
n = hidden.shape[0]
|
||||
|
||||
n_complete_blocks = n // m
|
||||
n_tail = n % m
|
||||
|
||||
kv_list = []
|
||||
indexer_kv_list = []
|
||||
|
||||
prev_hidden = state.prev_hidden # None on first ever call
|
||||
|
||||
for i in range(n_complete_blocks):
|
||||
cur = hidden[i * m : (i + 1) * m] # (m, d)
|
||||
# Absolute position of the last token in this block
|
||||
block_end = start_pos + i * m + (m - 1)
|
||||
|
||||
c_kv = self._compress_block(
|
||||
cur, prev_hidden, cos_sin_cache, block_end, for_indexer=False
|
||||
)
|
||||
c_I = self._compress_block(
|
||||
cur, prev_hidden, None, block_end, for_indexer=True
|
||||
)
|
||||
|
||||
kv_list.append(c_kv)
|
||||
indexer_kv_list.append(c_I)
|
||||
prev_hidden = cur # this block becomes the "previous" for the next
|
||||
|
||||
# Accumulate into state
|
||||
new_kv = torch.stack(kv_list, dim=0) if kv_list else None # (n_blocks, c)
|
||||
new_I = torch.stack(indexer_kv_list, dim=0) if indexer_kv_list else None
|
||||
|
||||
if state.compressed_kv is None:
|
||||
state.compressed_kv = new_kv
|
||||
state.compressed_indexer_kv = new_I
|
||||
elif new_kv is not None:
|
||||
state.compressed_kv = torch.cat([state.compressed_kv, new_kv], dim=0)
|
||||
state.compressed_indexer_kv = torch.cat([state.compressed_indexer_kv, new_I], dim=0)
|
||||
|
||||
state.num_blocks += n_complete_blocks
|
||||
state.prev_hidden = prev_hidden
|
||||
state.tail_hidden = hidden[n_complete_blocks * m :] if n_tail > 0 else None
|
||||
|
||||
return state
|
||||
|
||||
# ----------------------------------------------------------------
|
||||
# Decode: single new token, incremental
|
||||
# ----------------------------------------------------------------
|
||||
|
||||
def decode_step(
|
||||
self,
|
||||
hidden_new: torch.Tensor, # (1, d) or (d,)
|
||||
cos_sin_cache: Optional[torch.Tensor],
|
||||
current_pos: int,
|
||||
state: CompressorState,
|
||||
) -> tuple[CompressorState, bool]:
|
||||
"""
|
||||
Ingest one new token into the state.
|
||||
|
||||
Returns (updated_state, new_block_committed).
|
||||
|
||||
new_block_committed is True when the new token completes a block
|
||||
and a new compressed entry has been appended to state.compressed_kv.
|
||||
The caller only needs to re-run the Lightning Indexer when True.
|
||||
"""
|
||||
h = hidden_new.reshape(1, self.d)
|
||||
|
||||
# Append to tail
|
||||
if state.tail_hidden is None:
|
||||
state.tail_hidden = h
|
||||
else:
|
||||
state.tail_hidden = torch.cat([state.tail_hidden, h], dim=0)
|
||||
|
||||
tail_len = state.tail_hidden.shape[0]
|
||||
|
||||
if tail_len < self.m:
|
||||
# Block not yet complete — nothing to compress
|
||||
return state, False
|
||||
|
||||
# Tail is exactly m tokens — compress
|
||||
assert tail_len == self.m, f"tail_len={tail_len} should equal m={self.m}"
|
||||
cur = state.tail_hidden # (m, d)
|
||||
block_end = current_pos # last token in block
|
||||
|
||||
c_kv = self._compress_block(
|
||||
cur, state.prev_hidden, cos_sin_cache, block_end, for_indexer=False
|
||||
)
|
||||
c_I = self._compress_block(
|
||||
cur, state.prev_hidden, None, block_end, for_indexer=True
|
||||
)
|
||||
|
||||
# Append to accumulated KV
|
||||
c_kv_2d = c_kv.unsqueeze(0) # (1, c)
|
||||
c_I_2d = c_I.unsqueeze(0) # (1, c_I)
|
||||
|
||||
if state.compressed_kv is None:
|
||||
state.compressed_kv = c_kv_2d
|
||||
state.compressed_indexer_kv = c_I_2d
|
||||
else:
|
||||
state.compressed_kv = torch.cat([state.compressed_kv, c_kv_2d], dim=0)
|
||||
state.compressed_indexer_kv = torch.cat([state.compressed_indexer_kv, c_I_2d], dim=0)
|
||||
|
||||
state.num_blocks += 1
|
||||
state.prev_hidden = cur # save for next block's overlap
|
||||
state.tail_hidden = None # clear tail
|
||||
|
||||
return state, True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HCA compressor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class HCACompressor:
|
||||
"""
|
||||
Heavily Compressed Attention token-level compressor.
|
||||
|
||||
Paper Section 2.3.2. Simpler than CSA:
|
||||
- No overlap: each block of m' tokens is self-contained.
|
||||
- No indexer: HCA uses dense MQA over all compressed entries,
|
||||
so no top-k selection is needed and there are no indexer keys.
|
||||
|
||||
For compressed block i:
|
||||
C = H · W_KV (n, c)
|
||||
Z = H · W_Z (n, c)
|
||||
S_i = softmax_row(Z[m'i:m'(i+1)] + B) (m', c)
|
||||
C^Comp_i = (S_i * C[m'i:m'(i+1)]).sum(0) (c,)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_dim: int = 7168,
|
||||
head_dim: int = 512,
|
||||
compress_ratio: int = 128, # m'
|
||||
nope_dim: int = 448,
|
||||
rope_dim: int = 64,
|
||||
device: str = "cuda",
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
):
|
||||
self.d = hidden_dim
|
||||
self.c = head_dim
|
||||
self.m = compress_ratio # m' in the paper
|
||||
self.nope = nope_dim
|
||||
self.rope = rope_dim
|
||||
self.device = device
|
||||
self.dtype = dtype
|
||||
|
||||
# W_KV: (d, c) W_Z: (d, c)
|
||||
self.W_KV = torch.empty(hidden_dim, head_dim, dtype=dtype, device=device)
|
||||
self.W_Z = torch.empty(hidden_dim, head_dim, dtype=dtype, device=device)
|
||||
# Positional bias B: (m', c)
|
||||
self.B = torch.empty(compress_ratio, head_dim, dtype=dtype, device=device)
|
||||
|
||||
def load_weights(self, W_KV, W_Z, B):
|
||||
def _cvt(t): return t.to(device=self.device, dtype=self.dtype)
|
||||
self.W_KV = _cvt(W_KV)
|
||||
self.W_Z = _cvt(W_Z)
|
||||
self.B = _cvt(B)
|
||||
|
||||
def _compress_block(
|
||||
self,
|
||||
block_hidden: torch.Tensor, # (m', d)
|
||||
cos_sin_cache: Optional[torch.Tensor],
|
||||
block_end_pos: int,
|
||||
) -> torch.Tensor:
|
||||
"""Compress one block of m' tokens → (c,)."""
|
||||
m = self.m
|
||||
assert block_hidden.shape[0] == m
|
||||
|
||||
C = _proj(block_hidden, self.W_KV) # (m', c)
|
||||
Z = _proj(block_hidden, self.W_Z) # (m', c)
|
||||
|
||||
S = F.softmax((Z + self.B).float(), dim=0).to(self.dtype) # (m', c)
|
||||
C_out = (S * C).sum(dim=0) # (c,)
|
||||
|
||||
if cos_sin_cache is not None and self.rope > 0:
|
||||
pos_t = torch.tensor([block_end_pos], dtype=torch.long,
|
||||
device=block_hidden.device)
|
||||
C_out = _apply_partial_rope(
|
||||
C_out.unsqueeze(0), pos_t, cos_sin_cache,
|
||||
self.nope, self.rope
|
||||
).squeeze(0)
|
||||
|
||||
return C_out
|
||||
|
||||
def prefill(
|
||||
self,
|
||||
hidden: torch.Tensor, # (n, d)
|
||||
cos_sin_cache: Optional[torch.Tensor],
|
||||
start_pos: int = 0,
|
||||
state: Optional[CompressorState] = None,
|
||||
) -> CompressorState:
|
||||
"""Process all n tokens (prefill). Tail stored for later decode."""
|
||||
if state is None:
|
||||
state = CompressorState()
|
||||
|
||||
n = hidden.shape[0]
|
||||
|
||||
if state.tail_hidden is not None and state.tail_hidden.shape[0] > 0:
|
||||
hidden = torch.cat([state.tail_hidden, hidden], dim=0)
|
||||
n = hidden.shape[0]
|
||||
|
||||
m = self.m
|
||||
n_complete = n // m
|
||||
n_tail = n % m
|
||||
|
||||
kv_list = []
|
||||
for i in range(n_complete):
|
||||
block = hidden[i * m : (i + 1) * m]
|
||||
block_end = start_pos + i * m + (m - 1)
|
||||
kv_list.append(self._compress_block(block, cos_sin_cache, block_end))
|
||||
|
||||
new_kv = torch.stack(kv_list, dim=0) if kv_list else None
|
||||
|
||||
if state.compressed_kv is None:
|
||||
state.compressed_kv = new_kv
|
||||
elif new_kv is not None:
|
||||
state.compressed_kv = torch.cat([state.compressed_kv, new_kv], dim=0)
|
||||
|
||||
state.num_blocks += n_complete
|
||||
state.tail_hidden = hidden[n_complete * m :] if n_tail > 0 else None
|
||||
# HCA has no prev_hidden needed (no overlap)
|
||||
|
||||
return state
|
||||
|
||||
def decode_step(
|
||||
self,
|
||||
hidden_new: torch.Tensor, # (1, d)
|
||||
cos_sin_cache: Optional[torch.Tensor],
|
||||
current_pos: int,
|
||||
state: CompressorState,
|
||||
) -> tuple[CompressorState, bool]:
|
||||
"""Ingest one token. Returns (state, new_block_committed)."""
|
||||
h = hidden_new.reshape(1, self.d)
|
||||
|
||||
state.tail_hidden = h if state.tail_hidden is None else \
|
||||
torch.cat([state.tail_hidden, h], dim=0)
|
||||
|
||||
if state.tail_hidden.shape[0] < self.m:
|
||||
return state, False
|
||||
|
||||
# Full block ready
|
||||
block = state.tail_hidden # (m', d)
|
||||
c_kv = self._compress_block(block, cos_sin_cache, current_pos)
|
||||
c_kv_2d = c_kv.unsqueeze(0)
|
||||
|
||||
state.compressed_kv = c_kv_2d if state.compressed_kv is None else \
|
||||
torch.cat([state.compressed_kv, c_kv_2d], dim=0)
|
||||
|
||||
state.num_blocks += 1
|
||||
state.tail_hidden = None
|
||||
|
||||
return state, True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Quick smoke test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
torch.manual_seed(42)
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = torch.bfloat16
|
||||
|
||||
# ── V4-Pro dims ─────────────────────────────────────────────────
|
||||
D, C, M, C_I = 7168, 512, 4, 128
|
||||
M_HCA = 128
|
||||
NOPE, ROPE = 448, 64
|
||||
MAX_POS = 4096
|
||||
|
||||
cos_sin_cache = torch.randn(MAX_POS, ROPE, dtype=dtype, device=device)
|
||||
|
||||
# ── CSA ─────────────────────────────────────────────────────────
|
||||
print("=== CSA compressor ===")
|
||||
csa = CSACompressor(D, C, M, C_I, num_indexer_heads=64,
|
||||
nope_dim=NOPE, rope_dim=ROPE, device=device, dtype=dtype)
|
||||
|
||||
# Prefill 20 tokens (5 complete blocks, 0 tail)
|
||||
n_prefill = 20
|
||||
h_prefill = torch.randn(n_prefill, D, dtype=dtype, device=device)
|
||||
state = csa.prefill(h_prefill, cos_sin_cache, start_pos=0)
|
||||
|
||||
print(f" After prefill {n_prefill} tokens:")
|
||||
print(f" num_blocks: {state.num_blocks}") # 5
|
||||
print(f" compressed_kv: {state.compressed_kv.shape}") # (5, 512)
|
||||
print(f" indexer_kv: {state.compressed_indexer_kv.shape}") # (5, 128)
|
||||
print(f" tail_len: {0 if state.tail_hidden is None else state.tail_hidden.shape[0]}")
|
||||
|
||||
# Prefill 6 more (1 complete block + 2 tail)
|
||||
h2 = torch.randn(6, D, dtype=dtype, device=device)
|
||||
state = csa.prefill(h2, cos_sin_cache, start_pos=n_prefill, state=state)
|
||||
print(f"\n After prefill 6 more:")
|
||||
print(f" num_blocks: {state.num_blocks}") # 6
|
||||
print(f" compressed_kv: {state.compressed_kv.shape}") # (6, 512)
|
||||
print(f" tail_len: {state.tail_hidden.shape[0]}") # 2
|
||||
|
||||
# Decode 2 tokens (fills tail → 1 new block)
|
||||
for tok_i in range(2):
|
||||
h_tok = torch.randn(1, D, dtype=dtype, device=device)
|
||||
pos = n_prefill + 6 + tok_i
|
||||
state, committed = csa.decode_step(h_tok, cos_sin_cache, pos, state)
|
||||
print(f" decode tok {tok_i}: committed={committed}")
|
||||
print(f" num_blocks now: {state.num_blocks}") # 7
|
||||
|
||||
# ── HCA ─────────────────────────────────────────────────────────
|
||||
print("\n=== HCA compressor ===")
|
||||
hca = HCACompressor(D, C, M_HCA, nope_dim=NOPE, rope_dim=ROPE,
|
||||
device=device, dtype=dtype)
|
||||
|
||||
n_hca = 384 # 3 complete blocks + 0 tail
|
||||
h_hca = torch.randn(n_hca, D, dtype=dtype, device=device)
|
||||
hca_state = hca.prefill(h_hca, cos_sin_cache, start_pos=0)
|
||||
print(f" After prefill {n_hca} tokens:")
|
||||
print(f" num_blocks: {hca_state.num_blocks}") # 3
|
||||
print(f" compressed_kv: {hca_state.compressed_kv.shape}") # (3, 512)
|
||||
|
||||
# Decode 128 tokens → exactly 1 new HCA block
|
||||
for tok_i in range(M_HCA):
|
||||
h_tok = torch.randn(1, D, dtype=dtype, device=device)
|
||||
pos = n_hca + tok_i
|
||||
hca_state, committed = hca.decode_step(h_tok, cos_sin_cache, pos, hca_state)
|
||||
|
||||
print(f" After decode {M_HCA} tokens:")
|
||||
print(f" num_blocks: {hca_state.num_blocks}") # 4
|
||||
print(f" compressed_kv: {hca_state.compressed_kv.shape}") # (4, 512)
|
||||
|
||||
# ── Correctness sanity: prefill == incremental decode ───────────
|
||||
print("\n=== Equivalence check: prefill vs incremental decode ===")
|
||||
csa2 = CSACompressor(D, C, M, C_I, nope_dim=NOPE, rope_dim=ROPE,
|
||||
device=device, dtype=dtype)
|
||||
# Copy weights
|
||||
for attr in ("W_a_KV","W_b_KV","W_a_Z","W_b_Z","B_a","B_b",
|
||||
"W_I_a_KV","W_I_b_KV","W_I_a_Z","W_I_b_Z","B_I_a","B_I_b"):
|
||||
setattr(csa2, attr, getattr(csa, attr))
|
||||
|
||||
n_check = 8
|
||||
h_check = torch.randn(n_check, D, dtype=dtype, device=device)
|
||||
|
||||
# Batch prefill
|
||||
s_batch = csa2.prefill(h_check, cos_sin_cache, start_pos=0)
|
||||
|
||||
# Token-by-token decode
|
||||
s_incr = CompressorState()
|
||||
for i in range(n_check):
|
||||
s_incr, _ = csa2.decode_step(h_check[i:i+1], cos_sin_cache, i, s_incr)
|
||||
|
||||
if s_batch.compressed_kv is not None and s_incr.compressed_kv is not None:
|
||||
max_diff = (s_batch.compressed_kv - s_incr.compressed_kv).abs().max().item()
|
||||
print(f" max |prefill - decode| on compressed_kv: {max_diff:.6f}")
|
||||
assert max_diff < 1e-3, "Mismatch between prefill and decode paths!"
|
||||
print(" PASSED")
|
||||
else:
|
||||
print(" (no complete blocks produced in 8 tokens with m=4 — increase n_check)")
|
||||
|
||||
print("\nAll checks done.")
|
||||
377
tests/test_stage_b_v28.py
Normal file
377
tests/test_stage_b_v28.py
Normal file
@@ -0,0 +1,377 @@
|
||||
"""
|
||||
Stage B v28: Bug 4 fix — PV MMA with correct (128,64) tiler.
|
||||
|
||||
Key fixes over v27:
|
||||
1. PV MMA created with pv_mma_tiler[:2] = (128, 64) instead of mma_tiler_mn = (128, 128)
|
||||
-> A-fragment N_MMA=64 matches packed P layout (64 FP32 columns)
|
||||
2. epi_tile = pv_mma_tiler[:2] like FMHA (not from QK cta_tile)
|
||||
3. PV ACCUMULATE=False on first tile (FMHA: kphase_idx != 0)
|
||||
4. gC (output) tiling uses pv_mma_tiler for (M, head_dim) not qk_mma_tiler
|
||||
5. cta_tile_shape_mnk set to PV cta tile before epilogue
|
||||
"""
|
||||
import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.pipeline as pipeline
|
||||
from cutlass.cute.nvgpu import cpasync, tcgen05
|
||||
from cutlass import Float32, BFloat16, Int32, Boolean, const_expr
|
||||
from cutlass.utils import LayoutEnum
|
||||
from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset
|
||||
import cuda.bindings.driver as cuda
|
||||
|
||||
|
||||
class StageBIdentitySoftmax:
|
||||
def __init__(self, mma_tiler_mn, use_2cta_instrs=False, use_tma_store=True):
|
||||
self.acc_dtype = Float32; self.qk_acc_dtype = Float32
|
||||
self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16
|
||||
self.use_2cta_instrs = use_2cta_instrs; self.use_tma_store = use_tma_store
|
||||
self.mma_tiler_mn = mma_tiler_mn; self.mma_tiler = (*mma_tiler_mn, 1)
|
||||
self.cluster_shape_mn = (1, 1)
|
||||
self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE
|
||||
self.epilogue_warp_id = (0, 1, 2, 3)
|
||||
self.mma_warp_id = 4; self.tma_warp_id = 5
|
||||
self.threads_per_cta = 192
|
||||
self.epilog_sync_bar_id = 1; self.tmem_alloc_sync_bar_id = 2; self.tmem_dealloc_sync_bar_id = 3
|
||||
self.num_c_stage = 2
|
||||
|
||||
def _setup(self, qk_mma, pv_mma):
|
||||
qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2])
|
||||
self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4)
|
||||
self.pv_mma_tiler = (self.qk_mma_tiler[0], self.qk_mma_tiler[2], self.qk_mma_tiler[1])
|
||||
self.mma_tiler = self.qk_mma_tiler
|
||||
print(f"[v28] qk_mma_tiler = {self.qk_mma_tiler}")
|
||||
print(f"[v28] pv_mma_tiler = {self.pv_mma_tiler}")
|
||||
|
||||
self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,))
|
||||
|
||||
# cta_tile_shape_mnk from PV for epilogue
|
||||
self.cta_tile_shape_mnk = (
|
||||
self.pv_mma_tiler[0] // cute.size(pv_mma.thr_id.shape),
|
||||
self.pv_mma_tiler[1],
|
||||
self.pv_mma_tiler[2],
|
||||
)
|
||||
# FMHA: epi_tile from PV cta_tile, not QK
|
||||
self.epi_tile = utils.sm100.compute_epilogue_tile_shape(
|
||||
self.cta_tile_shape_mnk, self.use_2cta_instrs, self.c_layout, self.o_dtype)
|
||||
self.num_ab_stage = 1; self.num_acc_stage = 1
|
||||
|
||||
self.a_smem_s = utils.sm100.make_smem_layout_a(qk_mma, self.mma_tiler, self.q_dtype, 1)
|
||||
self.b_smem_s = utils.sm100.make_smem_layout_b(qk_mma, self.mma_tiler, self.b_dtype, 1)
|
||||
self.v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, self.b_dtype, 1)
|
||||
self.p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, self.q_dtype, 1)
|
||||
self.c_smem_s = utils.sm100.make_smem_layout_epi(self.o_dtype, self.c_layout, self.epi_tile, 2)
|
||||
|
||||
qk_thr = qk_mma.get_slice(0)
|
||||
qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2])
|
||||
tStS = qk_thr.make_fragment_C(qk_acc_shape)
|
||||
s_cols = find_tmem_tensor_col_offset(tStS)
|
||||
|
||||
pv_thr = pv_mma.get_slice(0)
|
||||
pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2])
|
||||
tOtO = pv_thr.make_fragment_C(pv_acc_shape)
|
||||
o_cols = find_tmem_tensor_col_offset(tOtO)
|
||||
|
||||
self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width
|
||||
self.tmem_s0_offset = 0
|
||||
self.tmem_p0_offset = 32
|
||||
self.tmem_o0_offset = s_cols
|
||||
|
||||
tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage))
|
||||
tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage))
|
||||
self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCtS_fake, tCtO_fake], arch="sm_100")
|
||||
|
||||
a_smem = cute.slice_(self.a_smem_s, (None, None, None, 0))
|
||||
b_smem = cute.slice_(self.b_smem_s, (None, None, None, 0))
|
||||
self.num_tma_load_bytes = (
|
||||
cute.size_in_bytes(self.q_dtype, a_smem) + cute.size_in_bytes(self.b_dtype, b_smem)
|
||||
) * cute.size(qk_mma.thr_id.shape)
|
||||
|
||||
@cute.jit
|
||||
def __call__(self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, c: cute.Tensor, stream: cuda.CUstream):
|
||||
self.q_dtype = q.element_type; self.b_dtype = k.element_type
|
||||
self.o_dtype = c.element_type; self.c_dtype = self.o_dtype
|
||||
self.a_major = LayoutEnum.from_tensor(q).mma_major_mode()
|
||||
self.b_major = LayoutEnum.from_tensor(k).mma_major_mode()
|
||||
self.v_major = LayoutEnum.from_tensor(v).mma_major_mode()
|
||||
self.c_layout = LayoutEnum.from_tensor(c)
|
||||
|
||||
qk_mma = utils.sm100.make_trivial_tiled_mma(
|
||||
self.q_dtype, self.b_dtype, self.a_major, self.b_major,
|
||||
self.qk_acc_dtype, self.cta_group, self.mma_tiler_mn, tcgen05.OperandSource.SMEM)
|
||||
|
||||
# Compute pv_mma_tiler[:2] BEFORE creating pv_mma
|
||||
qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2])
|
||||
qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4)
|
||||
pv_mma_tiler = (qk_mma_tiler[0], qk_mma_tiler[2], qk_mma_tiler[1])
|
||||
|
||||
# BUG 4 FIX: pv_mma_tiler[:2] = (128, 64) not (128, 128)
|
||||
pv_mma = utils.sm100.make_trivial_tiled_mma(
|
||||
self.q_dtype, self.b_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major,
|
||||
self.qk_acc_dtype, self.cta_group, pv_mma_tiler[:2], tcgen05.OperandSource.TMEM)
|
||||
|
||||
self._setup(qk_mma, pv_mma)
|
||||
|
||||
q_smem = cute.slice_(self.a_smem_s, (None, None, None, 0))
|
||||
k_smem = cute.slice_(self.b_smem_s, (None, None, None, 0))
|
||||
v_smem = cute.slice_(self.v_smem_s, (None, None, None, 0))
|
||||
|
||||
tma_q, tma_tq = cute.nvgpu.make_tiled_tma_atom_A(
|
||||
utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, qk_mma.thr_id),
|
||||
q, q_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape)
|
||||
tma_k, tma_tk = cute.nvgpu.make_tiled_tma_atom_B(
|
||||
utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, qk_mma.thr_id),
|
||||
k, k_smem, self.mma_tiler, qk_mma, self.cluster_layout_vmnk.shape)
|
||||
tma_v, tma_tv = cute.nvgpu.make_tiled_tma_atom_B(
|
||||
utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id),
|
||||
v, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape)
|
||||
epi_smem = cute.select(self.c_smem_s, mode=[0, 1])
|
||||
tma_c, tma_tc = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile)
|
||||
|
||||
self._kernel(qk_mma, pv_mma, tma_q, tma_tq, tma_k, tma_tk, tma_v, tma_tv,
|
||||
tma_c, tma_tc, self.cluster_layout_vmnk,
|
||||
self.a_smem_s, self.b_smem_s, self.v_smem_s, self.p_tmem_s, self.c_smem_s, self.epi_tile
|
||||
).launch(grid=(1,1,1), block=[self.threads_per_cta,1,1], stream=stream)
|
||||
|
||||
@cute.kernel
|
||||
def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV,
|
||||
tma_c, mC, cl_vmnk, a_smem_s, b_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile):
|
||||
warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
use_2cta = cute.size(qk_mma.thr_id.shape) == 2
|
||||
|
||||
if warp_idx == self.tma_warp_id:
|
||||
cpasync.prefetch_descriptor(tma_q); cpasync.prefetch_descriptor(tma_k)
|
||||
cpasync.prefetch_descriptor(tma_v); cpasync.prefetch_descriptor(tma_c)
|
||||
|
||||
@cute.struct
|
||||
class SS:
|
||||
ab_bar: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
|
||||
mma_si_bar: cute.struct.MemRange[cutlass.Int64, 2]
|
||||
acc_bar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
|
||||
tmem_dealloc: cutlass.Int64
|
||||
holding: cutlass.Int32
|
||||
|
||||
smem = utils.SmemAllocator(); st = smem.allocate(SS)
|
||||
|
||||
ab_p, ab_c = pipeline.PipelineTmaUmma.create(
|
||||
barrier_storage=st.ab_bar.data_ptr(), num_stages=self.num_ab_stage,
|
||||
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
|
||||
consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
|
||||
tx_count=self.num_tma_load_bytes, cta_layout_vmnk=cl_vmnk, defer_sync=True
|
||||
).make_participants()
|
||||
|
||||
mma_si_prod, mma_si_cons = pipeline.PipelineUmmaAsync.create(
|
||||
barrier_storage=st.mma_si_bar.data_ptr(), num_stages=1,
|
||||
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
|
||||
consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)),
|
||||
).make_participants()
|
||||
|
||||
acc_pipe = pipeline.PipelineUmmaAsync.create(
|
||||
barrier_storage=st.acc_bar.data_ptr(), num_stages=self.num_acc_stage,
|
||||
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
|
||||
consumer_group=pipeline.CooperativeGroup(
|
||||
pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta else 1)),
|
||||
cta_layout_vmnk=cl_vmnk, defer_sync=True)
|
||||
|
||||
tmem_bar = pipeline.NamedBarrier(barrier_id=self.tmem_alloc_sync_bar_id,
|
||||
num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)))
|
||||
tmem = utils.TmemAllocator(st.holding.ptr, barrier_for_retrieve=tmem_bar,
|
||||
allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta,
|
||||
two_cta_tmem_dealloc_mbar_ptr=st.tmem_dealloc.ptr)
|
||||
|
||||
pipeline.pipeline_init_arrive(cluster_shape_mn=cl_vmnk, is_relaxed=True)
|
||||
|
||||
sQ = smem.allocate_tensor(element_type=self.q_dtype, layout=a_smem_s.outer, byte_alignment=128, swizzle=a_smem_s.inner)
|
||||
sK = smem.allocate_tensor(element_type=self.b_dtype, layout=b_smem_s.outer, byte_alignment=128, swizzle=b_smem_s.inner)
|
||||
sV = smem.allocate_tensor(element_type=self.b_dtype, layout=v_smem_s.outer, byte_alignment=128, swizzle=v_smem_s.inner)
|
||||
sC = smem.allocate_tensor(element_type=self.o_dtype, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner)
|
||||
|
||||
gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None))
|
||||
gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None))
|
||||
# gC uses pv_mma_tiler output shape (M, head_dim)
|
||||
gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None, None, 0)), (None, None, None))
|
||||
k_cnt = cute.size(gQ, mode=[3])
|
||||
|
||||
qk_thr = qk_mma.get_slice(0)
|
||||
pv_thr = pv_mma.get_slice(0)
|
||||
tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK)
|
||||
# C partitioned by QK thr (for TMA store, FMHA uses the same pattern)
|
||||
tCgC = qk_thr.partition_C(gC)
|
||||
a_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape)
|
||||
tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_lay, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3))
|
||||
b_lay = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape)
|
||||
tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3))
|
||||
tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]
|
||||
|
||||
gV = cute.local_tile(mV, cute.slice_(self.pv_mma_tiler, (0,None,None)), (None,None,None))
|
||||
tCgV = pv_thr.partition_B(gV)
|
||||
tVsV, tVgV = cpasync.tma_partition(tma_v, 0, b_lay, cute.group_modes(sV,0,3), cute.group_modes(tCgV,0,3))
|
||||
tVgV = tVgV[(None,0,None,0)]
|
||||
|
||||
tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK)
|
||||
tCrV = pv_mma.make_fragment_B(sV)
|
||||
|
||||
qk_acc_shape = qk_thr.partition_shape_C(self.qk_mma_tiler[:2])
|
||||
tStS = qk_thr.make_fragment_C(qk_acc_shape)
|
||||
tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout)
|
||||
|
||||
pv_acc_shape = pv_thr.partition_shape_C(self.pv_mma_tiler[:2])
|
||||
tOtO = pv_thr.make_fragment_C(pv_acc_shape)
|
||||
tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout)
|
||||
|
||||
tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer)
|
||||
tOrP_base = pv_thr.make_fragment_A(tP)
|
||||
tOrP = tOrP_base[(None, None, None, 0)]
|
||||
tOrP0 = cute.make_tensor(
|
||||
tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset,
|
||||
tOrP.layout)
|
||||
|
||||
tCtS_fake = qk_mma.make_fragment_C(cute.append(qk_acc_shape, self.num_acc_stage))
|
||||
tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_acc_shape, self.num_acc_stage))
|
||||
|
||||
pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk)
|
||||
|
||||
# ===== TMA LOAD WARP =====
|
||||
if warp_idx == self.tma_warp_id:
|
||||
ab_p.reset(); peek = ab_p.try_acquire()
|
||||
for kt in cutlass.range(k_cnt, unroll=1):
|
||||
h = ab_p.acquire_and_advance(peek)
|
||||
cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier)
|
||||
cute.copy(tma_k, tBgK[(None,h.count)], tBsK[(None,h.index)], tma_bar_ptr=h.barrier)
|
||||
cute.copy(tma_v, tVgV[(None,h.count)], tVsV[(None,h.index)], tma_bar_ptr=h.barrier)
|
||||
peek = cutlass.Boolean(1)
|
||||
if h.count+1<k_cnt: peek = ab_p.try_acquire()
|
||||
ab_p.tail()
|
||||
|
||||
# ===== MMA WARP =====
|
||||
if warp_idx == self.mma_warp_id:
|
||||
tmem.wait_for_alloc()
|
||||
ab_c.reset(); peek = ab_c.try_wait()
|
||||
|
||||
s0_handle = mma_si_prod.acquire_and_advance()
|
||||
|
||||
acc_prod_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage)
|
||||
acc_pipe.producer_acquire(acc_prod_st)
|
||||
|
||||
# QK MMA
|
||||
qk_mma.set(tcgen05.Field.ACCUMULATE, False)
|
||||
for kt in range(k_cnt):
|
||||
h = ab_c.wait_and_advance(peek)
|
||||
nblk = cute.size(tCrQ, mode=[2])
|
||||
for kb in cutlass.range(nblk, unroll_full=True):
|
||||
cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,h.index)], tCrK[(None,None,kb,h.index)], tStS0)
|
||||
qk_mma.set(tcgen05.Field.ACCUMULATE, True)
|
||||
h.release(); peek = cutlass.Boolean(1)
|
||||
if h.count+1<k_cnt: peek = ab_c.try_wait()
|
||||
|
||||
cute.arch.fence_view_async_tmem_store()
|
||||
s0_handle.commit()
|
||||
|
||||
# Wait for softmax to complete
|
||||
s0_handle = mma_si_prod.acquire_and_advance()
|
||||
|
||||
# PV MMA — ACCUMULATE=False for first kphase (Bug 3 fix, FMHA: kphase_idx != 0)
|
||||
pv_mma.set(tcgen05.Field.ACCUMULATE, False)
|
||||
tCrV_s = tCrV[(None, None, None, 0)]
|
||||
nblk_pv = cute.size(tOrP0, mode=[2])
|
||||
for kb in cutlass.range(nblk_pv, unroll_full=True):
|
||||
cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV_s[(None,None,kb)], tOtO0)
|
||||
pv_mma.set(tcgen05.Field.ACCUMULATE, True)
|
||||
|
||||
acc_pipe.producer_commit(acc_prod_st)
|
||||
acc_prod_st.advance()
|
||||
acc_pipe.producer_tail(acc_prod_st)
|
||||
|
||||
# ===== EPILOGUE WARPS =====
|
||||
if warp_idx < self.mma_warp_id:
|
||||
tmem.allocate(self.num_tmem_alloc_cols)
|
||||
tmem.wait_for_alloc()
|
||||
tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype)
|
||||
sfw_idx = tidx % (32 * len(self.epilogue_warp_id))
|
||||
|
||||
tmem_load_atom = cute.make_copy_atom(
|
||||
tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
|
||||
tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS0)
|
||||
thr_load = tiled_tmem_load.get_slice(sfw_idx)
|
||||
tTMEM_LOADtS = thr_load.partition_S(tStS0)
|
||||
cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1]))
|
||||
tScS = qk_thr.partition_C(cS)
|
||||
tTMEM_LOADcS = thr_load.partition_D(tScS)
|
||||
|
||||
# Packed P store — composition collapses 128 FP32 cols to 64 packed FP32 cols
|
||||
tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, self.tilePlikeFP32)))
|
||||
tStS_P = cute.make_tensor(tStS.iterator + self.tmem_p0_offset, tStS_P_layout)
|
||||
tmem_store_atom = cute.make_copy_atom(
|
||||
tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
|
||||
tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P)
|
||||
thr_store = tiled_tmem_store.get_slice(sfw_idx)
|
||||
tTMEM_STOREtS_x4 = thr_store.partition_D(tStS_P)
|
||||
tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, self.tilePlikeFP32)))
|
||||
tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout)
|
||||
tTMEM_STOREcS = thr_store.partition_S(tScS_P)
|
||||
|
||||
si_handle = mma_si_cons.wait_and_advance()
|
||||
|
||||
tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype)
|
||||
cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS)
|
||||
|
||||
tTMEM_STORErS_x4 = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.qk_acc_dtype)
|
||||
tTMEM_STORErS_x4_e = cute.make_tensor(
|
||||
cute.recast_ptr(tTMEM_STORErS_x4.iterator, dtype=self.q_dtype),
|
||||
tTMEM_LOADrS.layout)
|
||||
|
||||
frg_cnt = 4
|
||||
frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt
|
||||
tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile))
|
||||
tTMEM_STORErS_x4_e_frg = cute.logical_divide(
|
||||
tTMEM_STORErS_x4_e, cute.make_layout(frg_tile))
|
||||
for j in range(frg_cnt):
|
||||
s_vec = tTMEM_LOADrS_frg[None, j].load()
|
||||
tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype))
|
||||
|
||||
cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4)
|
||||
cute.arch.fence_view_async_tmem_store()
|
||||
si_handle.release()
|
||||
|
||||
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)
|
||||
acc_cons_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage)
|
||||
c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))
|
||||
c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp)
|
||||
acc_cons_st = utils.gemm.sm100.epilogue_tma_store(
|
||||
self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC,
|
||||
epi_tile, 0, const_expr(lambda x: x), (0,0,0), acc_cons_st, acc_pipe, c_pipe)
|
||||
c_pipe.producer_tail()
|
||||
tmem.relinquish_alloc_permit()
|
||||
tmem.free(tmem_ptr)
|
||||
|
||||
|
||||
def test():
|
||||
torch.manual_seed(42)
|
||||
m, n, head_dim = 128, 128, 64
|
||||
q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda')
|
||||
k = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda')
|
||||
v_base = torch.randn(head_dim, n, dtype=torch.bfloat16, device='cuda')
|
||||
v = v_base.as_strided((head_dim, n), (1, head_dim)).unsqueeze(-1)
|
||||
c = torch.zeros(m, head_dim, 1, dtype=torch.bfloat16, device='cuda')
|
||||
|
||||
qf = q[:,:,0].float(); kf = k[:,:,0].float(); vf = v_base.float()
|
||||
ref = qf @ kf.T @ vf.T
|
||||
|
||||
import cutlass.torch as ct
|
||||
mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
|
||||
mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
|
||||
mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v))
|
||||
mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))
|
||||
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
||||
kernel = StageBIdentitySoftmax(mma_tiler_mn=(128, 128), use_2cta_instrs=False, use_tma_store=True)
|
||||
print('Compiling...', flush=True)
|
||||
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream)
|
||||
print('Running...', flush=True)
|
||||
compiled(mQ, mK, mV, mC, stream)
|
||||
torch.cuda.synchronize()
|
||||
out = c[:,:,0].float()
|
||||
cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item()
|
||||
max_err = (out - ref).abs().max().item()
|
||||
print('Stage B v28: Bug 4 fix — PV MMA (128,64), epi from PV, ACCUMULATE=False')
|
||||
print(' Cosine: {:.6f}, Max error: {:.6f}'.format(cos, max_err))
|
||||
print(' {}'.format('PASS' if cos >= 0.99 else 'FAIL'))
|
||||
|
||||
if __name__ == '__main__':
|
||||
test()
|
||||
Reference in New Issue
Block a user