diff --git a/vllm/patches/layers/csa_attention.py b/vllm/patches/layers/csa_attention.py index e26fabab..04ad23ff 100644 --- a/vllm/patches/layers/csa_attention.py +++ b/vllm/patches/layers/csa_attention.py @@ -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)