510 lines
20 KiB
Python
510 lines
20 KiB
Python
"""
|
|
XAttention Block Sparse Attention (BSA) Policy for nano-vllm.
|
|
|
|
This module implements XAttention-inspired block sparse attention for chunked prefill,
|
|
using block-level estimation to select important KV blocks for computation.
|
|
|
|
Reference: COMPASS/compass/src/Xattention.py
|
|
"""
|
|
|
|
import math
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from typing import List, Optional, Tuple
|
|
|
|
from nanovllm.kvcache.sparse.policy import SparsePolicy, PolicyContext
|
|
from nanovllm.utils.context import get_context
|
|
|
|
|
|
class XAttentionBSAPolicy(SparsePolicy):
|
|
"""
|
|
XAttention Block Sparse Attention policy for chunked prefill.
|
|
|
|
This policy uses block-level estimation to determine which KV blocks
|
|
are important for the current chunk's queries, enabling sparse computation.
|
|
|
|
Key features:
|
|
- Double-loading design: estimate phase loads samples, compute phase loads selected blocks
|
|
- Block-level granularity: 128-token blocks for estimation and computation
|
|
- Triton kernels for efficient estimation (optional, falls back to PyTorch)
|
|
|
|
Architecture:
|
|
1. Estimate Phase: Load samples from all historical chunks, compute importance scores
|
|
2. Selection Phase: Select top chunks by cumulative attention threshold
|
|
3. Compute Phase: Load selected chunks fully, apply block sparse attention
|
|
"""
|
|
|
|
supports_prefill = True
|
|
supports_decode = False # BSA is prefill-only
|
|
requires_block_selection = False # Selection happens at chunk level, not block level
|
|
|
|
def __init__(
|
|
self,
|
|
block_size: int = 128,
|
|
samples_per_chunk: int = 128,
|
|
threshold: float = 0.9,
|
|
use_triton: bool = True,
|
|
stride: int = 8,
|
|
):
|
|
"""
|
|
Initialize XAttention BSA policy.
|
|
|
|
Args:
|
|
block_size: Number of tokens per block (default: 128)
|
|
samples_per_chunk: Number of tokens to sample from each historical chunk for estimation
|
|
threshold: Cumulative attention threshold for chunk selection (0-1)
|
|
use_triton: Use Triton kernels for estimation (requires SM 80+)
|
|
stride: Stride for Q/K downsampling in estimation
|
|
"""
|
|
self.block_size = block_size
|
|
self.samples_per_chunk = samples_per_chunk
|
|
self.threshold = threshold
|
|
self.use_triton = use_triton
|
|
self.stride = stride
|
|
|
|
# Check Triton availability
|
|
if self.use_triton:
|
|
try:
|
|
import triton
|
|
props = torch.cuda.get_device_properties(torch.cuda.current_device())
|
|
if props.major < 8:
|
|
self.use_triton = False
|
|
print(f"[XAttentionBSA] Triton requires SM 80+, got SM {props.major}{props.minor}. Falling back to PyTorch.")
|
|
except ImportError:
|
|
self.use_triton = False
|
|
print("[XAttentionBSA] Triton not available. Using PyTorch implementation.")
|
|
|
|
def select_blocks(self, available_blocks: List[int], ctx: PolicyContext) -> List[int]:
|
|
"""
|
|
Select blocks to load from CPU (for decode compatibility, not used in prefill).
|
|
|
|
For prefill, BSA handles chunk-level selection internally.
|
|
"""
|
|
# For prefill, we return all blocks - selection happens in sparse_prefill_attention
|
|
return available_blocks
|
|
|
|
def sparse_prefill_attention(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
layer_id: int,
|
|
softmax_scale: float,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute XAttention block sparse attention for current chunk.
|
|
|
|
This implements a simplified version that loads all historical chunks
|
|
(sparse selection to be implemented in next phase).
|
|
|
|
Args:
|
|
q: Query tensor [seq_len, num_heads, head_dim]
|
|
k: Key tensor [seq_len, num_kv_heads, head_dim] (unused, we use prefill buffer)
|
|
v: Value tensor [seq_len, num_kv_heads, head_dim] (unused, we use prefill buffer)
|
|
layer_id: Current transformer layer index
|
|
softmax_scale: Softmax scaling factor from attention layer
|
|
|
|
Returns:
|
|
Attention output [seq_len, num_heads, head_dim]
|
|
"""
|
|
from nanovllm.kvcache.chunked_attention import flash_attn_with_lse, merge_attention_outputs
|
|
|
|
context = get_context()
|
|
kvcache_manager = context.kvcache_manager
|
|
offload_engine = kvcache_manager.offload_engine if kvcache_manager else None
|
|
|
|
if offload_engine is None:
|
|
# No offload engine, use standard attention with provided k, v
|
|
return self._full_attention(q, k, v, causal=True)
|
|
|
|
current_chunk_idx = getattr(context, 'current_chunk_idx', 0)
|
|
seq = getattr(context, 'chunked_seq', None)
|
|
num_tokens = q.shape[0]
|
|
|
|
if seq is None:
|
|
# No chunked sequence, fallback to full attention on current chunk only
|
|
return self._full_attention(q, k, v, causal=True)
|
|
|
|
# Get prefilled CPU blocks (historical chunks)
|
|
cpu_block_table = kvcache_manager.get_prefilled_cpu_blocks(seq)
|
|
|
|
q_batched = q.unsqueeze(0) # [1, seq_len, num_heads, head_dim]
|
|
o_acc = None
|
|
lse_acc = None
|
|
|
|
# Get compute stream for all attention operations
|
|
compute_stream = offload_engine.compute_stream
|
|
|
|
# Step 1: Load historical chunks from CPU using slot mechanism
|
|
if cpu_block_table:
|
|
load_slots = list(range(offload_engine.num_ring_slots))
|
|
num_blocks = len(cpu_block_table)
|
|
|
|
# Load ALL historical blocks (not just min(num_blocks, num_slots))
|
|
# Use synchronous mode like standard flow when pipeline_depth=1
|
|
if len(load_slots) == 1:
|
|
# Only 1 slot available, cannot pipeline - use synchronous mode
|
|
slot = load_slots[0]
|
|
for block_idx in range(num_blocks):
|
|
cpu_block_id = cpu_block_table[block_idx]
|
|
offload_engine.load_to_slot_layer(slot, layer_id, cpu_block_id)
|
|
offload_engine.wait_slot_layer(slot)
|
|
|
|
with torch.cuda.stream(compute_stream):
|
|
# Get KV from slot - returns [1, block_size, kv_heads, head_dim]
|
|
prev_k, prev_v = offload_engine.get_kv_for_slot(slot)
|
|
|
|
# Compute attention to historical chunk (non-causal, already processed)
|
|
prev_o, prev_lse = flash_attn_with_lse(
|
|
q_batched, prev_k, prev_v,
|
|
softmax_scale=softmax_scale,
|
|
causal=False,
|
|
)
|
|
|
|
# Merge results
|
|
if o_acc is None:
|
|
o_acc, lse_acc = prev_o, prev_lse
|
|
else:
|
|
o_acc, lse_acc = merge_attention_outputs(o_acc, lse_acc, prev_o, prev_lse)
|
|
|
|
# Record compute done so slot can be reused
|
|
offload_engine.record_slot_compute_done(slot)
|
|
else:
|
|
# Multiple slots available - use pipeline
|
|
num_slots = len(load_slots)
|
|
|
|
# Phase 1: Pre-load up to num_slots blocks to fill the pipeline
|
|
num_preload = min(num_slots, num_blocks)
|
|
for i in range(num_preload):
|
|
offload_engine.load_to_slot_layer(load_slots[i], layer_id, cpu_block_table[i])
|
|
|
|
# Phase 2: Main loop - compute and immediately reuse slot for next transfer
|
|
for block_idx in range(num_blocks):
|
|
# Cycle through slots: slot[block_idx % num_slots]
|
|
current_slot = load_slots[block_idx % num_slots]
|
|
cpu_block_id = cpu_block_table[block_idx]
|
|
|
|
# Wait for current slot's transfer to complete
|
|
offload_engine.wait_slot_layer(current_slot)
|
|
|
|
# Compute attention on current slot's data
|
|
with torch.cuda.stream(compute_stream):
|
|
# Get KV from slot - returns [1, block_size, kv_heads, head_dim]
|
|
prev_k, prev_v = offload_engine.get_kv_for_slot(current_slot)
|
|
|
|
# Compute attention to historical chunk (non-causal, already processed)
|
|
prev_o, prev_lse = flash_attn_with_lse(
|
|
q_batched, prev_k, prev_v,
|
|
softmax_scale=softmax_scale,
|
|
causal=False,
|
|
)
|
|
|
|
# Merge results
|
|
if o_acc is None:
|
|
o_acc, lse_acc = prev_o, prev_lse
|
|
else:
|
|
o_acc, lse_acc = merge_attention_outputs(o_acc, lse_acc, prev_o, prev_lse)
|
|
|
|
# Record compute done so slot can be reused
|
|
offload_engine.record_slot_compute_done(current_slot)
|
|
|
|
# Issue next transfer if there are more blocks
|
|
next_block_idx = block_idx + num_slots
|
|
if next_block_idx < num_blocks:
|
|
next_slot = load_slots[next_block_idx % num_slots]
|
|
next_cpu_block_id = cpu_block_table[next_block_idx]
|
|
offload_engine.load_to_slot_layer(next_slot, layer_id, next_cpu_block_id)
|
|
|
|
# Step 2: Compute attention to current chunk (causal mask) - use prefill buffer on compute_stream
|
|
with torch.cuda.stream(compute_stream):
|
|
k_curr, v_curr = offload_engine.get_prefill_buffer_slice(layer_id, num_tokens)
|
|
|
|
current_o, current_lse = flash_attn_with_lse(
|
|
q_batched,
|
|
k_curr,
|
|
v_curr,
|
|
softmax_scale=softmax_scale,
|
|
causal=True,
|
|
)
|
|
|
|
# Step 3: Merge historical and current attention
|
|
with torch.cuda.stream(compute_stream):
|
|
if o_acc is None:
|
|
# No historical chunks processed
|
|
final_o = current_o
|
|
else:
|
|
final_o, _ = merge_attention_outputs(o_acc, lse_acc, current_o, current_lse)
|
|
|
|
# Sync default stream with compute_stream before returning
|
|
torch.cuda.default_stream().wait_stream(compute_stream)
|
|
|
|
# Remove batch dimension: [1, seq_len, num_heads, head_dim] -> [seq_len, num_heads, head_dim]
|
|
return final_o.squeeze(0)
|
|
|
|
def _estimate_historical_chunks(
|
|
self,
|
|
q: torch.Tensor,
|
|
historical_blocks: List[int],
|
|
layer_id: int,
|
|
current_chunk_idx: int,
|
|
) -> Tuple[List[float], bool]:
|
|
"""
|
|
Estimate importance of each historical chunk for current Q.
|
|
|
|
First load: Load samples from each historical chunk for estimation.
|
|
|
|
Args:
|
|
q: Current chunk queries [chunk_size, num_heads, head_dim]
|
|
historical_blocks: List of historical CPU block IDs
|
|
layer_id: Current layer index
|
|
current_chunk_idx: Current chunk index
|
|
|
|
Returns:
|
|
(List of importance scores (one per historical chunk), has_valid_data flag)
|
|
has_valid_data is True if at least one block had non-zero data
|
|
"""
|
|
chunk_estimates = []
|
|
has_valid_data = False
|
|
|
|
for block_idx, cpu_block_id in enumerate(historical_blocks):
|
|
# First load: Load sample from this historical chunk
|
|
k_sample, v_sample = self._load_block_sample(
|
|
cpu_block_id, layer_id, self.samples_per_chunk
|
|
)
|
|
|
|
# Check if loaded data is valid (non-zero)
|
|
if k_sample.abs().max().item() > 0:
|
|
has_valid_data = True
|
|
|
|
# Quick estimation: Compute Q attention to this chunk's sample
|
|
# q [chunk_size, H, D] @ k_sample [samples, H, D]
|
|
# Result: Aggregate to chunk-level score
|
|
estimate = self._compute_chunk_estimate(q, k_sample)
|
|
chunk_estimates.append(estimate)
|
|
|
|
return chunk_estimates, has_valid_data
|
|
|
|
def _select_important_chunks(
|
|
self,
|
|
chunk_estimates: List[float],
|
|
) -> List[int]:
|
|
"""
|
|
Select important chunks based on cumulative attention threshold.
|
|
|
|
Args:
|
|
chunk_estimates: Importance scores for each historical chunk
|
|
|
|
Returns:
|
|
Indices of selected chunks
|
|
"""
|
|
if not chunk_estimates:
|
|
return []
|
|
|
|
scores = torch.tensor(chunk_estimates, device='cpu')
|
|
threshold_value = scores.max() * self.threshold
|
|
|
|
# Select chunks that contribute to cumulative attention threshold
|
|
selected_indices = []
|
|
cumulative = 0.0
|
|
sorted_indices = torch.argsort(scores, descending=True)
|
|
|
|
for idx in sorted_indices:
|
|
cumulative += scores[idx].item()
|
|
selected_indices.append(idx.item())
|
|
if cumulative >= threshold_value:
|
|
break
|
|
|
|
return selected_indices
|
|
|
|
def _compute_with_selected_chunks(
|
|
self,
|
|
q: torch.Tensor,
|
|
historical_blocks: List[int],
|
|
selected_indices: List[int],
|
|
layer_id: int,
|
|
current_chunk_idx: int,
|
|
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
|
|
"""
|
|
Compute attention to selected historical chunks.
|
|
|
|
Second load: Load full data for selected chunks.
|
|
|
|
Args:
|
|
q: Current chunk queries
|
|
historical_blocks: All historical block IDs
|
|
selected_indices: Indices of selected blocks
|
|
layer_id: Current layer index
|
|
current_chunk_idx: Current chunk index
|
|
|
|
Returns:
|
|
(accumulated_output, accumulated_lse) or (None, None)
|
|
"""
|
|
if not selected_indices:
|
|
return None, None
|
|
|
|
o_acc = None
|
|
lse_acc = None
|
|
|
|
for chunk_idx in selected_indices:
|
|
cpu_block_id = historical_blocks[chunk_idx]
|
|
|
|
# Second load: Load full data for this selected chunk
|
|
k_full, v_full = self._load_block_full(
|
|
cpu_block_id, layer_id
|
|
)
|
|
|
|
# Compute attention (non-causal, already processed)
|
|
o, lse = self._full_attention(
|
|
q.unsqueeze(0), k_full.unsqueeze(0),
|
|
v_full.unsqueeze(0), causal=False, return_lse=True
|
|
)
|
|
|
|
# Merge results
|
|
if o_acc is None:
|
|
o_acc, lse_acc = o.squeeze(0), lse
|
|
else:
|
|
from nanovllm.kvcache.chunked_attention import merge_attention_outputs
|
|
o_acc, lse_acc = merge_attention_outputs(
|
|
o_acc.unsqueeze(0), lse_acc,
|
|
o.unsqueeze(0), lse
|
|
)
|
|
o_acc = o_acc.squeeze(0)
|
|
|
|
return o_acc, lse_acc
|
|
|
|
def _load_block_sample(
|
|
self,
|
|
cpu_block_id: int,
|
|
layer_id: int,
|
|
num_samples: int,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""Load sample tokens from a CPU block."""
|
|
offload_engine = get_context().kvcache_manager.offload_engine
|
|
|
|
k_sample, v_sample = offload_engine.load_block_sample_from_cpu(
|
|
cpu_block_id, layer_id, num_samples
|
|
)
|
|
return k_sample, v_sample
|
|
|
|
def _load_block_full(
|
|
self,
|
|
cpu_block_id: int,
|
|
layer_id: int,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""Load full tokens from a CPU block."""
|
|
offload_engine = get_context().kvcache_manager.offload_engine
|
|
return offload_engine.load_block_full_from_cpu(
|
|
cpu_block_id, layer_id
|
|
)
|
|
|
|
def _compute_chunk_estimate(
|
|
self,
|
|
q: torch.Tensor,
|
|
k_sample: torch.Tensor,
|
|
) -> float:
|
|
"""
|
|
Compute chunk-level importance estimate.
|
|
|
|
Args:
|
|
q: [chunk_size, num_heads, head_dim]
|
|
k_sample: [num_samples, num_kv_heads, head_dim]
|
|
|
|
Returns:
|
|
Aggregate importance score for this chunk
|
|
"""
|
|
# Expand K to match Q's head count (GQA support)
|
|
num_heads = q.shape[1]
|
|
num_kv_heads = k_sample.shape[1]
|
|
head_dim = q.shape[2] # Last dimension is head_dim
|
|
if num_heads != num_kv_heads:
|
|
repeat_factor = num_heads // num_kv_heads
|
|
k_sample = k_sample.repeat_interleave(repeat_factor, dim=1)
|
|
|
|
# Compute attention scores: Q @ K.T with proper scaling
|
|
# q [chunk_size, H, D], k [samples, H, D] -> need to compute per-head attention
|
|
# Use scaled dot-product attention: (Q @ K.T) / sqrt(D)
|
|
scale = 1.0 / (head_dim ** 0.5)
|
|
|
|
# Reshape to 2D: [chunk_size * H, D] @ [D, samples * H] then aggregate
|
|
chunk_size = q.shape[0]
|
|
num_samples = k_sample.shape[0]
|
|
|
|
# Reshape for batched matmul: merge heads and seq dims
|
|
q_2d = q.reshape(chunk_size * num_heads, head_dim) # [chunk_size*H, D]
|
|
k_2d = k_sample.reshape(num_samples * num_heads, head_dim) # [samples*H, D]
|
|
|
|
# Compute scaled Q @ K.T: [chunk_size*H, D] @ [D, samples*H] = [chunk_size*H, samples*H]
|
|
attn_scores_2d = torch.matmul(q_2d, k_2d.T) * scale
|
|
|
|
# Use max absolute value as importance (captures both positive and negative attention)
|
|
importance = attn_scores_2d.abs().max().item()
|
|
|
|
return importance
|
|
|
|
def _full_attention(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
causal: bool = False,
|
|
return_lse: bool = False,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute full FlashAttention (fallback when sparse not applicable).
|
|
|
|
Args:
|
|
q: [batch_size, seq_len, num_heads, head_dim] or [seq_len, num_heads, head_dim]
|
|
k, v: Same shape as q
|
|
causal: Apply causal mask
|
|
return_lse: Whether to return log-sum-exp
|
|
|
|
Returns:
|
|
attention output [batch_size, seq_len, num_heads, head_dim] or [seq_len, num_heads, head_dim]
|
|
"""
|
|
from nanovllm.kvcache.chunked_attention import flash_attn_with_lse
|
|
|
|
# Handle 3D input: add batch dimension
|
|
input_3d = q.dim() == 3
|
|
if input_3d:
|
|
q = q.unsqueeze(0) # [seq_len, H, D] -> [1, seq_len, H, D]
|
|
k = k.unsqueeze(0)
|
|
v = v.unsqueeze(0)
|
|
|
|
if return_lse:
|
|
o, lse = flash_attn_with_lse(q, k, v, softmax_scale=self.scale, causal=causal)
|
|
result = (o, lse)
|
|
else:
|
|
o, _ = flash_attn_with_lse(q, k, v, softmax_scale=self.scale, causal=causal)
|
|
result = o
|
|
|
|
# Remove batch dimension if input was 3D
|
|
if input_3d:
|
|
if return_lse:
|
|
result = (result[0].squeeze(0), result[1])
|
|
else:
|
|
result = result.squeeze(0)
|
|
|
|
return result
|
|
|
|
@property
|
|
def scale(self) -> float:
|
|
"""Get softmax scale factor from Attention layer."""
|
|
context = get_context()
|
|
# Get scale from current Attention layer in the model
|
|
if hasattr(context, 'current_attention') and context.current_attention is not None:
|
|
return context.current_attention.scale
|
|
# Fallback: try to get from model runner
|
|
if hasattr(context, 'model_runner') and context.model_runner is not None:
|
|
model_runner = context.model_runner
|
|
if hasattr(model_runner, 'model') and hasattr(model_runner.model, 'layers'):
|
|
# Get scale from first attention layer
|
|
first_layer = model_runner.model.layers[0]
|
|
if hasattr(first_layer, 'self_attn'):
|
|
return first_layer.self_attn.scaling
|
|
# Default: 1 / sqrt(128) for Qwen models
|
|
return 1.0 / 128.0 ** 0.5
|
|
|
|
def reset(self) -> None:
|
|
"""Reset policy state."""
|
|
pass
|