250 lines
9.6 KiB
Python
250 lines
9.6 KiB
Python
import torch
|
|
from torch import nn
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
from flash_attn.flash_attn_interface import flash_attn_varlen_func, flash_attn_with_kvcache
|
|
from nanovllm.utils.context import get_context
|
|
|
|
|
|
@triton.jit
|
|
def store_kvcache_kernel(
|
|
key_ptr,
|
|
key_stride,
|
|
value_ptr,
|
|
value_stride,
|
|
k_cache_ptr,
|
|
v_cache_ptr,
|
|
slot_mapping_ptr,
|
|
D: tl.constexpr,
|
|
):
|
|
idx = tl.program_id(0)
|
|
slot = tl.load(slot_mapping_ptr + idx)
|
|
if slot == -1: return
|
|
key_offsets = idx * key_stride + tl.arange(0, D)
|
|
value_offsets = idx * value_stride + tl.arange(0, D)
|
|
key = tl.load(key_ptr + key_offsets)
|
|
value = tl.load(value_ptr + value_offsets)
|
|
cache_offsets = slot * D + tl.arange(0, D)
|
|
tl.store(k_cache_ptr + cache_offsets, key)
|
|
tl.store(v_cache_ptr + cache_offsets, value)
|
|
|
|
|
|
def store_kvcache(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, slot_mapping: torch.Tensor):
|
|
N, num_heads, head_dim = key.shape
|
|
D = num_heads * head_dim
|
|
assert key.stride(-1) == 1 and value.stride(-1) == 1
|
|
assert key.stride(1) == head_dim and value.stride(1) == head_dim
|
|
assert k_cache.stride(1) == D and v_cache.stride(1) == D
|
|
assert slot_mapping.numel() == N
|
|
store_kvcache_kernel[(N,)](key, key.stride(0), value, value.stride(0), k_cache, v_cache, slot_mapping, D)
|
|
|
|
|
|
class Attention(nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
num_heads,
|
|
head_dim,
|
|
scale,
|
|
num_kv_heads,
|
|
):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.head_dim = head_dim
|
|
self.scale = scale
|
|
self.num_kv_heads = num_kv_heads
|
|
self.k_cache = self.v_cache = torch.tensor([])
|
|
# Layer ID set by model_runner after model creation
|
|
self.layer_id: int = -1
|
|
|
|
def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
|
|
context = get_context()
|
|
k_cache, v_cache = self.k_cache, self.v_cache
|
|
if k_cache.numel() and v_cache.numel():
|
|
store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
|
|
|
|
if context.is_prefill:
|
|
if context.is_chunked_prefill:
|
|
# Chunked prefill: merge attention from previous KV
|
|
o = self._chunked_prefill_attention(q, k, v, context)
|
|
elif context.block_tables is not None: # prefix cache
|
|
k, v = k_cache, v_cache
|
|
o = flash_attn_varlen_func(q, k, v,
|
|
max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q,
|
|
max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k,
|
|
softmax_scale=self.scale, causal=True, block_table=context.block_tables)
|
|
else:
|
|
o = flash_attn_varlen_func(q, k, v,
|
|
max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q,
|
|
max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k,
|
|
softmax_scale=self.scale, causal=True, block_table=context.block_tables)
|
|
else: # decode
|
|
if context.is_chunked_prefill:
|
|
# Chunked decode: need to load all KV from CPU+GPU
|
|
o = self._chunked_decode_attention(q, k, v, context)
|
|
else:
|
|
o = flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache,
|
|
cache_seqlens=context.context_lens, block_table=context.block_tables,
|
|
softmax_scale=self.scale, causal=True)
|
|
return o
|
|
|
|
def _chunked_prefill_attention(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
context,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute attention with chunked KV from CPU cache.
|
|
|
|
For chunked prefill:
|
|
1. Load previous KV from CPU for this layer
|
|
2. Compute attention against previous KV (no causal mask)
|
|
3. Compute attention against current chunk's KV (causal)
|
|
4. Merge results using online softmax
|
|
"""
|
|
from nanovllm.kvcache.chunked_attention import flash_attn_with_lse, merge_attention_outputs
|
|
|
|
# q, k, v shape: [total_tokens, num_heads, head_dim]
|
|
total_tokens = q.shape[0]
|
|
|
|
# Reshape for flash attention: [batch, seq, heads, dim]
|
|
q_batched = q.unsqueeze(0) # [1, total_tokens, heads, dim]
|
|
k_batched = k.unsqueeze(0)
|
|
v_batched = v.unsqueeze(0)
|
|
|
|
accumulated_o = None
|
|
accumulated_lse = None
|
|
|
|
# Load previous KV from CPU for this layer
|
|
if context.offload_engine is not None and self.layer_id >= 0:
|
|
# Get the kvcache_manager from context
|
|
kvcache_manager = context.offload_engine
|
|
|
|
# For each sequence in the chunk, load previous KV
|
|
# Currently assuming single sequence
|
|
if hasattr(context, 'chunked_seq') and context.chunked_seq is not None:
|
|
prev_k, prev_v = kvcache_manager.load_prev_kv_for_layer(
|
|
context.chunked_seq,
|
|
self.layer_id,
|
|
)
|
|
|
|
if prev_k is not None and prev_v is not None:
|
|
# Compute attention against previous KV (no causal mask)
|
|
prev_o, prev_lse = flash_attn_with_lse(
|
|
q_batched,
|
|
prev_k,
|
|
prev_v,
|
|
softmax_scale=self.scale,
|
|
causal=False, # No causal mask for previous context
|
|
)
|
|
accumulated_o = prev_o
|
|
accumulated_lse = prev_lse
|
|
|
|
# Compute attention against current chunk's KV (with causal mask)
|
|
current_o, current_lse = flash_attn_with_lse(
|
|
q_batched,
|
|
k_batched,
|
|
v_batched,
|
|
softmax_scale=self.scale,
|
|
causal=True, # Causal mask for current chunk
|
|
)
|
|
|
|
# Merge with accumulated
|
|
if accumulated_o is None:
|
|
final_o = current_o
|
|
else:
|
|
final_o, _ = merge_attention_outputs(
|
|
accumulated_o, accumulated_lse,
|
|
current_o, current_lse,
|
|
)
|
|
|
|
# Remove batch dimension: [1, total_tokens, heads, dim] -> [total_tokens, heads, dim]
|
|
return final_o.squeeze(0)
|
|
|
|
def _chunked_decode_attention(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
context,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute decode attention with KV spread across CPU and GPU.
|
|
|
|
Uses chunked attention similar to chunked prefill:
|
|
1. Process blocks on GPU first (if any)
|
|
2. Load CPU blocks in chunks to GPU slots (per-layer)
|
|
3. Compute attention for each chunk, merge with online softmax
|
|
"""
|
|
from nanovllm.kvcache.chunked_attention import flash_attn_with_lse, merge_attention_outputs
|
|
|
|
# q shape: [batch_size, num_heads, head_dim] (single decode token per sequence)
|
|
# Need: [batch, seqlen, heads, dim]
|
|
q_batched = q.unsqueeze(1) # [batch, 1, heads, dim]
|
|
|
|
kvcache_manager = context.offload_engine
|
|
seq = context.chunked_seq
|
|
|
|
o_acc = None
|
|
lse_acc = None
|
|
|
|
# Step 1: Process blocks already on GPU (if any)
|
|
gpu_slots, _ = kvcache_manager.get_gpu_blocks_for_decode(seq)
|
|
if gpu_slots:
|
|
k_gpu, v_gpu = kvcache_manager.get_kv_for_gpu_slots(self.layer_id, gpu_slots)
|
|
o_gpu, lse_gpu = flash_attn_with_lse(
|
|
q_batched, k_gpu, v_gpu,
|
|
softmax_scale=self.scale,
|
|
causal=False,
|
|
)
|
|
o_acc, lse_acc = o_gpu, lse_gpu
|
|
|
|
# Step 2: Process CPU blocks in chunks
|
|
# Get chunk info from kvcache_manager
|
|
cpu_block_ids, cpu_logical_ids, num_chunks = kvcache_manager.get_decode_chunk_info(seq)
|
|
|
|
if num_chunks > 0:
|
|
# Use num_gpu_slots - 1 to avoid the reserved slot (used for write block)
|
|
chunk_size = kvcache_manager.num_gpu_slots - 1
|
|
|
|
for chunk_idx in range(num_chunks):
|
|
start = chunk_idx * chunk_size
|
|
end = min(start + chunk_size, len(cpu_block_ids))
|
|
chunk_cpu_ids = cpu_block_ids[start:end]
|
|
|
|
# Load this chunk to GPU slots 0, 1, 2, ... for THIS LAYER
|
|
# (slot num_gpu_slots-1 is reserved for write block)
|
|
gpu_slots_for_chunk = list(range(len(chunk_cpu_ids)))
|
|
kvcache_manager.offload_engine.load_cpu_blocks_to_gpu_slots(
|
|
self.layer_id,
|
|
chunk_cpu_ids,
|
|
gpu_slots_for_chunk,
|
|
)
|
|
|
|
# Get KV for this chunk
|
|
k_chunk, v_chunk = kvcache_manager.get_kv_for_gpu_slots(
|
|
self.layer_id, gpu_slots_for_chunk
|
|
)
|
|
|
|
# Compute attention for this chunk
|
|
o_chunk, lse_chunk = flash_attn_with_lse(
|
|
q_batched, k_chunk, v_chunk,
|
|
softmax_scale=self.scale,
|
|
causal=False,
|
|
)
|
|
|
|
# Merge with accumulated
|
|
if o_acc is None:
|
|
o_acc, lse_acc = o_chunk, lse_chunk
|
|
else:
|
|
o_acc, lse_acc = merge_attention_outputs(o_acc, lse_acc, o_chunk, lse_chunk)
|
|
|
|
if o_acc is None:
|
|
raise RuntimeError("Chunked decode attention failed: no KV available")
|
|
|
|
# Output shape: [batch, 1, heads, dim] -> [batch, heads, dim]
|
|
return o_acc.squeeze(1)
|