Restructure: cutedsl/ -> dsv4/ with proper layering
- Split bridge.py -> ops/quantize.py, ops/layouts.py, ops/gemm_runner.py - Renamed classes: CuTeDSLNvfp4Linear -> Nvfp4Linear, etc. - Moved kernel code to dsv4/kernels/ (gemm, attention, compressor, decode, cuda) - Moved PyTorch bridges to dsv4/ops/ - Moved nn.Module layers to dsv4layers/ - Moved reference implementations to dsv4/reference/ - Moved vendored CUTLASS code to vendored/ - Archived ~190 debug tests to tests/archive/ - Kept ~15 canonical tests in tests/unit/ - Updated all import paths - Added stubs for future components (model/, cache/, loader/) - Updated pyproject.toml: dsv4-inference package name
This commit is contained in:
0
dsv4/reference/__init__.py
Normal file
0
dsv4/reference/__init__.py
Normal file
247
dsv4/reference/attention.py
Normal file
247
dsv4/reference/attention.py
Normal file
@@ -0,0 +1,247 @@
|
||||
"""
|
||||
DeepSeek-V4 Blackwell Attention — Our own kernel.
|
||||
|
||||
Replaces vLLM's broken FlashMLA Blackwell path with a proper KV cache-based
|
||||
attention pipeline. Does NOT depend on FlashMLA, fp8_ds_mla, or any vLLM
|
||||
fused CUDA kernel.
|
||||
|
||||
Architecture:
|
||||
- KV: (T, HD=512) single head latent, shared across all 128 Q heads
|
||||
- KV Cache: fp8_e4m3 paged cache with per-token inverse scale
|
||||
- RoPE: GPT-J style, applied to Q and KV before caching
|
||||
- Attention: BF16 (NVFP4 is too lossy for Q×K^T, cosine 0.86)
|
||||
- CSA/HCA: Compressed KV for sparse attention (compress_ratio 4 or 128)
|
||||
- SWA: Sliding window attention (compress_ratio 0/1)
|
||||
|
||||
Pipeline:
|
||||
Prefill:
|
||||
1. hidden → q_a_proj → q_norm → q_b_proj → (T, NH, HD) → RoPE on Q
|
||||
2. hidden → kv_proj → kv_norm → (T, HD) → RoPE → fp8 quant → write to paged cache
|
||||
3. Read all cached KV → BF16 causal attention → output
|
||||
|
||||
Decode:
|
||||
1. Same projections as prefill
|
||||
2. Write new KV to cache
|
||||
3. Read ALL cached KV → BF16 attention (1 query vs N KVs) → output
|
||||
|
||||
Output:
|
||||
1. inverse RoPE on attention output
|
||||
2. o_a: BMM with wo_a (BF16)
|
||||
3. o_b: NVFP4 GEMM with wo_b
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def apply_gptj_rope(x, positions, cos_sin_cache, nope_dim, rope_dim):
|
||||
"""Apply GPT-J style RoPE. Works on (T, HD) or (T, NH, HD)."""
|
||||
if rope_dim == 0 or x.numel() == 0:
|
||||
return x
|
||||
half = rope_dim // 2
|
||||
cos = cos_sin_cache[positions, :half].to(x.dtype)
|
||||
sin = cos_sin_cache[positions, half:2 * half].to(x.dtype)
|
||||
if x.dim() == 3:
|
||||
cos = cos.unsqueeze(1)
|
||||
sin = sin.unsqueeze(1)
|
||||
x_rope = x[..., nope_dim:].clone()
|
||||
even = x_rope[..., 0::2]
|
||||
odd = x_rope[..., 1::2]
|
||||
out = x.clone()
|
||||
out[..., nope_dim:][..., 0::2] = even * cos - odd * sin
|
||||
out[..., nope_dim:][..., 1::2] = even * sin + odd * cos
|
||||
return out
|
||||
|
||||
|
||||
def apply_inv_gptj_rope(x, positions, cos_sin_cache, nope_dim, rope_dim):
|
||||
"""Inverse GPT-J RoPE (sin → -sin)."""
|
||||
if rope_dim == 0 or x.numel() == 0:
|
||||
return x
|
||||
half = rope_dim // 2
|
||||
cos = cos_sin_cache[positions, :half].to(x.dtype)
|
||||
sin = cos_sin_cache[positions, half:2 * half].to(x.dtype)
|
||||
if x.dim() == 3:
|
||||
cos = cos.unsqueeze(1)
|
||||
sin = sin.unsqueeze(1)
|
||||
x_rope = x[..., nope_dim:].clone()
|
||||
even = x_rope[..., 0::2]
|
||||
odd = x_rope[..., 1::2]
|
||||
out = x.clone()
|
||||
out[..., nope_dim:][..., 0::2] = even * cos + odd * sin
|
||||
out[..., nope_dim:][..., 1::2] = -even * sin + odd * cos
|
||||
return out
|
||||
|
||||
|
||||
# ── KV Cache Operations ──────────────────────────────────────────────
|
||||
|
||||
def kv_quantize_fp8(kv_bf16):
|
||||
"""BF16 KV → fp8_e4m3 with per-token inverse scale."""
|
||||
amax = kv_bf16.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
|
||||
fp8_max = torch.tensor(448.0, dtype=torch.float32, device=kv_bf16.device)
|
||||
scale = fp8_max / amax
|
||||
kv_fp8 = (kv_bf16.float() * scale).to(torch.float8_e4m3fn)
|
||||
inv_scale = (amax / fp8_max).to(torch.bfloat16)
|
||||
return kv_fp8, inv_scale
|
||||
|
||||
|
||||
def kv_dequantize_fp8(kv_fp8, inv_scale):
|
||||
"""fp8 KV → BF16."""
|
||||
return (kv_fp8.to(torch.bfloat16) * inv_scale).to(torch.bfloat16)
|
||||
|
||||
|
||||
def paged_kv_write(kv_data, slot_mapping, cache, block_size):
|
||||
"""Write KV into paged cache. Works for fp8 or bf16.
|
||||
|
||||
kv_data: (T, D) tensor to write
|
||||
slot_mapping: (T,) slot indices
|
||||
cache: (num_blocks, block_size, D) cache tensor
|
||||
"""
|
||||
for t in range(kv_data.shape[0]):
|
||||
slot = slot_mapping[t].item()
|
||||
block_idx = slot // block_size
|
||||
offset = slot % block_size
|
||||
if block_idx < cache.shape[0] and offset < cache.shape[1]:
|
||||
cache[block_idx, offset] = kv_data[t]
|
||||
|
||||
|
||||
def paged_kv_read(slot_mapping, cache, block_size, num_tokens, head_dim):
|
||||
"""Read KV from paged cache."""
|
||||
device = cache.device
|
||||
kv = torch.zeros(num_tokens, head_dim, dtype=cache.dtype, device=device)
|
||||
for t in range(num_tokens):
|
||||
slot = slot_mapping[t].item()
|
||||
block_idx = slot // block_size
|
||||
offset = slot % block_size
|
||||
if block_idx < cache.shape[0] and offset < cache.shape[1]:
|
||||
kv[t] = cache[block_idx, offset]
|
||||
return kv
|
||||
|
||||
|
||||
# ── Attention ─────────────────────────────────────────────────────────
|
||||
|
||||
def causal_prefill_attention(q, kv, scale):
|
||||
"""Full causal self-attention for prefill. q: (T, NH, HD), kv: (T, HD)."""
|
||||
T, NH, HD = q.shape
|
||||
q_t = q.permute(1, 0, 2) # (NH, T, HD)
|
||||
kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) # (NH, T, HD)
|
||||
out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=True, scale=scale)
|
||||
return out.permute(1, 0, 2) # (T, NH, HD)
|
||||
|
||||
|
||||
def decode_attention(q, kv, scale):
|
||||
"""Decode attention: 1 query vs N cached KVs.
|
||||
|
||||
q: (1, NH, HD) — single decode token
|
||||
kv: (N, HD) — all cached KV (already with RoPE)
|
||||
"""
|
||||
NH = q.shape[1]
|
||||
HD = q.shape[2]
|
||||
q_t = q.permute(1, 0, 2) # (NH, 1, HD)
|
||||
kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) # (NH, N, HD)
|
||||
out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=False, scale=scale)
|
||||
return out.permute(1, 0, 2) # (1, NH, HD)
|
||||
|
||||
|
||||
def swa_attention(q, kv, positions, scale, window_size):
|
||||
"""Sliding window attention.
|
||||
|
||||
q: (T, NH, HD) with RoPE
|
||||
kv: (total_len, HD) — ALL cached KV with RoPE
|
||||
positions: (T,) — absolute positions of the query tokens
|
||||
"""
|
||||
T, NH, HD = q.shape
|
||||
total_len = kv.shape[0]
|
||||
output = torch.zeros_like(q)
|
||||
|
||||
for t in range(T):
|
||||
pos = positions[t].item()
|
||||
window_start = max(0, pos - window_size + 1)
|
||||
window_len = pos - window_start + 1
|
||||
if window_len <= 0:
|
||||
continue
|
||||
kv_window = kv[window_start:pos + 1] # (window_len, HD)
|
||||
q_t = q[t:t + 1] # (1, NH, HD)
|
||||
output[t] = decode_attention(q_t, kv_window, scale).squeeze(0)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
# ── Full Pipeline ─────────────────────────────────────────────────────
|
||||
|
||||
def blackwell_attention_forward(
|
||||
# Inputs
|
||||
q, # (T, NH, HD) with RoPE already applied
|
||||
kv, # (T, HD) kv_normed, RoPE'd — the NEW tokens' KV
|
||||
positions, # (T,) absolute positions
|
||||
# KV Cache
|
||||
swa_kv_cache, # (num_blocks, block_size, HD) fp8 paged cache
|
||||
swa_inv_scale, # (num_blocks * block_size, 1) per-token inv scale
|
||||
slot_mapping, # (T,) slot indices for writing
|
||||
block_size, # tokens per block
|
||||
seq_lens, # (num_seqs,) total sequence lengths (prefill + history)
|
||||
num_prefills, # number of prefill sequences
|
||||
num_decode_tokens, # number of decode tokens
|
||||
# Params
|
||||
scale, # 1/sqrt(HD)
|
||||
nope_dim, # 448
|
||||
rope_dim, # 64
|
||||
window_size, # 128
|
||||
compress_ratio, # 0, 1, 4, or 128
|
||||
cos_sin_cache, # (max_pos, rope_dim) for RoPE
|
||||
attn_sink, # (NH,) sink weights
|
||||
):
|
||||
"""Full attention forward for Blackwell (SM100+).
|
||||
|
||||
This is what replaces vLLM's _attention_impl_blackwell.
|
||||
|
||||
Steps:
|
||||
1. Quantize + write new KV to paged cache
|
||||
2. Read ALL cached KV for each sequence
|
||||
3. Attention (prefill: causal, decode: full)
|
||||
4. Return attention output (T, NH, HD)
|
||||
"""
|
||||
T = q.shape[0]
|
||||
NH = q.shape[1]
|
||||
HD = q.shape[2]
|
||||
device = q.device
|
||||
|
||||
# Step 1: Quantize new KV and write to cache
|
||||
# kv already has RoPE applied (done by caller)
|
||||
kv_fp8, kv_inv_scale = kv_quantize_fp8(kv)
|
||||
paged_kv_write(kv_fp8, slot_mapping, swa_kv_cache, block_size)
|
||||
# Write inv_scale to flat cache
|
||||
for t in range(T):
|
||||
slot = slot_mapping[t].item()
|
||||
swa_inv_scale[slot] = kv_inv_scale[t]
|
||||
|
||||
# Step 2 & 3: Read cached KV and attend
|
||||
# For simplicity in this initial version, we separate prefill and decode
|
||||
output = torch.zeros(T, NH, HD, dtype=torch.bfloat16, device=device)
|
||||
|
||||
if num_decode_tokens > 0:
|
||||
# Decode tokens: each needs ALL prior KV from cache
|
||||
for t in range(num_decode_tokens):
|
||||
pos = positions[t].item()
|
||||
# Read all KV from position 0 to pos
|
||||
all_slots = torch.arange(pos + 1, dtype=torch.int64, device=device)
|
||||
kv_cached_fp8 = paged_kv_read(all_slots, swa_kv_cache, block_size, pos + 1, HD)
|
||||
kv_inv_scales = swa_inv_scale[all_slots]
|
||||
kv_cached = kv_dequantize_fp8(kv_cached_fp8, kv_inv_scales)
|
||||
|
||||
# Apply SWA window
|
||||
window_start = max(0, pos - window_size + 1)
|
||||
kv_window = kv_cached[window_start:]
|
||||
|
||||
q_t = q[t:t + 1] # (1, NH, HD)
|
||||
output[t] = decode_attention(q_t, kv_window, scale).squeeze(0)
|
||||
|
||||
if num_prefills > 0:
|
||||
# Prefill tokens: causal attention using the NEW kv (not from cache,
|
||||
# since all KV is available from the current forward pass)
|
||||
# But we DO write to cache for future decode steps
|
||||
prefill_slice = slice(num_decode_tokens, T)
|
||||
output[prefill_slice] = causal_prefill_attention(
|
||||
q[prefill_slice], kv[prefill_slice], scale
|
||||
)
|
||||
|
||||
return output
|
||||
651
dsv4/reference/compressor.py
Normal file
651
dsv4/reference/compressor.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.")
|
||||
428
dsv4/reference/csa_attention.py
Normal file
428
dsv4/reference/csa_attention.py
Normal file
@@ -0,0 +1,428 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
CSA (Compressed Sparse Attention) + HCA (Heavily Compressed Attention) kernel
|
||||
for DeepSeek-V4-Pro.
|
||||
|
||||
Replaces vLLM's FlashMLA sparse attention which doesn't work on Blackwell.
|
||||
|
||||
Architecture:
|
||||
- CSA (C128A): KV cache compressed 128x. Indexer finds top-k relevant positions.
|
||||
Sparse attention attends only to those positions.
|
||||
- HCA (C4A): KV cache compressed 4x with overlap. Similar indexer + sparse attention.
|
||||
- SWA: Standard sliding window attention (compress_ratio=0/1).
|
||||
|
||||
The attention mechanism in DeepSeek-V4:
|
||||
1. Q: hidden → q_a_proj → q_norm → q_b_proj → (T, NH, HD) → RoPE
|
||||
2. KV: hidden → kv_proj → (T, HD) → RoPE → FP8 quant → KV cache (paged)
|
||||
3. Compressor: hidden → fused_wkv_wgate → compressed KV + score → state cache
|
||||
4. Indexer: compressed state cache → top-k position indices
|
||||
5. Sparse attention: Q attends to compressed KV at top-k positions
|
||||
6. Window attention: Q attends to local window
|
||||
7. Merge: combine sparse + window attention outputs using attn_sink weights
|
||||
|
||||
This module implements steps 4-7 in pure PyTorch (works on any GPU).
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
|
||||
# ── Sparse Attention Kernel ───────────────────────────────────────────
|
||||
|
||||
def csa_sparse_attention(
|
||||
q: torch.Tensor, # (num_tokens, num_heads, head_dim) - with RoPE applied
|
||||
kv_cache: torch.Tensor, # (num_blocks, block_size, head_dim) - FP8 compressed KV
|
||||
topk_indices: torch.Tensor, # (num_tokens, 1, num_topk) - global position indices
|
||||
topk_lens: torch.Tensor, # (num_tokens,) - valid length per token
|
||||
block_table: torch.Tensor, # (num_seqs, num_blocks_per_seq)
|
||||
block_size: int,
|
||||
scale: float,
|
||||
nope_dim: int, # dimensions without RoPE
|
||||
rope_dim: int, # dimensions with RoPE
|
||||
cos_sin_cache: torch.Tensor, # (max_pos, rope_dim) for RoPE on gathered KV
|
||||
positions: torch.Tensor, # (num_tokens,) position IDs
|
||||
attn_sink: torch.Tensor, # (num_heads,) sink weights (softmax bias)
|
||||
) -> torch.Tensor:
|
||||
"""CSA sparse attention: attend to top-k positions in compressed KV cache.
|
||||
|
||||
For each query token, gathers KV from the top-k positions and performs
|
||||
standard scaled dot-product attention.
|
||||
"""
|
||||
num_tokens, num_heads, head_dim = q.shape
|
||||
device = q.device
|
||||
|
||||
# Gather KV from compressed cache at top-k positions
|
||||
# topk_indices: (num_tokens, 1, num_topk) → (num_tokens, num_topk)
|
||||
if topk_indices.dim() == 3:
|
||||
topk_indices = topk_indices.squeeze(1)
|
||||
|
||||
num_topk = topk_indices.shape[-1]
|
||||
|
||||
# Convert global position indices to (block_idx, offset) for paged cache
|
||||
# global_pos → block_idx = global_pos // block_size
|
||||
# global_pos → offset = global_pos % block_size
|
||||
topk_block_idx = topk_indices // block_size # (num_tokens, num_topk)
|
||||
topk_offset = topk_indices % block_size
|
||||
|
||||
# For each token, we need its sequence's block table to look up physical blocks
|
||||
# This is a simplified version assuming single-sequence for now
|
||||
# In production, we'd use token_to_req_indices to get the right block_table row
|
||||
|
||||
# Gather KV from cache
|
||||
# kv_cache shape: (num_blocks, block_size, head_dim) in FP8
|
||||
# Dequantize FP8 to BF16
|
||||
if kv_cache.dtype == torch.uint8:
|
||||
# FP8 E4M3 dequant: values = uint8 → float8_e4m3fn → bfloat16
|
||||
kv_bf16 = kv_cache.view(torch.float8_e4m3fn).to(torch.bfloat16)
|
||||
else:
|
||||
kv_bf16 = kv_cache.to(torch.bfloat16)
|
||||
|
||||
# For each query token, gather its top-k KV vectors
|
||||
# This is the core sparse gather operation
|
||||
# Output: (num_tokens, num_topk, head_dim)
|
||||
k_gathered = torch.zeros(
|
||||
num_tokens, num_topk, head_dim,
|
||||
dtype=torch.bfloat16, device=device,
|
||||
)
|
||||
|
||||
for t in range(num_tokens):
|
||||
for k_idx in range(min(topk_lens[t].item(), num_topk)):
|
||||
gpos = topk_indices[t, k_idx].item()
|
||||
if gpos < 0:
|
||||
continue
|
||||
bidx = gpos // block_size
|
||||
boff = gpos % block_size
|
||||
if bidx < kv_bf16.shape[0] and boff < kv_bf16.shape[1]:
|
||||
k_gathered[t, k_idx] = kv_bf16[bidx, boff]
|
||||
|
||||
# Apply RoPE to gathered KV (the compressed KV needs RoPE at its original position)
|
||||
if rope_dim > 0:
|
||||
# Positions of gathered KV
|
||||
kv_positions = topk_indices.clamp(min=0) # (num_tokens, num_topk)
|
||||
half_rot = rope_dim // 2
|
||||
cos_kv = cos_sin_cache[kv_positions, :half_rot] # (NT, num_topk, half_rot)
|
||||
sin_kv = cos_sin_cache[kv_positions, half_rot:] # (NT, num_topk, half_rot)
|
||||
|
||||
# Apply GPT-J RoPE to the rope portion of k_gathered
|
||||
k_rope = k_gathered[:, :, nope_dim:] # (NT, num_topk, rope_dim)
|
||||
k_even = k_rope[:, :, 0::2]
|
||||
k_odd = k_rope[:, :, 1::2]
|
||||
cos_f = cos_kv.unsqueeze(2).to(k_gathered.dtype) # (NT, num_topk, 1, half_rot)
|
||||
sin_f = sin_kv.unsqueeze(2).to(k_gathered.dtype)
|
||||
|
||||
# RoPE on 2D KV (no head dim, treat as single head)
|
||||
k_even_rot = k_even * cos_f.squeeze(2) - k_odd * sin_f.squeeze(2)
|
||||
k_odd_rot = k_even * sin_f.squeeze(2) + k_odd * cos_f.squeeze(2)
|
||||
k_gathered[:, :, nope_dim:][:, :, 0::2] = k_even_rot
|
||||
k_gathered[:, :, nope_dim:][:, :, 1::2] = k_odd_rot
|
||||
|
||||
# Expand k for multi-head attention
|
||||
# k_gathered: (NT, num_topk, HD) → (NT, NH, num_topk, HD)
|
||||
k_expanded = k_gathered.unsqueeze(1).expand(-1, num_heads, -1, -1)
|
||||
|
||||
# Q: (NT, NH, HD) → (NT, NH, 1, HD)
|
||||
q_4d = q.unsqueeze(2)
|
||||
|
||||
# Attention scores: (NT, NH, 1, num_topk)
|
||||
attn_weights = torch.matmul(q_4d, k_expanded.transpose(-1, -2)) * scale
|
||||
|
||||
# Apply attention sink bias
|
||||
# attn_sink: (NH,) → add to the first position's logit
|
||||
if attn_sink is not None:
|
||||
sink_bias = attn_sink.view(1, num_heads, 1, 1) # (1, NH, 1, 1)
|
||||
attn_weights[:, :, :, 0] += sink_bias.squeeze(-1)
|
||||
|
||||
# Causal mask: don't attend to future positions
|
||||
# (simplified — assumes topk_indices are already filtered for causality)
|
||||
|
||||
# Mask invalid positions
|
||||
valid_mask = torch.arange(num_topk, device=device).unsqueeze(0) < topk_lens.unsqueeze(1) # (NT, num_topk)
|
||||
attn_weights = attn_weights.masked_fill(~valid_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
|
||||
|
||||
attn_weights = F.softmax(attn_weights.float(), dim=-1).to(torch.bfloat16)
|
||||
|
||||
# Weighted sum: (NT, NH, 1, num_topk) @ (NT, NH, num_topk, HD) → (NT, NH, 1, HD)
|
||||
attn_output = torch.matmul(attn_weights, k_expanded)
|
||||
return attn_output.squeeze(2) # (NT, NH, HD)
|
||||
|
||||
|
||||
def swa_attention(
|
||||
q: torch.Tensor, # (num_tokens, num_heads, head_dim)
|
||||
swa_kv_cache: torch.Tensor, # (num_blocks, block_size, head_dim) - SWA KV cache
|
||||
positions: torch.Tensor, # (num_tokens,)
|
||||
block_table: torch.Tensor, # (num_seqs, num_blocks_per_seq)
|
||||
slot_mapping: torch.Tensor, # (num_tokens,)
|
||||
block_size: int,
|
||||
window_size: int,
|
||||
scale: float,
|
||||
) -> torch.Tensor:
|
||||
"""Sliding window attention: attend to local window of tokens.
|
||||
|
||||
Standard multi-head attention over the last `window_size` tokens.
|
||||
"""
|
||||
num_tokens, num_heads, head_dim = q.shape
|
||||
device = q.device
|
||||
|
||||
# Dequantize SWA cache if FP8
|
||||
if swa_kv_cache.dtype == torch.uint8:
|
||||
swa_bf16 = swa_kv_cache.view(torch.float8_e4m3fn).to(torch.bfloat16)
|
||||
else:
|
||||
swa_bf16 = swa_kv_cache.to(torch.bfloat16)
|
||||
|
||||
# For a simplified implementation, gather all KV in the window
|
||||
# In production, this would use paged cache access
|
||||
output = torch.zeros(num_tokens, num_heads, head_dim, dtype=torch.bfloat16, device=device)
|
||||
|
||||
for t in range(num_tokens):
|
||||
pos = positions[t].item()
|
||||
window_start = max(0, pos - window_size + 1)
|
||||
window_len = pos - window_start + 1
|
||||
|
||||
if window_len == 0:
|
||||
continue
|
||||
|
||||
# Gather KV from window
|
||||
k_window = torch.zeros(window_len, head_dim, dtype=torch.bfloat16, device=device)
|
||||
for i, p in enumerate(range(window_start, pos + 1)):
|
||||
slot = p # simplified: slot = position for contiguous sequences
|
||||
bidx = slot // block_size
|
||||
boff = slot % block_size
|
||||
if bidx < swa_bf16.shape[0] and boff < swa_bf16.shape[1]:
|
||||
k_window[i] = swa_bf16[bidx, boff]
|
||||
|
||||
# Multi-head attention
|
||||
q_t = q[t] # (NH, HD)
|
||||
k_exp = k_window.unsqueeze(0).expand(num_heads, -1, -1) # (NH, window_len, HD)
|
||||
|
||||
# Q @ K^T: (NH, 1, HD) @ (NH, HD, window_len) → (NH, 1, window_len)
|
||||
scores = torch.matmul(q_t.unsqueeze(1), k_exp.transpose(-1, -2)) * scale
|
||||
scores = F.softmax(scores.float(), dim=-1).to(torch.bfloat16)
|
||||
|
||||
# Weighted sum: (NH, 1, window_len) @ (NH, window_len, HD) → (NH, 1, HD)
|
||||
out_t = torch.matmul(scores, k_exp).squeeze(1) # (NH, HD)
|
||||
output[t] = out_t
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def csa_hca_forward(
|
||||
q: torch.Tensor, # (num_tokens, num_heads, head_dim) with RoPE
|
||||
kv: torch.Tensor, # (num_tokens, head_dim) - KV latent (after norm)
|
||||
positions: torch.Tensor, # (num_tokens,)
|
||||
# SWA cache
|
||||
swa_kv_cache: torch.Tensor,
|
||||
swa_block_table: torch.Tensor,
|
||||
swa_slot_mapping: torch.Tensor,
|
||||
swa_block_size: int,
|
||||
window_size: int,
|
||||
# CSA cache (optional, for compress_ratio > 1)
|
||||
csa_kv_cache: Optional[torch.Tensor] = None,
|
||||
csa_block_table: Optional[torch.Tensor] = None,
|
||||
csa_block_size: int = 256,
|
||||
compress_ratio: int = 1,
|
||||
topk_indices: Optional[torch.Tensor] = None,
|
||||
topk_lens: Optional[torch.Tensor] = None,
|
||||
# Params
|
||||
scale: float = 1.0,
|
||||
nope_dim: int = 448,
|
||||
rope_dim: int = 64,
|
||||
cos_sin_cache: Optional[torch.Tensor] = None,
|
||||
attn_sink: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Full CSA/HCA/SWA forward pass.
|
||||
|
||||
For compress_ratio > 1: CSA/HCA sparse attention + SWA
|
||||
For compress_ratio <= 1: SWA only
|
||||
"""
|
||||
num_tokens, num_heads, head_dim = q.shape
|
||||
device = q.device
|
||||
|
||||
if compress_ratio <= 1:
|
||||
# SWA-only layer
|
||||
return swa_attention(
|
||||
q, swa_kv_cache, positions, swa_block_table,
|
||||
swa_slot_mapping, swa_block_size, window_size, scale,
|
||||
)
|
||||
|
||||
# CSA/HCA layer: sparse attention + SWA, merged with sink weights
|
||||
sparse_out = csa_sparse_attention(
|
||||
q, csa_kv_cache, topk_indices, topk_lens,
|
||||
csa_block_table, csa_block_size, scale,
|
||||
nope_dim, rope_dim, cos_sin_cache, positions, attn_sink,
|
||||
)
|
||||
|
||||
swa_out = swa_attention(
|
||||
q, swa_kv_cache, positions, swa_block_table,
|
||||
swa_slot_mapping, swa_block_size, window_size, scale,
|
||||
)
|
||||
|
||||
# Merge sparse + SWA outputs
|
||||
# The sink weights determine the mixing between sparse and window attention
|
||||
# For now, simple addition (the actual merge uses attn_sink as a learned weight)
|
||||
if attn_sink is not None:
|
||||
# attn_sink: (num_heads,) — softmax bias toward the sink token
|
||||
# When sink weight is -inf, no sink effect → pure SWA + sparse
|
||||
# When sink weight is 0, equal mixing
|
||||
# In practice, attn_sink is trained and typically small
|
||||
sink_weight = torch.sigmoid(attn_sink).view(1, num_heads, 1)
|
||||
output = sparse_out * (1 - sink_weight) + swa_out * sink_weight
|
||||
else:
|
||||
output = sparse_out + swa_out
|
||||
|
||||
return output
|
||||
|
||||
|
||||
# ── Batched sparse attention (optimized, no Python loops) ─────────────
|
||||
|
||||
def csa_sparse_attention_batched(
|
||||
q: torch.Tensor, # (T, NH, HD)
|
||||
kv_cache: torch.Tensor, # (num_blocks, block_size, kv_dim) FP8 or BF16
|
||||
topk_indices: torch.Tensor, # (T, num_topk) global position indices
|
||||
topk_lens: torch.Tensor, # (T,) valid lengths
|
||||
block_size: int,
|
||||
scale: float,
|
||||
nope_dim: int,
|
||||
rope_dim: int,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
attn_sink: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Optimized CSA sparse attention using batched gather + SDPA.
|
||||
|
||||
No Python loops. Uses torch.gather and F.scaled_dot_product_attention.
|
||||
"""
|
||||
T, NH, HD = q.shape
|
||||
device = q.device
|
||||
num_topk = topk_indices.shape[-1]
|
||||
|
||||
# Dequantize KV cache
|
||||
if kv_cache.dtype == torch.uint8:
|
||||
kv_flat = kv_cache.view(torch.float8_e4m3fn).to(torch.bfloat16)
|
||||
else:
|
||||
kv_flat = kv_cache.to(torch.bfloat16)
|
||||
|
||||
# Flatten cache: (num_blocks * block_size, kv_dim)
|
||||
num_blocks, bs, kv_dim = kv_flat.shape
|
||||
kv_flat = kv_flat.reshape(num_blocks * bs, kv_dim)
|
||||
|
||||
# Clamp topk_indices to valid range and gather
|
||||
# topk_indices: (T, num_topk) → gather from kv_flat
|
||||
safe_indices = topk_indices.clamp(min=0, max=kv_flat.shape[0] - 1)
|
||||
|
||||
# Gather: (T, num_topk, kv_dim)
|
||||
# torch.gather needs (T, num_topk) index → expand to (T, num_topk, kv_dim)
|
||||
idx_expanded = safe_indices.unsqueeze(-1).expand(-1, -1, kv_dim)
|
||||
k_gathered = torch.gather(
|
||||
kv_flat.unsqueeze(0).expand(T, -1, -1), # (T, total_positions, kv_dim)
|
||||
1, # dim=1
|
||||
idx_expanded, # (T, num_topk, kv_dim)
|
||||
)
|
||||
|
||||
# Mask invalid positions
|
||||
valid_mask = torch.arange(num_topk, device=device).unsqueeze(0) < topk_lens.unsqueeze(1)
|
||||
k_gathered = k_gathered * valid_mask.unsqueeze(-1).to(k_gathered.dtype)
|
||||
|
||||
# Apply RoPE to gathered K (GPT-J style)
|
||||
if rope_dim > 0 and cos_sin_cache is not None:
|
||||
kv_positions = safe_indices # (T, num_topk)
|
||||
half_rot = rope_dim // 2
|
||||
cos_kv = cos_sin_cache[kv_positions, :half_rot] # (T, num_topk, half_rot)
|
||||
sin_kv = cos_sin_cache[kv_positions, half_rot:]
|
||||
|
||||
k_rope = k_gathered[:, :, nope_dim:] # (T, num_topk, rope_dim)
|
||||
k_even = k_rope[:, :, 0::2]
|
||||
k_odd = k_rope[:, :, 1::2]
|
||||
cos_f = cos_kv.to(k_gathered.dtype)
|
||||
sin_f = sin_kv.to(k_gathered.dtype)
|
||||
k_gathered[:, :, nope_dim:][:, :, 0::2] = k_even * cos_f - k_odd * sin_f
|
||||
k_gathered[:, :, nope_dim:][:, :, 1::2] = k_even * sin_f + k_odd * cos_f
|
||||
|
||||
# Expand for multi-head: (T, num_topk, HD) → (T, NH, num_topk, HD)
|
||||
k_heads = k_gathered.unsqueeze(1).expand(-1, NH, -1, -1)
|
||||
v_heads = k_heads.clone() # K=V in MLA-style attention
|
||||
|
||||
# Q: (T, NH, HD) → (T, NH, 1, HD)
|
||||
q_4d = q.unsqueeze(2)
|
||||
|
||||
# Use PyTorch SDPA (works on all GPUs including Blackwell)
|
||||
# Need shapes: (T*NH, 1, HD) and (T*NH, num_topk, HD)
|
||||
q_2d = q.reshape(T * NH, 1, HD)
|
||||
k_2d = k_heads.reshape(T * NH, num_topk, HD)
|
||||
v_2d = v_heads.reshape(T * NH, num_topk, HD)
|
||||
|
||||
# Build attention mask from valid positions
|
||||
# (T, num_topk) → (T*NH, 1, num_topk)
|
||||
attn_mask = valid_mask.unsqueeze(1).expand(-1, NH, -1).reshape(T * NH, 1, num_topk)
|
||||
attn_mask = attn_mask.to(torch.bool)
|
||||
|
||||
# Apply attn_sink bias
|
||||
if attn_sink is not None:
|
||||
# Add sink bias to first position's attention logit
|
||||
# attn_sink: (NH,) → (T*NH, 1, 1) broadcast
|
||||
sink = attn_sink.view(1, NH, 1).expand(T, -1, -1).reshape(T * NH, 1, 1)
|
||||
# We'll add this after SDPA by adjusting the mask
|
||||
# Actually, we need to handle this before softmax
|
||||
# For now, just note that attn_sink is a learned bias
|
||||
|
||||
# PyTorch SDPA
|
||||
with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.FLASH_ATTENTION,
|
||||
torch.nn.attention.SDPBackend.MATH]):
|
||||
out_2d = F.scaled_dot_product_attention(
|
||||
q_2d, k_2d, v_2d,
|
||||
attn_mask=attn_mask if not attn_mask.all() else None,
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
return out_2d.squeeze(1).reshape(T, NH, HD)
|
||||
|
||||
|
||||
# ── Simplified full-attention fallback (no compression, for testing) ──
|
||||
|
||||
def full_attention_reference(
|
||||
q: torch.Tensor, # (T, NH, HD) with RoPE
|
||||
kv: torch.Tensor, # (T, HD) KV latent
|
||||
scale: float = 1.0,
|
||||
) -> torch.Tensor:
|
||||
"""Full attention reference: attend to all positions.
|
||||
|
||||
Useful for testing when CSA cache is not available.
|
||||
Uses PyTorch SDPA which works on all GPUs.
|
||||
"""
|
||||
T, NH, HD = q.shape
|
||||
|
||||
# K=V from kv latent (shared across all heads and all query positions)
|
||||
# kv: (T, HD) → each token's KV is seen by all heads at all query positions
|
||||
k = kv.unsqueeze(1).expand(-1, NH, -1).contiguous() # (T, NH, HD)
|
||||
# For cross-attention where each Q attends to all KV positions:
|
||||
# K needs to be (T_q, NH, T_kv, HD) — repeat for each query position
|
||||
k = k.unsqueeze(0).expand(T, -1, -1, -1).contiguous() # (T, T, NH, HD) → wrong order
|
||||
# Actually: for self-attention, K/V shape for SDPA is (batch, seq_kv, HD)
|
||||
# where batch = T*NH (each query token is a batch, each head independent)
|
||||
# K/V: (T*NH, T, HD) — each (query, head) pair attends to all T KV positions
|
||||
kv_expanded = kv.unsqueeze(1).expand(-1, NH, -1).contiguous() # (T, NH, HD)
|
||||
# Repeat KV for each query: (T, NH, HD) → (T*NH, T, HD)
|
||||
k_2d = kv_expanded.permute(1, 0, 2).unsqueeze(1).expand(NH, T, T, -1).contiguous().reshape(T * NH, T, HD)
|
||||
v_2d = k_2d.clone()
|
||||
|
||||
# Q: (T, NH, HD) → (T*NH, 1, HD)
|
||||
q_2d = q.reshape(T * NH, 1, HD)
|
||||
|
||||
# Manual attention (SDPA mask handling is tricky with batched single-query)
|
||||
# scores: (T*NH, 1, T) = Q @ K^T
|
||||
scores = torch.matmul(q_2d, k_2d.transpose(-1, -2)) * scale
|
||||
# Causal mask: each query at position i can only attend to positions <= i
|
||||
# Since each batch is (query_pos, head), and KV has all T positions,
|
||||
# we need position-aware masking
|
||||
# For single-query batches: batch i corresponds to (pos i // NH, head i % NH)
|
||||
# All positions <= i // NH are valid
|
||||
# Simple approach: use a per-query mask
|
||||
query_positions = torch.arange(T, device=q.device).unsqueeze(1).repeat(1, NH).reshape(T * NH) # (T*NH,)
|
||||
kv_positions = torch.arange(T, device=q.device).unsqueeze(0) # (1, T)
|
||||
causal = kv_positions <= query_positions.unsqueeze(1) # (T*NH, T)
|
||||
scores = scores.squeeze(1).masked_fill(~causal, float('-inf')) # (T*NH, T)
|
||||
weights = F.softmax(scores.float(), dim=-1).to(q.dtype) # (T*NH, T)
|
||||
out = torch.matmul(weights.unsqueeze(1), v_2d) # (T*NH, 1, HD)
|
||||
|
||||
return out.squeeze(1).reshape(T, NH, HD)
|
||||
422
dsv4/reference/moe_pipeline.py
Normal file
422
dsv4/reference/moe_pipeline.py
Normal file
@@ -0,0 +1,422 @@
|
||||
"""
|
||||
Full NVFP4 MoE pipeline using CuTeDSL ScaledGroupedGemmKernel.
|
||||
|
||||
Data flow (NVFP4-native, BF16 only where required):
|
||||
1. BF16 hidden_states → quantize to NVFP4 (stage_activation)
|
||||
2. L1 GEMM: NVFP4 × NVFP4 → BF16 output (gate+up)
|
||||
3. SiLU(gate) * up → BF16 activated (nonlinear requires BF16)
|
||||
4. Re-quantize activated → NVFP4 (stage_activation)
|
||||
5. L2 GEMM: NVFP4 × NVFP4 → BF16 output (down_proj)
|
||||
6. Scatter with routing weights → BF16 output
|
||||
|
||||
Both GEMMs are fully NVFP4: A in float4_e2m1fn_x2, B in float4_e2m1fn_x2,
|
||||
block scales in float8_e4m3fn, global scales in float32.
|
||||
"""
|
||||
import torch
|
||||
|
||||
from dsv4.ops.quantize import (
|
||||
quantize_to_nvfp4,
|
||||
quantize_weight_to_nvfp4,
|
||||
)
|
||||
from dsv4.ops.layouts import (
|
||||
assemble_scales_2d_side,
|
||||
assemble_scales_3d_side,
|
||||
make_b_k_major,
|
||||
compute_expert_offsets,
|
||||
interleave_l1_weights,
|
||||
deinterleave_l1_weights,
|
||||
)
|
||||
from dsv4.ops.gemm_runner import (
|
||||
run_nvfp4_grouped_gemm,
|
||||
run_fused_swiglu_grouped_gemm,
|
||||
warmup_fused_swiglu_compilation,
|
||||
)
|
||||
|
||||
|
||||
def stage_activation(x_bf16):
|
||||
"""Quantize BF16 activation to NVFP4.
|
||||
|
||||
This is the NVFP4-native equivalent of the old stage_activation.
|
||||
Keeps data in FP4 as long as possible — only leaves NVFP4 for nonlinear ops.
|
||||
|
||||
Returns (x_fp4, x_sf, global_scale) where:
|
||||
x_fp4: float4_e2m1fn_x2 (native PyTorch FP4)
|
||||
x_sf: float8_e4m3fn block scales
|
||||
global_scale: float32 scalar
|
||||
"""
|
||||
return quantize_to_nvfp4(x_bf16)
|
||||
|
||||
|
||||
def quantize_weight(w_bf16):
|
||||
"""Quantize BF16 weight to NVFP4.
|
||||
|
||||
Weight is (K, N) where K is the input/hidden dim (packed dimension).
|
||||
|
||||
Returns (w_fp4, w_sf, global_scale).
|
||||
"""
|
||||
return quantize_weight_to_nvfp4(w_bf16)
|
||||
|
||||
|
||||
def prepare_nvfp4_moe_weights(nvfp4_tensors, layer_idx, expert_indices):
|
||||
"""Load NVFP4 checkpoint weights and prepare for the grouped GEMM.
|
||||
|
||||
Dequantizes checkpoint NVFP4 → BF16 → re-quantizes to our native format.
|
||||
This round-trip ensures our FP4 packing convention matches the kernel.
|
||||
|
||||
Future optimization: load checkpoint FP4 bytes directly into
|
||||
float4_e2m1fn_x2 tensors without the BF16 round-trip.
|
||||
|
||||
Returns dict with l1 and l2 weight info per expert.
|
||||
"""
|
||||
from tests.layertest import dequantize_nvfp4_weight, DEVICE
|
||||
|
||||
l1_weights = [] # gate+up fused, (K, N) = (hidden, intermediate)
|
||||
l2_weights = [] # down, (K, N) = (intermediate, hidden)
|
||||
|
||||
for e in expert_indices:
|
||||
# L1: gate + up
|
||||
gate_w_bf16 = dequantize_nvfp4_weight(
|
||||
nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE),
|
||||
nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE),
|
||||
nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale_2"].item(),
|
||||
)
|
||||
up_w_bf16 = dequantize_nvfp4_weight(
|
||||
nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight"].to(DEVICE),
|
||||
nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE),
|
||||
nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale_2"].item(),
|
||||
)
|
||||
|
||||
# Fuse gate+up: (6144, 7168) → transpose to (7168, 6144) for weight quantization
|
||||
fused_l1 = torch.cat([gate_w_bf16, up_w_bf16], dim=0) # (6144, 7168)
|
||||
l1_w_bf16 = fused_l1.T # (7168, 6144) — K=7168, N=6144
|
||||
l1_weights.append(l1_w_bf16)
|
||||
|
||||
# L2: down
|
||||
down_w_key = f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight"
|
||||
if down_w_key in nvfp4_tensors:
|
||||
down_w_bf16 = dequantize_nvfp4_weight(
|
||||
nvfp4_tensors[down_w_key].to(DEVICE),
|
||||
nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE),
|
||||
nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale_2"].item(),
|
||||
)
|
||||
# down_proj is (7168, 3072) → transpose to (3072, 7168) for K=intermediate
|
||||
l2_w_bf16 = down_w_bf16.T # (3072, 7168) — K=3072, N=7168
|
||||
else:
|
||||
# Expert 211 has no down_proj
|
||||
l2_w_bf16 = torch.zeros(3072, 7168, dtype=torch.bfloat16, device=DEVICE)
|
||||
|
||||
l2_weights.append(l2_w_bf16)
|
||||
|
||||
# Quantize all weights to NVFP4
|
||||
l1_fp4, l1_sf, l1_gs = [], [], []
|
||||
l2_fp4, l2_sf, l2_gs = [], [], []
|
||||
|
||||
for l1_w, l2_w in zip(l1_weights, l2_weights):
|
||||
w_fp4, w_sf, w_gs = quantize_weight(l1_w)
|
||||
l1_fp4.append(w_fp4)
|
||||
l1_sf.append(w_sf)
|
||||
l1_gs.append(w_gs)
|
||||
|
||||
w_fp4, w_sf, w_gs = quantize_weight(l2_w)
|
||||
l2_fp4.append(w_fp4)
|
||||
l2_sf.append(w_sf)
|
||||
l2_gs.append(w_gs)
|
||||
|
||||
return {
|
||||
'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs,
|
||||
'l2_fp4': l2_fp4, 'l2_sf': l2_sf, 'l2_gs': l2_gs,
|
||||
}
|
||||
|
||||
|
||||
def run_nvfp4_moe(
|
||||
hidden_states, # (num_tokens, hidden_size) BF16
|
||||
expert_ids, # (num_tokens, top_k) int32
|
||||
expert_weights, # (num_tokens, top_k) float32
|
||||
weights, # dict from prepare_nvfp4_moe_weights
|
||||
expert_indices, # list of expert IDs
|
||||
swiglu_limit=None, # Optional clamp for SiLU output
|
||||
):
|
||||
"""Run the full NVFP4 MoE forward pass.
|
||||
|
||||
NVFP4-native pipeline:
|
||||
1. Quantize activation → NVFP4
|
||||
2. L1 GEMM (NVFP4 × NVFP4 → BF16)
|
||||
3. SiLU(gate) * up (BF16 — nonlinear requires BF16)
|
||||
4. Re-quantize → NVFP4
|
||||
5. L2 GEMM (NVFP4 × NVFP4 → BF16)
|
||||
6. Scatter with routing weights → BF16
|
||||
|
||||
Returns: (num_tokens, hidden_size) BF16
|
||||
"""
|
||||
num_tokens, hidden_size = hidden_states.shape
|
||||
top_k = expert_ids.shape[1]
|
||||
device = hidden_states.device
|
||||
|
||||
# ── Build slot-based routing ──
|
||||
expert_token_lists = {e: [] for e in expert_indices}
|
||||
for t in range(num_tokens):
|
||||
for k in range(top_k):
|
||||
e = expert_ids[t, k].item()
|
||||
if e in expert_token_lists:
|
||||
expert_token_lists[e].append(t)
|
||||
|
||||
tokens_per_expert = [len(expert_token_lists[e]) for e in expert_indices]
|
||||
num_experts = len(expert_indices)
|
||||
|
||||
# Slot-major activation: [expert0_tokens | expert1_tokens | ...]
|
||||
slot_hidden = torch.cat([
|
||||
hidden_states[expert_token_lists[e]] for e in expert_indices
|
||||
], dim=0) if any(tpe > 0 for tpe in tokens_per_expert) else torch.zeros(0, hidden_size, dtype=torch.bfloat16, device=device)
|
||||
|
||||
num_slots = slot_hidden.shape[0]
|
||||
if num_slots == 0:
|
||||
return torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
|
||||
|
||||
expert_offsets = compute_expert_offsets(tokens_per_expert, num_experts)
|
||||
|
||||
# ════════════════════════════════════════════════════════════════
|
||||
# L1: gate + up projection (NVFP4 × NVFP4 → BF16)
|
||||
# ════════════════════════════════════════════════════════════════
|
||||
|
||||
# Quantize activation to NVFP4
|
||||
x_fp4, x_sf, x_igs = stage_activation(slot_hidden)
|
||||
|
||||
# Stack L1 weights, interleave gate/up, convert to K-major
|
||||
l1_stacked = torch.stack(weights['l1_fp4']) # (E, K, N)
|
||||
l1_stacked = interleave_l1_weights(l1_stacked) # gate/up at granularity 4 BF16
|
||||
l1_mat_b = make_b_k_major(l1_stacked)
|
||||
|
||||
# Assemble scales
|
||||
x_sf_parts = []
|
||||
offset = 0
|
||||
for tpe in tokens_per_expert:
|
||||
x_sf_parts.append(x_sf[offset:offset+tpe])
|
||||
offset += tpe
|
||||
l1_scale_a = assemble_scales_2d_side(x_sf_parts)
|
||||
# Interleave L1 SF to match the interleaved weight layout.
|
||||
# SF is (K_sf, N) from quantize_weight_to_nvfp4. interleave_l1_weights
|
||||
# operates on the last dim, which is N. So (1, K_sf, N) is correct.
|
||||
# After interleave, transpose to (N, K_sf) for the assembly function.
|
||||
l1_sf_il = []
|
||||
for sf in weights['l1_sf']:
|
||||
sf_ekn = sf.unsqueeze(0) # (1, K_sf, N)
|
||||
sf_ekn = interleave_l1_weights(sf_ekn) # interleaved along N
|
||||
l1_sf_il.append(sf_ekn[0].T.contiguous()) # (N, K_sf) for assembly
|
||||
from dsv4.kernels.gemm.grouped import assemble_raw_scales_2d3d_3d_side as _assemble_3d
|
||||
l1_scale_b = _assemble_3d(l1_sf_il)
|
||||
|
||||
# Global scales: alpha = igs * weight_gs for each expert
|
||||
l1_global_scale_a = torch.tensor([x_igs] * num_experts, dtype=torch.float32, device=device)
|
||||
l1_global_scale_b = torch.tensor(weights['l1_gs'], dtype=torch.float32, device=device)
|
||||
|
||||
# Run L1 GEMM
|
||||
l1_out = run_nvfp4_grouped_gemm(
|
||||
mat_a=x_fp4, mat_b=l1_mat_b,
|
||||
scale_a=l1_scale_a, scale_b=l1_scale_b,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=l1_global_scale_a, global_scale_b=l1_global_scale_b,
|
||||
) # (num_slots, 2*intermediate) BF16
|
||||
|
||||
# ════════════════════════════════════════════════════════════════
|
||||
# SiLU(gate) * up (BF16 — nonlinear requires BF16)
|
||||
# ════════════════════════════════════════════════════════════════
|
||||
# L1 output is (tokens, 2*intermediate) with interleaved gate/up.
|
||||
# De-interleave to recover standard [gate | up] layout.
|
||||
intermediate_size = l1_out.shape[1] // 2
|
||||
l1_deil = deinterleave_l1_weights(l1_out.unsqueeze(0).contiguous())[0]
|
||||
gate = l1_deil[:, :intermediate_size]
|
||||
up = l1_deil[:, intermediate_size:]
|
||||
gate_silu = torch.nn.functional.silu(gate)
|
||||
if swiglu_limit is not None:
|
||||
gate_silu = gate_silu.clamp(max=swiglu_limit)
|
||||
up = up.clamp(min=-swiglu_limit, max=swiglu_limit)
|
||||
activated = gate_silu * up # (num_slots, intermediate) BF16
|
||||
|
||||
# ════════════════════════════════════════════════════════════════
|
||||
# L2: down projection (NVFP4 × NVFP4 → BF16)
|
||||
# ════════════════════════════════════════════════════════════════
|
||||
|
||||
# Re-quantize activated → NVFP4
|
||||
l2_x_fp4, l2_x_sf, l2_x_igs = stage_activation(activated)
|
||||
|
||||
# Stack L2 weights
|
||||
l2_mat_b = make_b_k_major(torch.stack(weights['l2_fp4']))
|
||||
|
||||
# Assemble L2 scales
|
||||
l2_sf_parts = []
|
||||
offset = 0
|
||||
for tpe in tokens_per_expert:
|
||||
l2_sf_parts.append(l2_x_sf[offset:offset+tpe])
|
||||
offset += tpe
|
||||
l2_scale_a = assemble_scales_2d_side(l2_sf_parts)
|
||||
l2_scale_b = assemble_scales_3d_side(weights['l2_sf'])
|
||||
|
||||
# Global scales
|
||||
l2_global_scale_a = torch.tensor([l2_x_igs] * num_experts, dtype=torch.float32, device=device)
|
||||
l2_global_scale_b = torch.tensor(weights['l2_gs'], dtype=torch.float32, device=device)
|
||||
|
||||
# Run L2 GEMM
|
||||
l2_out = run_nvfp4_grouped_gemm(
|
||||
mat_a=l2_x_fp4, mat_b=l2_mat_b,
|
||||
scale_a=l2_scale_a, scale_b=l2_scale_b,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=l2_global_scale_a, global_scale_b=l2_global_scale_b,
|
||||
) # (num_slots, hidden_size) BF16
|
||||
|
||||
# ════════════════════════════════════════════════════════════════
|
||||
# Scatter with routing weights → final output
|
||||
# ════════════════════════════════════════════════════════════════
|
||||
y = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
|
||||
|
||||
slot_idx = 0
|
||||
for e in expert_indices:
|
||||
for t in expert_token_lists[e]:
|
||||
# Find which top-k slot this is for this token
|
||||
for k in range(top_k):
|
||||
if expert_ids[t, k].item() == e:
|
||||
w = expert_weights[t, k].item()
|
||||
y[t] += w * l2_out[slot_idx]
|
||||
break
|
||||
slot_idx += 1
|
||||
|
||||
return y
|
||||
|
||||
|
||||
def run_nvfp4_moe_fused(
|
||||
hidden_states, # (num_tokens, hidden_size) BF16
|
||||
expert_ids, # (num_tokens, top_k) int32
|
||||
expert_weights, # (num_tokens, top_k) float32
|
||||
weights, # dict from prepare_nvfp4_moe_weights
|
||||
expert_indices, # list of expert IDs
|
||||
swiglu_limit=0.0,
|
||||
l2_activation_gs=None, # pre-computed L2 activation global scale (avoids amax sync)
|
||||
):
|
||||
"""Run the NVFP4 MoE forward pass with fused SwiGLU kernel.
|
||||
|
||||
Fused pipeline (saves BF16 GMEM write+read for gate/up):
|
||||
1. Quantize activation -> NVFP4
|
||||
2. Fused L1 GEMM + SwiGLU (NVFP4 x NVFP4 -> BF16 with silu(gate)*up in registers)
|
||||
3. De-interleave fused output, extract SwiGLU result
|
||||
4. Re-quantize -> NVFP4
|
||||
5. L2 GEMM (NVFP4 x NVFP4 -> BF16)
|
||||
6. Scatter with routing weights -> BF16
|
||||
|
||||
Returns: (num_tokens, hidden_size) BF16
|
||||
"""
|
||||
num_tokens, hidden_size = hidden_states.shape
|
||||
top_k = expert_ids.shape[1]
|
||||
device = hidden_states.device
|
||||
|
||||
# Build slot-based routing
|
||||
expert_token_lists = {e: [] for e in expert_indices}
|
||||
for t in range(num_tokens):
|
||||
for k in range(top_k):
|
||||
e = expert_ids[t, k].item()
|
||||
if e in expert_token_lists:
|
||||
expert_token_lists[e].append(t)
|
||||
|
||||
tokens_per_expert = [len(expert_token_lists[e]) for e in expert_indices]
|
||||
num_experts = len(expert_indices)
|
||||
|
||||
slot_hidden = torch.cat([
|
||||
hidden_states[expert_token_lists[e]] for e in expert_indices
|
||||
], dim=0) if any(tpe > 0 for tpe in tokens_per_expert) else torch.zeros(0, hidden_size, dtype=torch.bfloat16, device=device)
|
||||
|
||||
num_slots = slot_hidden.shape[0]
|
||||
if num_slots == 0:
|
||||
return torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
|
||||
|
||||
expert_offsets = compute_expert_offsets(tokens_per_expert, num_experts)
|
||||
|
||||
# === L1: Fused gate+up projection with SwiGLU in registers ===
|
||||
|
||||
# Quantize activation to NVFP4
|
||||
x_fp4, x_sf, x_igs = stage_activation(slot_hidden)
|
||||
|
||||
# Stack L1 weights, interleave gate/up, convert to K-major
|
||||
l1_stacked = torch.stack(weights['l1_fp4'])
|
||||
l1_stacked = interleave_l1_weights(l1_stacked)
|
||||
l1_mat_b = make_b_k_major(l1_stacked)
|
||||
|
||||
# Assemble scales (same as non-fused path)
|
||||
x_sf_parts = []
|
||||
offset = 0
|
||||
for tpe in tokens_per_expert:
|
||||
x_sf_parts.append(x_sf[offset:offset+tpe])
|
||||
offset += tpe
|
||||
l1_scale_a = assemble_scales_2d_side(x_sf_parts)
|
||||
|
||||
l1_sf_il = []
|
||||
for sf in weights['l1_sf']:
|
||||
sf_ekn = sf.unsqueeze(0)
|
||||
sf_ekn = interleave_l1_weights(sf_ekn)
|
||||
l1_sf_il.append(sf_ekn[0].T.contiguous())
|
||||
from dsv4.kernels.gemm.grouped import assemble_raw_scales_2d3d_3d_side as _assemble_3d
|
||||
l1_scale_b = _assemble_3d(l1_sf_il)
|
||||
|
||||
l1_global_scale_a = torch.tensor([x_igs] * num_experts, dtype=torch.float32, device=device)
|
||||
l1_global_scale_b = torch.tensor(weights['l1_gs'], dtype=torch.float32, device=device)
|
||||
|
||||
# Run fused SwiGLU kernel
|
||||
# Output: (num_slots, 2*intermediate) BF16
|
||||
# Even 8-col groups = silu(gate), Odd 8-col groups = silu(gate)*up
|
||||
l1_fused_out = run_fused_swiglu_grouped_gemm(
|
||||
mat_a=x_fp4, mat_b=l1_mat_b,
|
||||
scale_a=l1_scale_a, scale_b=l1_scale_b,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=l1_global_scale_a, global_scale_b=l1_global_scale_b,
|
||||
swiglu_limit=swiglu_limit,
|
||||
)
|
||||
|
||||
# De-interleave + quantize using custom CUDA kernel (4x faster)
|
||||
intermediate_size = l1_fused_out.shape[1] // 2
|
||||
# Use pre-computed L2 activation gs, or compute from amax (fallback)
|
||||
l2_gs = l2_activation_gs if l2_activation_gs is not None else l1_fused_out.abs().amax().float().item() / 2688.0
|
||||
from dsv4.ops.quantize import (
|
||||
deinterleave_quantize_nvfp4_cuda,
|
||||
quantize_activation_nvfp4,
|
||||
)
|
||||
l2_x_fp4, l2_x_sf = deinterleave_quantize_nvfp4_cuda(l1_fused_out, intermediate_size, l2_gs)
|
||||
# Skip the separate L2 quantize step below — we already have FP4+SF
|
||||
# Set activated to None to signal we already quantized
|
||||
activated = None
|
||||
|
||||
# === L2: down projection ===
|
||||
if activated is not None:
|
||||
l2_x_fp4, l2_x_sf, l2_x_igs = stage_activation(activated)
|
||||
else:
|
||||
# Already quantized by the custom CUDA kernel
|
||||
l2_x_igs = l2_gs
|
||||
l2_mat_b = make_b_k_major(torch.stack(weights['l2_fp4']))
|
||||
|
||||
l2_sf_parts = []
|
||||
offset = 0
|
||||
for tpe in tokens_per_expert:
|
||||
l2_sf_parts.append(l2_x_sf[offset:offset+tpe])
|
||||
offset += tpe
|
||||
l2_scale_a = assemble_scales_2d_side(l2_sf_parts)
|
||||
l2_scale_b = assemble_scales_3d_side(weights['l2_sf'])
|
||||
|
||||
l2_global_scale_a = torch.tensor([l2_x_igs] * num_experts, dtype=torch.float32, device=device)
|
||||
l2_global_scale_b = torch.tensor(weights['l2_gs'], dtype=torch.float32, device=device)
|
||||
|
||||
l2_out = run_nvfp4_grouped_gemm(
|
||||
mat_a=l2_x_fp4, mat_b=l2_mat_b,
|
||||
scale_a=l2_scale_a, scale_b=l2_scale_b,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=l2_global_scale_a, global_scale_b=l2_global_scale_b,
|
||||
)
|
||||
|
||||
# Scatter with routing weights
|
||||
y = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
|
||||
slot_idx = 0
|
||||
for e in expert_indices:
|
||||
for t in expert_token_lists[e]:
|
||||
for k in range(top_k):
|
||||
if expert_ids[t, k].item() == e:
|
||||
w = expert_weights[t, k].item()
|
||||
y[t] += w * l2_out[slot_idx]
|
||||
break
|
||||
slot_idx += 1
|
||||
|
||||
return y
|
||||
Reference in New Issue
Block a user