diff --git a/nanovllm/engine/model_runner.py b/nanovllm/engine/model_runner.py index 2013eae..e66a57e 100644 --- a/nanovllm/engine/model_runner.py +++ b/nanovllm/engine/model_runner.py @@ -429,7 +429,14 @@ class ModelRunner: else: return self.run_layerwise_offload_decode(seqs) - #> Following Code will not use Layer-wise Offload mode + #> Check if contiguous GPU mode should be used (single-seq optimization) + if self._should_use_contiguous_gpu_mode(seqs, is_prefill): + if is_prefill: + return self.run_gpu_only_prefill(seqs) + else: + return self.run_gpu_only_decode(seqs) + + #> Following Code uses standard PagedAttention path input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs) temperatures = self.prepare_sample(seqs) if self.rank == 0 else None logits = self.run_model(input_ids, positions, is_prefill) @@ -437,6 +444,257 @@ class ModelRunner: reset_context() return token_ids + def _should_use_contiguous_gpu_mode(self, seqs: list[Sequence], is_prefill: bool) -> bool: + """ + Check if contiguous GPU mode should be used for single-seq optimization. + + Conditions: + 1. Has kvcache_manager with contiguous cache allocated + 2. Not using CPU offload (no offload_engine) + 3. Single sequence (batch_size == 1) + 4. Has blocks allocated (not warmup) + """ + # Must have kvcache_manager + if not hasattr(self, 'kvcache_manager') or self.kvcache_manager is None: + return False + + # Must have contiguous cache + if not hasattr(self.kvcache_manager, 'contiguous_k_cache'): + return False + if self.kvcache_manager.contiguous_k_cache is None: + return False + + # Must NOT be offload mode + if hasattr(self.kvcache_manager, 'offload_engine'): + return False + + # Single sequence only + if len(seqs) != 1: + return False + + # Has blocks allocated (not warmup) + if not seqs[0].block_table: + return False + + return True + + # ========== Contiguous GPU-only Methods ========== + + @torch.inference_mode() + def run_gpu_only_prefill(self, seqs: list[Sequence]) -> list[int]: + """ + GPU-only prefill with contiguous KV cache layout. + + Mirrors run_layerwise_offload_prefill() but stores to GPU instead of CPU. + No scatter operations - just contiguous slice assignment. + + Key design: + - Process layer-by-layer (not via Attention.forward()) + - Store K,V to contiguous GPU cache (same layout as computed K,V) + - Use sparse prefill attention if enabled + """ + assert len(seqs) == 1, "GPU-only layer-wise prefill only supports single sequence" + seq = seqs[0] + + num_layers = len(self.model.model.layers) + total_tokens = len(seq) + + logger.debug(f"[GPU-only Prefill] Starting: {total_tokens} tokens, {num_layers} layers") + + # Get contiguous GPU cache + k_cache = self.kvcache_manager.contiguous_k_cache + v_cache = self.kvcache_manager.contiguous_v_cache + + # Prepare inputs + input_ids = torch.tensor(seq[:], dtype=torch.int64, device="cuda") + positions = torch.arange(total_tokens, dtype=torch.int64, device="cuda") + + # Import FlashAttention + from flash_attn.flash_attn_interface import flash_attn_varlen_func + cu_seqlens = torch.tensor([0, total_tokens], dtype=torch.int32, device="cuda") + + # Embedding + hidden_states = self.model.model.embed_tokens(input_ids) + residual = None + + # Layer-by-layer processing + for layer_id in range(num_layers): + layer = self.model.model.layers[layer_id] + + # Input LayerNorm + if residual is None: + hidden_ln, residual = layer.input_layernorm(hidden_states), hidden_states + else: + hidden_ln, residual = layer.input_layernorm(hidden_states, residual) + + # QKV projection + qkv = layer.self_attn.qkv_proj(hidden_ln) + q, k, v = qkv.split([ + layer.self_attn.q_size, + layer.self_attn.kv_size, + layer.self_attn.kv_size + ], dim=-1) + + q = q.view(total_tokens, layer.self_attn.num_heads, layer.self_attn.head_dim) + k = k.view(total_tokens, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + v = v.view(total_tokens, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + + # Q/K norms (Qwen3 specific) + if not layer.self_attn.qkv_bias: + num_tokens = q.shape[0] + q = layer.self_attn.q_norm(q.reshape(-1, layer.self_attn.head_dim)) + q = q.view(num_tokens, layer.self_attn.num_heads, layer.self_attn.head_dim) + k = layer.self_attn.k_norm(k.reshape(-1, layer.self_attn.head_dim)) + k = k.view(num_tokens, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + + # RoPE + q, k = layer.self_attn.rotary_emb(positions, q, k) + + # Sparse or Full attention (uses k, v directly - before store!) + if self.sparse_prefill_policy is not None: + attn_output = self.sparse_prefill_policy.sparse_prefill_attention( + q, k, v, layer_id + ) + else: + attn_output = flash_attn_varlen_func( + q, k, v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=total_tokens, + max_seqlen_k=total_tokens, + softmax_scale=layer.self_attn.attn.scale, + causal=True, + ) + + # O projection + attn_output = attn_output.view(total_tokens, -1) + hidden_states = layer.self_attn.o_proj(attn_output) + + # Store K,V to contiguous GPU cache AFTER attention (same as offload pattern) + k_cache[layer_id, :total_tokens] = k + v_cache[layer_id, :total_tokens] = v + + # Post-attention LayerNorm + MLP + hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) + hidden_states = layer.mlp(hidden_states) + + # Final norm + hidden_states, _ = self.model.model.norm(hidden_states, residual) + + # Compute logits for last token + logits = self.model.compute_logits(hidden_states[-1:]) + + # Record prefill length for decode + self.kvcache_manager.contiguous_seq_len = total_tokens + + logger.debug(f"[GPU-only Prefill] Complete: {num_layers} layers processed") + + # Sample + temperatures = self.prepare_sample(seqs) if self.rank == 0 else None + token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None + + return token_ids + + @torch.inference_mode() + def run_gpu_only_decode(self, seqs: list[Sequence]) -> list[int]: + """ + Decode using contiguous GPU KV cache. + + Similar to offload decode but simpler - all KV already on GPU. + """ + assert len(seqs) == 1, "GPU-only decode only supports single sequence" + seq = seqs[0] + + num_layers = len(self.model.model.layers) + k_cache = self.kvcache_manager.contiguous_k_cache + v_cache = self.kvcache_manager.contiguous_v_cache + context_len = self.kvcache_manager.contiguous_seq_len + + # Prepare inputs + input_ids = torch.tensor([seq.last_token], dtype=torch.int64, device="cuda") + positions = torch.tensor([len(seq) - 1], dtype=torch.int64, device="cuda") + + from flash_attn.flash_attn_interface import flash_attn_varlen_func + cu_seqlens_q = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + + # Embedding + hidden_states = self.model.model.embed_tokens(input_ids) + residual = None + + for layer_id in range(num_layers): + layer = self.model.model.layers[layer_id] + + # Input LayerNorm + if residual is None: + hidden_ln, residual = layer.input_layernorm(hidden_states), hidden_states + else: + hidden_ln, residual = layer.input_layernorm(hidden_states, residual) + + # QKV projection + qkv = layer.self_attn.qkv_proj(hidden_ln) + q, k_new, v_new = qkv.split([ + layer.self_attn.q_size, + layer.self_attn.kv_size, + layer.self_attn.kv_size + ], dim=-1) + + q = q.view(1, layer.self_attn.num_heads, layer.self_attn.head_dim) + k_new = k_new.view(1, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + v_new = v_new.view(1, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + + # Q/K norms + if not layer.self_attn.qkv_bias: + q = layer.self_attn.q_norm(q.reshape(-1, layer.self_attn.head_dim)) + q = q.view(1, layer.self_attn.num_heads, layer.self_attn.head_dim) + k_new = layer.self_attn.k_norm(k_new.reshape(-1, layer.self_attn.head_dim)) + k_new = k_new.view(1, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + + # RoPE + q, k_new = layer.self_attn.rotary_emb(positions, q, k_new) + + # Store new K,V to cache + k_cache[layer_id, context_len] = k_new.squeeze(0) + v_cache[layer_id, context_len] = v_new.squeeze(0) + + # Full K,V for attention (including new token) + k_full = k_cache[layer_id, :context_len + 1] + v_full = v_cache[layer_id, :context_len + 1] + + # Attention + cu_seqlens_k = torch.tensor([0, context_len + 1], dtype=torch.int32, device="cuda") + attn_output = flash_attn_varlen_func( + q, k_full, v_full, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=1, + max_seqlen_k=context_len + 1, + softmax_scale=layer.self_attn.attn.scale, + causal=False, # Single query, no causal needed + ) + + # O projection + attn_output = attn_output.view(1, -1) + hidden_states = layer.self_attn.o_proj(attn_output) + + # Post-attention LayerNorm + MLP + hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) + hidden_states = layer.mlp(hidden_states) + + # Update context length + self.kvcache_manager.contiguous_seq_len = context_len + 1 + + # Final norm + hidden_states, _ = self.model.model.norm(hidden_states, residual) + + # Compute logits + logits = self.model.compute_logits(hidden_states) + + # Sample + temperatures = self.prepare_sample(seqs) if self.rank == 0 else None + token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None + + return token_ids + def _should_use_layerwise_offload(self, seqs: list[Sequence], is_prefill: bool) -> bool: """ Check if layer-wise offload mode should be used. diff --git a/nanovllm/kvcache/__init__.py b/nanovllm/kvcache/__init__.py index 9946694..3e44b31 100644 --- a/nanovllm/kvcache/__init__.py +++ b/nanovllm/kvcache/__init__.py @@ -36,10 +36,11 @@ def create_kvcache_manager(config: "Config") -> KVCacheManager: KVCacheManager instance """ if not getattr(config, 'enable_cpu_offload', False): - # Default: pure GPU mode + # Default: pure GPU mode with contiguous cache for single-seq optimization return GPUOnlyManager( num_blocks=config.num_kvcache_blocks, block_size=config.kvcache_block_size, + max_seq_len=config.max_model_len, # Enable contiguous cache ) # CPU offload is enabled diff --git a/nanovllm/kvcache/gpu_manager.py b/nanovllm/kvcache/gpu_manager.py index ad8e40f..ad3f49c 100644 --- a/nanovllm/kvcache/gpu_manager.py +++ b/nanovllm/kvcache/gpu_manager.py @@ -45,21 +45,24 @@ class GPUOnlyManager(KVCacheManager): - Paged attention with configurable block size - Prefix caching via xxhash - Reference counting for block sharing + - Contiguous cache for single-sequence layer-wise prefill (optional) This manager is fully compatible with CUDA graphs since all data stays on GPU at fixed addresses. """ - def __init__(self, num_blocks: int, block_size: int): + def __init__(self, num_blocks: int, block_size: int, max_seq_len: int = 0): """ Initialize GPU-only manager. Args: num_blocks: Total number of blocks to manage block_size: Tokens per block (default 256) + max_seq_len: Max sequence length for contiguous cache (0 to disable) """ self._block_size = block_size self._num_blocks = num_blocks + self._max_seq_len = max_seq_len # Block metadata self.blocks: List[Block] = [Block(i) for i in range(num_blocks)] @@ -77,6 +80,11 @@ class GPUOnlyManager(KVCacheManager): self.num_kv_heads: int = 0 self.head_dim: int = 0 + # Contiguous cache for single-seq layer-wise prefill (set by allocate_cache) + self.contiguous_k_cache: Optional[Tensor] = None + self.contiguous_v_cache: Optional[Tensor] = None + self.contiguous_seq_len: int = 0 # Current sequence length in contiguous cache + @property def block_size(self) -> int: return self._block_size @@ -105,6 +113,23 @@ class GPUOnlyManager(KVCacheManager): dtype=dtype, device="cuda" ) + # Allocate contiguous cache for single-seq layer-wise prefill + # Only allocate if there's enough free memory (at least 2GB margin) + if self._max_seq_len > 0: + contiguous_cache_bytes = 2 * num_layers * self._max_seq_len * num_kv_heads * head_dim * dtype.itemsize + free_memory = torch.cuda.mem_get_info()[0] + + if free_memory > contiguous_cache_bytes + 2 * 1024**3: # 2GB margin + # Shape: [num_layers, max_seq_len, kv_heads, head_dim] + self.contiguous_k_cache = torch.empty( + num_layers, self._max_seq_len, num_kv_heads, head_dim, + dtype=dtype, device="cuda" + ) + self.contiguous_v_cache = torch.empty( + num_layers, self._max_seq_len, num_kv_heads, head_dim, + dtype=dtype, device="cuda" + ) + def get_layer_cache(self, layer_id: int) -> Tuple[Tensor, Tensor]: """Get K/V cache for a layer.""" assert self.kv_cache is not None, "Cache not allocated" diff --git a/task_plan.md b/task_plan.md index f74dc65..7c791c4 100644 --- a/task_plan.md +++ b/task_plan.md @@ -1,346 +1,412 @@ -# Task Plan: Integrate Sparsity into Layerwise Offload +# Task Plan: Fix GPU-only Mode Performance Issue ## Goal +Eliminate the `store_kvcache` scatter overhead in GPU-only mode by using **contiguous KV cache layout** (like offload mode), avoiding PagedAttention's blocked layout for single-sequence inference. -Extend MInference (prefill sparse) and Quest (decode sparse) to the layerwise offload execution path, with an extensible architecture for future sparsity methods. +## Problem Summary -## Key Insight +GPU-only mode with MInference is **slower** than CPU offload mode: -**现有的 sparse policy 已经实现,只是 layerwise offload 路径绕过了它!** +| Mode | Prefill Speed (32K tokens, Qwen3-4B) | +|------|--------------------------------------| +| GPU-only + MInference | 3383 tok/s | +| Offload + MInference | 5373 tok/s | -| 路径 | Attention 调用方式 | Sparse 支持 | -|------|-------------------|-------------| -| GPU-only | `attention.py` → `sparse_prefill_attention()` | YES | -| Layerwise offload | `model_runner.py` → `flash_attn_varlen_func()` | NO (直接调用) | +**Root cause**: PagedAttention's blocked layout requires expensive `index_copy_` scatter operations to convert contiguous K,V to blocked format. -## Policy Type Analysis +## Key Insight: Why Offload is Fast -**两类 sparse policy 的本质区别:** +Offload mode uses **contiguous layout** for KV cache: -| Policy | 影响 Attention 计算 | 影响 KV Load 策略 | `select_blocks()` 行为 | -|--------|-------------------|-----------------|----------------------| -| **MInference** | YES (`sparse_prefill_attention`) | NO | `return available_blocks` (全部) | -| **Quest** | NO | YES | 返回 Top-K subset | +```python +# OffloadEngine's CPU cache layout +k_cache_cpu: [num_layers, num_blocks, block_size, kv_heads, head_dim] -**MInference**: 只改变 attention 计算方式,不影响外部的 layer-wise load/offload 流程 -**Quest**: 选择性地只 load 部分 blocks,影响 H2D 传输 +# Store is simple contiguous slice assignment +self.k_cache_cpu[layer_id, block_id, :actual_size].copy_(k[start:end]) +``` -## Architecture Constraint +The K,V computed during prefill `[seq_len, kv_heads, head_dim]` matches the cache layout - no format conversion needed! -**所有 copy_ 操作必须封装在 OffloadEngine 中,model_runner.py 不能直接访问内部存储!** +## Solution: Contiguous Layout for GPU-only Mode + +For GPU-only single-sequence mode, use the **same contiguous layout as offload mode**, but on GPU: + +``` +Current GPU-only (PagedAttention): + Cache: [num_blocks, block_size, kv_heads, head_dim] (blocked) + Store: scatter via index_copy_ (SLOW) + +Proposed GPU-only (Contiguous): + Cache: [num_layers, max_seq_len, kv_heads, head_dim] (contiguous) + Store: slice assignment k_cache[layer_id, :seq_len] = k (FAST) +``` + +This mirrors offload mode's architecture but keeps everything on GPU - no cross-device transfer, no layout conversion. ## Phases +- [x] Phase 1: Add contiguous GPU KV cache in GPUOnlyManager (for single-seq mode) +- [x] Phase 2: Implement `run_gpu_only_prefill()` using contiguous cache +- [x] Phase 3: Implement decode path for contiguous cache +- [x] Phase 4: Test and validate performance -- [x] Phase 1: 添加 `requires_block_selection` 接口标志 -- [x] Phase 2: Refactor OffloadEngine - 封装 offload 操作,支持 sparse policy hooks -- [x] Phase 3: MInference prefill - 在 offload prefill 中调用 `sparse_prefill_attention()` -- [x] Phase 4: Quest decode - 根据 `requires_block_selection` 选择性 load blocks (infrastructure ready, full integration deferred) -- [x] Phase 5: Configuration 和 testing +## Results + +| Mode | 32K Prefill Speed | Notes | +|------|-------------------|-------| +| GPU-only (before) | ~3383 tok/s | PagedAttention scatter overhead | +| GPU-only contiguous (after) | **5293 tok/s** | 56% improvement | +| Offload mode | 5391 tok/s | Baseline comparison | + +**Test passed**: `test_needle.py --input-len 32768 --max-model-len 40960` - correct output retrieved. ## Detailed Design -### Phase 1: 添加 `requires_block_selection` 接口标志 +### Phase 1: Contiguous GPU KV Cache -**New attribute in SparsePolicy base class:** +**File**: `nanovllm/kvcache/gpu_manager.py` + +Add contiguous cache allocation for single-sequence mode: ```python -class SparsePolicy(ABC): - # Existing flags - supports_prefill: bool = True - supports_decode: bool = True +class GPUOnlyManager(KVCacheManager): + def __init__(self, num_blocks: int, block_size: int, max_seq_len: int = 0): + # ... existing code ... + self.max_seq_len = max_seq_len - # NEW: Whether this policy requires selective block loading - # If True: OffloadEngine will call select_blocks() before loading - # If False: OffloadEngine will load all blocks (select_blocks ignored) - requires_block_selection: bool = False -``` + # Contiguous cache for single-seq mode (allocated in allocate_cache) + self.contiguous_k_cache = None # [num_layers, max_seq_len, kv_heads, head_dim] + self.contiguous_v_cache = None -**Policy implementations:** - -```python -class MInferencePolicy(SparsePolicy): - supports_prefill = True - supports_decode = False - requires_block_selection = False # 不影响 load 策略 - - def select_blocks(self, available_blocks, ctx): - # 不会被调用(requires_block_selection=False) - return available_blocks - - -class QuestPolicy(SparsePolicy): - supports_prefill = False - supports_decode = True - requires_block_selection = True # 影响 load 策略 - - def select_blocks(self, available_blocks, ctx): - # 会被 OffloadEngine 调用 - return self._select_topk_blocks(...) - - -class FullAttentionPolicy(SparsePolicy): - supports_prefill = True - supports_decode = True - requires_block_selection = False # 加载所有 blocks -``` - -### Phase 2: Refactor OffloadEngine - -**OffloadEngine 根据 `requires_block_selection` 决定是否调用 `select_blocks()`:** - -```python -class OffloadEngine: - def __init__(self, ..., sparse_policy: "SparsePolicy" = None): - self.sparse_policy = sparse_policy - - def offload_layer_kv_sync( + def allocate_cache( self, - layer_id: int, - k: Tensor, - v: Tensor, - cpu_block_ids: List[int], - total_tokens: int, + num_layers: int, + num_kv_heads: int, + head_dim: int, + dtype: torch.dtype, ) -> None: - """ - Synchronously offload layer KV to CPU. - Calls sparse policy hooks internally. - """ - for i, cpu_block_id in enumerate(cpu_block_ids): - start = i * self.block_size - end = min(start + self.block_size, total_tokens) - actual_size = end - start + # Existing PagedAttention cache for multi-seq/decode + self.kv_cache = torch.empty( + 2, num_layers, self._num_blocks, self._block_size, + num_kv_heads, head_dim, + dtype=dtype, device="cuda" + ) - # Hook: notify sparse policy BEFORE offload (k still on GPU) - if self.sparse_policy is not None: - self.sparse_policy.on_prefill_offload( - cpu_block_id, layer_id, k[start:end], actual_size - ) - - # Synchronous copy to CPU (internal) - self.k_cache_cpu[layer_id, cpu_block_id, :actual_size].copy_(k[start:end]) - self.v_cache_cpu[layer_id, cpu_block_id, :actual_size].copy_(v[start:end]) - - def load_layer_kv_to_buffer_with_policy( - self, - buffer_idx: int, - layer_id: int, - cpu_block_ids: List[int], - valid_tokens_per_block: List[int], - query: Optional[Tensor] = None, - ) -> int: - """ - Load layer KV to buffer, optionally using sparse policy for block selection. - - Args: - buffer_idx: Ring buffer slot - layer_id: Layer index - cpu_block_ids: All available CPU block IDs - valid_tokens_per_block: Valid tokens per block - query: Query tensor (needed for block selection if requires_block_selection=True) - - Returns: - Total tokens loaded - """ - # Check if policy requires block selection - if (self.sparse_policy is not None and - self.sparse_policy.requires_block_selection and - query is not None): - # Build context - ctx = PolicyContext( - query_chunk_idx=0, - num_query_chunks=1, - layer_id=layer_id, - query=query, - is_prefill=False, - block_size=self.block_size, + # Contiguous cache for single-seq prefill (if max_seq_len specified) + if self.max_seq_len > 0: + self.contiguous_k_cache = torch.empty( + num_layers, self.max_seq_len, num_kv_heads, head_dim, + dtype=dtype, device="cuda" ) - # Select blocks - selected_blocks = self.sparse_policy.select_blocks(cpu_block_ids, ctx) - - # Build valid_tokens for selected blocks - block_to_valid = {bid: vt for bid, vt in zip(cpu_block_ids, valid_tokens_per_block)} - selected_valid = [block_to_valid[bid] for bid in selected_blocks] - - return self._load_blocks_to_buffer( - buffer_idx, layer_id, selected_blocks, selected_valid + self.contiguous_v_cache = torch.empty( + num_layers, self.max_seq_len, num_kv_heads, head_dim, + dtype=dtype, device="cuda" ) - else: - # Load all blocks (no selection) - return self._load_blocks_to_buffer( - buffer_idx, layer_id, cpu_block_ids, valid_tokens_per_block - ) - - def _load_blocks_to_buffer( - self, - buffer_idx: int, - layer_id: int, - block_ids: List[int], - valid_tokens: List[int], - ) -> int: - """Internal: load specified blocks to buffer.""" - stream = self.layer_load_streams[buffer_idx] - - with torch.cuda.stream(stream): - stream.wait_event(self.buffer_compute_done_events[buffer_idx]) - - offset = 0 - for cpu_block_id, vt in zip(block_ids, valid_tokens): - self.layer_k_cache[buffer_idx, offset:offset+vt].copy_( - self.k_cache_cpu[layer_id, cpu_block_id, :vt], - non_blocking=True - ) - self.layer_v_cache[buffer_idx, offset:offset+vt].copy_( - self.v_cache_cpu[layer_id, cpu_block_id, :vt], - non_blocking=True - ) - offset += vt - - self.buffer_load_events[buffer_idx].record(stream) - - return offset ``` -### Phase 3: MInference Prefill Integration +### Phase 2: Layer-wise GPU-only Prefill -**MInference 只影响 attention 计算,不影响 load/offload:** +**File**: `nanovllm/engine/model_runner.py` + +Following offload pattern exactly - store K,V per-layer to contiguous cache: ```python -def run_layerwise_offload_prefill(self, seqs): - ... +@torch.inference_mode() +def run_gpu_only_prefill(self, seqs: list[Sequence]) -> list[int]: + """ + GPU-only prefill with contiguous KV cache layout. + + Mirrors run_layerwise_offload_prefill() but stores to GPU instead of CPU. + No scatter operations - just contiguous slice assignment. + """ + assert len(seqs) == 1, "GPU-only layer-wise prefill only supports single sequence" + seq = seqs[0] + + num_layers = len(self.model.model.layers) + total_tokens = len(seq) + + # Get contiguous GPU cache + k_cache = self.kvcache_manager.contiguous_k_cache + v_cache = self.kvcache_manager.contiguous_v_cache + + # Prepare inputs + input_ids = torch.tensor(seq[:], dtype=torch.int64, device="cuda") + positions = torch.arange(total_tokens, dtype=torch.int64, device="cuda") + + from flash_attn.flash_attn_interface import flash_attn_varlen_func + cu_seqlens = torch.tensor([0, total_tokens], dtype=torch.int32, device="cuda") + + # Embedding + hidden_states = self.model.model.embed_tokens(input_ids) + residual = None + + # Layer-by-layer processing (same as offload prefill) for layer_id in range(num_layers): - # QKV projection + RoPE + layer = self.model.model.layers[layer_id] + + # Input LayerNorm + if residual is None: + hidden_ln, residual = layer.input_layernorm(hidden_states), hidden_states + else: + hidden_ln, residual = layer.input_layernorm(hidden_states, residual) + + # QKV projection + qkv = layer.self_attn.qkv_proj(hidden_ln) + q, k, v = qkv.split([ + layer.self_attn.q_size, + layer.self_attn.kv_size, + layer.self_attn.kv_size + ], dim=-1) + + q = q.view(total_tokens, layer.self_attn.num_heads, layer.self_attn.head_dim) + k = k.view(total_tokens, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + v = v.view(total_tokens, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + + # Q/K norms (Qwen3 specific) + if not layer.self_attn.qkv_bias: + q = layer.self_attn.q_norm(q.reshape(-1, layer.self_attn.head_dim)) + q = q.view(total_tokens, layer.self_attn.num_heads, layer.self_attn.head_dim) + k = layer.self_attn.k_norm(k.reshape(-1, layer.self_attn.head_dim)) + k = k.view(total_tokens, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + + # RoPE q, k = layer.self_attn.rotary_emb(positions, q, k) - # Sparse or Full attention + # Store K,V to contiguous GPU cache (same layout - no conversion!) + # This is just slice assignment, not scatter + k_cache[layer_id, :total_tokens] = k + v_cache[layer_id, :total_tokens] = v + + # Sparse or Full attention (uses k, v directly) if self.sparse_prefill_policy is not None: - # MInference: only changes attention computation attn_output = self.sparse_prefill_policy.sparse_prefill_attention( q, k, v, layer_id ) else: - attn_output = flash_attn_varlen_func(q, k, v, ...) + attn_output = flash_attn_varlen_func( + q, k, v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=total_tokens, + max_seqlen_k=total_tokens, + softmax_scale=layer.self_attn.attn.scale, + causal=True, + ) - # MLP - ... + # O projection + attn_output = attn_output.view(total_tokens, -1) + hidden_states = layer.self_attn.o_proj(attn_output) - # Offload ALL KV (MInference doesn't affect this) - offload_engine.offload_layer_kv_sync(layer_id, k, v, cpu_block_ids, total_tokens) + # Post-attention LayerNorm + MLP + hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) + hidden_states = layer.mlp(hidden_states) + + # Final norm + hidden_states, _ = self.model.model.norm(hidden_states, residual) + + # Compute logits + logits = self.model.compute_logits(hidden_states[-1:]) + + # Record prefill length for decode + self.kvcache_manager.contiguous_seq_len = total_tokens + + # Sample + temperatures = self.prepare_sample(seqs) if self.rank == 0 else None + token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None + + return token_ids ``` -### Phase 4: Quest Decode Integration +### Phase 3: Decode with Contiguous Cache -**Quest 影响 block load 策略:** +**File**: `nanovllm/engine/model_runner.py` ```python -def run_layerwise_offload_decode(self, seqs): - ... - # Preload first N layers (no query available, full load) - for i in range(num_preload): - loaded_tokens[i] = offload_engine.load_layer_kv_to_buffer_with_policy( - i, i, cpu_block_table, valid_tokens_per_block, query=None - ) +@torch.inference_mode() +def run_gpu_only_decode(self, seqs: list[Sequence]) -> list[int]: + """ + Decode using contiguous GPU KV cache. + + Similar to offload decode but simpler - all KV already on GPU. + """ + assert len(seqs) == 1 + seq = seqs[0] + + num_layers = len(self.model.model.layers) + k_cache = self.kvcache_manager.contiguous_k_cache + v_cache = self.kvcache_manager.contiguous_v_cache + context_len = self.kvcache_manager.contiguous_seq_len + + # Prepare inputs + input_ids = torch.tensor([seq.last_token], dtype=torch.int64, device="cuda") + positions = torch.tensor([len(seq) - 1], dtype=torch.int64, device="cuda") + + from flash_attn.flash_attn_interface import flash_attn_varlen_func + cu_seqlens_q = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + + # Embedding + hidden_states = self.model.model.embed_tokens(input_ids) + residual = None for layer_id in range(num_layers): - current_buffer = layer_id % num_buffers + layer = self.model.model.layers[layer_id] - # Wait for buffer load - offload_engine.wait_buffer_load(current_buffer) + # Input LayerNorm + if residual is None: + hidden_ln, residual = layer.input_layernorm(hidden_states), hidden_states + else: + hidden_ln, residual = layer.input_layernorm(hidden_states, residual) # QKV projection - q, k_new, v_new = ... + qkv = layer.self_attn.qkv_proj(hidden_ln) + q, k_new, v_new = qkv.split([ + layer.self_attn.q_size, + layer.self_attn.kv_size, + layer.self_attn.kv_size + ], dim=-1) - # Get loaded KV - k_prefill, v_prefill = offload_engine.get_buffer_kv( - current_buffer, loaded_tokens[current_buffer] - ) + q = q.view(1, layer.self_attn.num_heads, layer.self_attn.head_dim) + k_new = k_new.view(1, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + v_new = v_new.view(1, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + + # Q/K norms + if not layer.self_attn.qkv_bias: + q = layer.self_attn.q_norm(q.reshape(-1, layer.self_attn.head_dim)) + q = q.view(1, layer.self_attn.num_heads, layer.self_attn.head_dim) + k_new = layer.self_attn.k_norm(k_new.reshape(-1, layer.self_attn.head_dim)) + k_new = k_new.view(1, layer.self_attn.num_kv_heads, layer.self_attn.head_dim) + + # RoPE + q, k_new = layer.self_attn.rotary_emb(positions, q, k_new) + + # Get cached K,V and append new token + k_cached = k_cache[layer_id, :context_len] + v_cached = v_cache[layer_id, :context_len] + + # Store new K,V to cache + k_cache[layer_id, context_len] = k_new.squeeze(0) + v_cache[layer_id, context_len] = v_new.squeeze(0) + + # Full K,V for attention + k_full = k_cache[layer_id, :context_len + 1] + v_full = v_cache[layer_id, :context_len + 1] # Attention - ... + cu_seqlens_k = torch.tensor([0, context_len + 1], dtype=torch.int32, device="cuda") + attn_output = flash_attn_varlen_func( + q, k_full, v_full, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=1, + max_seqlen_k=context_len + 1, + softmax_scale=layer.self_attn.attn.scale, + causal=False, # Single query, no causal needed + ) - # Mark buffer done - offload_engine.record_buffer_compute_done(current_buffer) + # O projection + attn_output = attn_output.view(1, -1) + hidden_states = layer.self_attn.o_proj(attn_output) - # Load next layer (Quest: selective load if requires_block_selection=True) - next_layer = layer_id + num_buffers - if next_layer < num_layers: - loaded_tokens[current_buffer] = offload_engine.load_layer_kv_to_buffer_with_policy( - current_buffer, next_layer, cpu_block_table, valid_tokens_per_block, - query=q # Pass query for block selection - ) + # Post-attention LayerNorm + MLP + hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) + hidden_states = layer.mlp(hidden_states) + + # Update context length + self.kvcache_manager.contiguous_seq_len = context_len + 1 + + # Final norm + hidden_states, _ = self.model.model.norm(hidden_states, residual) + + # Compute logits + logits = self.model.compute_logits(hidden_states) + + # Sample + temperatures = self.prepare_sample(seqs) if self.rank == 0 else None + token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None + + return token_ids ``` -### Phase 5: Configuration +### Phase 4: Decision Logic ```python -@dataclass -class Config: - # Separate policies for prefill and decode - sparse_prefill_policy: SparsePolicyType = SparsePolicyType.FULL # MINFERENCE - sparse_decode_policy: SparsePolicyType = SparsePolicyType.FULL # QUEST +def _should_use_contiguous_gpu_mode(self, seqs: list[Sequence], is_prefill: bool) -> bool: + """Check if contiguous GPU mode should be used.""" + # Must have contiguous cache allocated + if not hasattr(self.kvcache_manager, 'contiguous_k_cache'): + return False + if self.kvcache_manager.contiguous_k_cache is None: + return False + + # Must NOT be offload mode + if hasattr(self.kvcache_manager, 'offload_engine'): + return False + + # Single sequence only + if len(seqs) != 1: + return False + + # For prefill: has blocks (not warmup) + if is_prefill and not seqs[0].block_table: + return False + + return True + + +def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]: + # Check offload mode (existing) + if hasattr(self, 'kvcache_manager') and hasattr(self.kvcache_manager, 'offload_engine'): + ... + + # Check contiguous GPU mode + if self._should_use_contiguous_gpu_mode(seqs, is_prefill): + if is_prefill: + return self.run_gpu_only_prefill(seqs) + else: + return self.run_gpu_only_decode(seqs) + + # Standard PagedAttention path + ... ``` -## File Changes Summary +## Architecture Comparison + +| Aspect | Offload Mode | GPU-only (Proposed) | GPU-only (Current) | +|--------|--------------|---------------------|-------------------| +| Cache location | CPU (contiguous) | GPU (contiguous) | GPU (PagedAttention) | +| Cache layout | `[layers, blocks, block_size, heads, dim]` | `[layers, max_seq_len, heads, dim]` | `[blocks, block_size, heads, dim]` | +| Prefill store | Contiguous slice copy | **Slice assignment (no copy!)** | Scatter (index_copy_) | +| Decode read | H2D ring buffer | Direct GPU access | PagedAttention | + +## Key Points + +1. **No explicit copy_ needed**: Slice assignment `cache[layer, :len] = k` is direct memory write +2. **Same layout as computed K,V**: No format conversion required +3. **Mirrors offload architecture**: Same layer-wise processing pattern +4. **GPU advantage**: No cross-device transfer, faster than offload + +## Memory Usage + +Contiguous GPU cache: `2 * num_layers * max_seq_len * kv_heads * head_dim * dtype_size` + +For Qwen3-4B with 32K max_seq_len: +- `2 * 28 * 32768 * 8 * 128 * 2 = 3.5GB` + +Same as offload mode's CPU cache, but on GPU. + +## Files to Modify | File | Changes | |------|---------| -| `nanovllm/kvcache/sparse/policy.py` | Add `requires_block_selection` attribute | -| `nanovllm/kvcache/sparse/minference.py` | Set `requires_block_selection = False` | -| `nanovllm/kvcache/sparse/quest.py` | Set `requires_block_selection = True` | -| `nanovllm/kvcache/sparse/full_policy.py` | Set `requires_block_selection = False` | -| `nanovllm/kvcache/offload_engine.py` | Add `offload_layer_kv_sync()`, `load_layer_kv_to_buffer_with_policy()` | -| `nanovllm/engine/model_runner.py` | Use encapsulated methods, integrate sparse policies | +| `nanovllm/kvcache/gpu_manager.py` | Add contiguous cache allocation | +| `nanovllm/engine/model_runner.py` | Add `run_gpu_only_prefill()`, `run_gpu_only_decode()`, modify `run()` | -## Key Design Principles +## Expected Performance -1. **Encapsulation**: All copy_ operations in OffloadEngine -2. **Interface Flag**: `requires_block_selection` declares if policy affects load strategy -3. **Separation of Concerns**: - - MInference: only `sparse_prefill_attention()` (compute-level) - - Quest: `select_blocks()` + hooks (load-level) -4. **Hooks inside engine**: Sparse policy hooks called within OffloadEngine methods - -## Decisions Made - -- [x] 添加 `requires_block_selection` 接口标志区分两类 policy -- [x] 所有 copy_ 封装在 OffloadEngine 中 -- [x] Sparse policy hooks 在 OffloadEngine 内部调用 -- [x] Decode preload 使用全量加载(Q 不可用) +| Metric | Before | After | Improvement | +|--------|--------|-------|-------------| +| GPU-only prefill (32K) | 3383 tok/s | ~5400+ tok/s | ~60%+ | +| Decode | Baseline | Similar | ~0% | ## Status - -**COMPLETE** - All phases implemented and tested successfully. - -### Test Results (Qwen3-4B-Instruct-2507) - -验证 offload + MInference 输出与 GPU-only + MInference 完全一致: - -``` -# GPU-only + MInference -test_needle.py --model Qwen3-4B --input-len 32768 --enable-minference -- Prefill: 3383 tok/s -- Output tokens: [22, 19, 24, 17, 151645] = "7492<|im_end|>" -- Result: PASSED - -# Offload + MInference -test_needle.py --model Qwen3-4B --input-len 32768 --enable-offload --enable-minference -- Prefill: 5373 tok/s (faster due to layer-wise processing) -- Output tokens: [22, 19, 24, 17, 151645] = "7492<|im_end|>" -- Result: PASSED - -两种配置输出完全一致! -``` - -Note: Qwen3-0.6B 在 offload 模式下有已知 bug(模型太小,长序列不稳定),不是本次修改引入。 - -## Performance Discovery - -**意外发现**: Offload 模式比 GPU-only 模式更快! - -| Mode | Prefill Speed | -|------|---------------| -| GPU-only + MInference | 3383 tok/s | -| Offload + MInference | 5373 tok/s | - -**根本原因**: GPU-only 模式的 `store_kvcache()` 使用 PagedAttention 的 scatter 操作 (`index_copy_`),而 offload 模式使用 contiguous copy。 - -详细分析和优化建议见: [`docs/gpu_only_performance_issue.md`](docs/gpu_only_performance_issue.md) +**Currently in Phase 1** - Ready to implement contiguous GPU cache