Fix KV cache crash: skip SWA cache write on Blackwell
The SWA KV cache uses fp8_ds_mla packed layout (37376 bytes per slot, not 512). Our naive FP8 quant + write had a shape mismatch. Fix: skip the SWA cache write entirely. The compressor (Triton) handles the compressed cache. For full SDPA attention, we use the raw kv tensor directly — we don't need the paged cache at all during prefill.
This commit is contained in:
@@ -1,31 +1,15 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-License: Apache-2.0
|
||||
"""
|
||||
CSA (Compressed Sparse Attention) + HCA (Heavily Compressed Attention)
|
||||
replacement for vLLM's FlashMLA sparse attention on Blackwell (SM100+).
|
||||
CSA/HCA attention for Blackwell (SM100+).
|
||||
|
||||
FlashMLA's compiled CUDA kernels don't work on SM100. This module provides
|
||||
a pure PyTorch implementation using torch.nn.functional.scaled_dot_product_attention
|
||||
which works on all GPUs including Blackwell.
|
||||
|
||||
The architecture:
|
||||
- CSA (C128A, compress_ratio=128): KV cache compressed 128x.
|
||||
Indexer finds top-k relevant positions. Sparse attention on those.
|
||||
- HCA (C4A, compress_ratio=4): KV cache compressed 4x with overlap.
|
||||
Similar indexer + sparse attention.
|
||||
- SWA: Sliding window attention (compress_ratio=0/1).
|
||||
Replaces vLLM's FlashMLA + fused CUDA kernels with pure PyTorch.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def apply_gptj_rope(
|
||||
x: torch.Tensor, # (..., head_dim)
|
||||
cos: torch.Tensor, # (..., rope_dim // 2)
|
||||
sin: torch.Tensor, # (..., rope_dim // 2)
|
||||
nope_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""Apply GPT-J style RoPE to the last dimensions of x."""
|
||||
def apply_gptj_rope(x, cos, sin, nope_dim):
|
||||
out = x.clone()
|
||||
even = x[..., nope_dim:][..., 0::2]
|
||||
odd = x[..., nope_dim:][..., 1::2]
|
||||
@@ -34,13 +18,7 @@ def apply_gptj_rope(
|
||||
return out
|
||||
|
||||
|
||||
def apply_inv_gptj_rope(
|
||||
x: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
nope_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""Apply inverse GPT-J RoPE (sin -> -sin)."""
|
||||
def apply_inv_gptj_rope(x, cos, sin, nope_dim):
|
||||
out = x.clone()
|
||||
even = x[..., nope_dim:][..., 0::2]
|
||||
odd = x[..., nope_dim:][..., 1::2]
|
||||
@@ -50,141 +28,56 @@ def apply_inv_gptj_rope(
|
||||
|
||||
|
||||
def fused_qnorm_rope_kv_insert_py(
|
||||
q: torch.Tensor, # (T, num_heads, head_dim) — modified in-place
|
||||
kv: torch.Tensor, # (T, head_dim) — not modified
|
||||
swa_kv_cache_2d: torch.Tensor, # (num_blocks * block_size, head_dim) — written to
|
||||
slot_mapping: torch.Tensor, # (T,) — maps token to slot in cache
|
||||
positions: torch.Tensor, # (T,)
|
||||
cos_sin_cache: torch.Tensor, # (max_pos, rope_dim)
|
||||
eps: float,
|
||||
block_size: int,
|
||||
nope_dim: int,
|
||||
rope_dim: int,
|
||||
quant_to_fp8: bool = True,
|
||||
q, # (T, num_heads, head_dim) — modified in-place
|
||||
kv, # (T, head_dim) — not modified
|
||||
swa_kv_cache_2d, # paged KV cache — NOT written (cache uses fp8_ds_mla packed format)
|
||||
slot_mapping,
|
||||
positions,
|
||||
cos_sin_cache,
|
||||
eps,
|
||||
block_size,
|
||||
nope_dim,
|
||||
rope_dim,
|
||||
) -> None:
|
||||
"""Pure PyTorch replacement for fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.
|
||||
"""Pure PyTorch: per-head RMS norm on Q + GPT-J RoPE on Q.
|
||||
|
||||
Does: per-head RMS norm on Q + GPT-J RoPE on Q + GPT-J RoPE on KV + FP8 quant + KV cache insert.
|
||||
The C++ kernel also writes KV (RoPE'd + FP8 quantized) to the SWA paged
|
||||
cache, but the cache uses fp8_ds_mla packed layout which is not a simple
|
||||
[slot, head_dim] uint8 array. We skip the cache write here — the
|
||||
compressor (Triton) handles the compressed cache, and for full SDPA
|
||||
attention we use the raw kv tensor directly.
|
||||
"""
|
||||
T = q.shape[0]
|
||||
if T == 0:
|
||||
return
|
||||
|
||||
# Per-head RMS norm on Q (no learned weight — just normalize each head)
|
||||
# Per-head RMS norm on Q (no learned weight)
|
||||
q_f32 = q.float()
|
||||
q_rms = q_f32.pow(2).mean(-1, keepdim=True)
|
||||
q.copy_(torch.rsqrt(q_rms + eps) * q_f32) # in-place, keeps BF16
|
||||
q.copy_(torch.rsqrt(q_rms + eps) * q_f32)
|
||||
|
||||
# GPT-J RoPE on Q
|
||||
half = rope_dim // 2
|
||||
cos_q = cos_sin_cache[positions, :half].unsqueeze(1).to(q.dtype) # (T, 1, half)
|
||||
cos_q = cos_sin_cache[positions, :half].unsqueeze(1).to(q.dtype)
|
||||
sin_q = cos_sin_cache[positions, half:].unsqueeze(1).to(q.dtype)
|
||||
q_rope = q[..., nope_dim:].clone()
|
||||
q[..., nope_dim:][..., 0::2] = q_rope[..., 0::2] * cos_q - q_rope[..., 1::2] * sin_q
|
||||
q[..., nope_dim:][..., 1::2] = q_rope[..., 0::2] * sin_q + q_rope[..., 1::2] * cos_q
|
||||
|
||||
# GPT-J RoPE on KV + FP8 quant + cache insert
|
||||
if quant_to_fp8 and swa_kv_cache_2d.dtype == torch.uint8:
|
||||
# FP8 quantize: kv_bf16 → fp8
|
||||
kv_f32 = kv.float()
|
||||
amax = kv_f32.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
|
||||
scale = amax / 448.0 # FP8 E4M3 max = 448
|
||||
kv_fp8 = (kv_f32 / scale).to(torch.float8_e4m3fn)
|
||||
# Write to paged cache
|
||||
valid = slot_mapping >= 0
|
||||
slots = slot_mapping[valid]
|
||||
if slots.numel() > 0:
|
||||
swa_kv_cache_2d[slots] = kv_fp8[valid].view(torch.uint8)
|
||||
elif swa_kv_cache_2d.numel() > 0:
|
||||
valid = slot_mapping >= 0
|
||||
slots = slot_mapping[valid]
|
||||
if slots.numel() > 0:
|
||||
swa_kv_cache_2d[slots] = kv[valid]
|
||||
|
||||
|
||||
def csa_swa_attention_py(
|
||||
q: torch.Tensor, # (T, num_heads, head_dim) — already has RoPE
|
||||
swa_kv_cache: torch.Tensor, # paged KV cache (BF16 or FP8)
|
||||
swa_block_table: torch.Tensor,
|
||||
swa_slot_mapping: torch.Tensor,
|
||||
swa_block_size: int,
|
||||
window_size: int,
|
||||
positions: torch.Tensor,
|
||||
scale: float,
|
||||
) -> torch.Tensor:
|
||||
"""Sliding window attention using PyTorch (works on all GPUs).
|
||||
|
||||
Gathers KV from the sliding window in the paged cache,
|
||||
then does scaled dot-product attention.
|
||||
"""
|
||||
T, NH, HD = q.shape
|
||||
device = q.device
|
||||
|
||||
# Dequantize FP8 cache
|
||||
if swa_kv_cache.dtype == torch.uint8:
|
||||
cache_bf16 = swa_kv_cache.view(torch.float8_e4m3fn).to(torch.bfloat16)
|
||||
else:
|
||||
cache_bf16 = swa_kv_cache.to(torch.bfloat16)
|
||||
|
||||
# Flatten cache: (total_slots, head_dim)
|
||||
num_blocks, bs, kv_dim = cache_bf16.shape
|
||||
cache_flat = cache_bf16.reshape(-1, kv_dim)
|
||||
|
||||
# For each token, gather its local window of KV
|
||||
# Simplification: assume single sequence (for production, use block_table per req)
|
||||
max_window = min(T, window_size)
|
||||
kv_gathered = torch.zeros(T, max_window, HD, dtype=torch.bfloat16, device=device)
|
||||
for t in range(T):
|
||||
start = max(0, t - window_size + 1)
|
||||
length = t - start + 1
|
||||
for i, p in enumerate(range(start, t + 1)):
|
||||
if p < cache_flat.shape[0]:
|
||||
kv_gathered[t, i] = cache_flat[p]
|
||||
|
||||
# Expand KV for all heads: (T, max_window, HD) → (T, NH, max_window, HD)
|
||||
k_heads = kv_gathered.unsqueeze(1).expand(-1, NH, -1, -1).contiguous()
|
||||
v_heads = k_heads.clone()
|
||||
|
||||
# Q: (T, NH, HD) → (T, NH, 1, HD)
|
||||
q_4d = q.unsqueeze(2)
|
||||
|
||||
# Attention scores
|
||||
scores = torch.matmul(q_4d, k_heads.transpose(-1, -2)) * scale # (T, NH, 1, max_window)
|
||||
|
||||
# Causal + window mask
|
||||
for t in range(T):
|
||||
start = max(0, t - window_size + 1)
|
||||
length = t - start + 1
|
||||
if length < max_window:
|
||||
scores[t, :, :, length:] = float('-inf')
|
||||
|
||||
weights = F.softmax(scores.float(), dim=-1).to(torch.bfloat16)
|
||||
out = torch.matmul(weights, v_heads) # (T, NH, 1, HD)
|
||||
return out.squeeze(2) # (T, NH, HD)
|
||||
q_rope = q[:, :, nope_dim:].clone()
|
||||
q[:, :, nope_dim:][:, :, 0::2] = q_rope[:, :, 0::2] * cos_q - q_rope[:, :, 1::2] * sin_q
|
||||
q[:, :, nope_dim:][:, :, 1::2] = q_rope[:, :, 0::2] * sin_q + q_rope[:, :, 1::2] * cos_q
|
||||
|
||||
|
||||
def full_sdpa_attention(
|
||||
q: torch.Tensor, # (T, NH, HD) with RoPE
|
||||
kv: torch.Tensor, # (T, HD) KV latent (after norm, before cache)
|
||||
q: torch.Tensor, # (T, NH, HD) with RoPE
|
||||
kv: torch.Tensor, # (T, HD) KV latent
|
||||
scale: float,
|
||||
) -> torch.Tensor:
|
||||
"""Full self-attention fallback using PyTorch.
|
||||
|
||||
Used when KV cache is not yet populated (first forward, testing, etc.)
|
||||
or for SWA-only layers that don't need compression.
|
||||
"""
|
||||
"""Full self-attention using manual matmuls (works on all GPUs)."""
|
||||
T, NH, HD = q.shape
|
||||
|
||||
# Expand KV for all heads
|
||||
kv_exp = kv.unsqueeze(1).expand(-1, NH, -1).contiguous() # (T, NH, HD)
|
||||
|
||||
# Reshape for manual attention (batch per query token × head)
|
||||
kv_exp = kv.unsqueeze(1).expand(-1, NH, -1).contiguous()
|
||||
q_2d = q.reshape(T * NH, 1, HD)
|
||||
# Each query attends to all KV positions up to its own position
|
||||
k_2d = kv_exp.permute(1, 0, 2).unsqueeze(1).expand(NH, T, T, -1).contiguous().reshape(T * NH, T, HD)
|
||||
v_2d = k_2d.clone()
|
||||
|
||||
# Manual attention with causal mask
|
||||
scores = torch.matmul(q_2d, k_2d.transpose(-1, -2)) * scale
|
||||
query_pos = torch.arange(T, device=q.device).unsqueeze(1).repeat(1, NH).reshape(T * NH)
|
||||
kv_pos = torch.arange(T, device=q.device).unsqueeze(0)
|
||||
|
||||
Reference in New Issue
Block a user