[WIP] Before refactor the nanovllm sparse policy.

.claude/rules/planning-with-files.md

@@ -23,7 +23,7 @@ rm -f task_plan_*.md findings_*.md progress_*.md
 
 ```bash
 # Step 1: 清理旧计划文件
-rm -f task_plan.md findings.md progress.md task_plan_*.md findings_*.md progress_*.md
+rm -f task_plan.md findings.md progress.md
 
 # Step 2: 启动 planning-with-files 技能
 # 在 Claude 中调用 /planning-with-files 或 Skill tool

findings.md

@@ -1,160 +0,0 @@
-# Findings: Multi-Model Support Analysis
-
-## Current Architecture Analysis
-
-### Model Loading Flow
-```
-LLM(model_path)
-  → LLMEngine.__init__()
-    → Config.__post_init__()
-      → hf_config = AutoConfig.from_pretrained(model)
-    → ModelRunner.__init__()
-      → model = Qwen3ForCausalLM(hf_config)  ← HARDCODED
-      → load_model(model, config.model)
-```
-
-### Key Files
-| File | Purpose |
-|------|---------|
-| `nanovllm/engine/model_runner.py` | 模型加载和运行 |
-| `nanovllm/models/qwen3.py` | Qwen3 模型定义 |
-| `nanovllm/utils/loader.py` | safetensors 权重加载 |
-| `nanovllm/layers/rotary_embedding.py` | RoPE 实现 |
-
----
-
-## Llama 3.1 Config Analysis
-
-```json
-{
-  "architectures": ["LlamaForCausalLM"],
-  "model_type": "llama",
-  "attention_bias": false,
-  "mlp_bias": false,
-  "head_dim": 128,
-  "hidden_size": 4096,
-  "intermediate_size": 14336,
-  "num_attention_heads": 32,
-  "num_hidden_layers": 32,
-  "num_key_value_heads": 8,
-  "hidden_act": "silu",
-  "rms_norm_eps": 1e-05,
-  "rope_theta": 500000.0,
-  "rope_scaling": {
-    "factor": 8.0,
-    "high_freq_factor": 4.0,
-    "low_freq_factor": 1.0,
-    "original_max_position_embeddings": 8192,
-    "rope_type": "llama3"
-  },
-  "max_position_embeddings": 131072,
-  "tie_word_embeddings": false,
-  "vocab_size": 128256
-}
-```
-
-### Llama 3 RoPE Scaling
-Llama 3 使用特殊的 RoPE scaling 策略 (`rope_type: "llama3"`)：
-- 低频分量保持不变（对应短距离依赖）
-- 高频分量线性插值（对应长距离依赖）
-- 参数: `factor`, `low_freq_factor`, `high_freq_factor`, `original_max_position_embeddings`
-
-参考实现 (transformers):
-```python
-def _compute_llama3_parameters(config, device, inv_freq):
-    factor = config.factor
-    low_freq_factor = config.low_freq_factor
-    high_freq_factor = config.high_freq_factor
-    old_context_len = config.original_max_position_embeddings
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    inv_freq_llama = torch.where(
-        wavelen > low_freq_wavelen,
-        inv_freq / factor,
-        inv_freq
-    )
-    smooth_factor = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (1 - smooth_factor) * inv_freq_llama + smooth_factor * inv_freq
-    is_medium_freq = (wavelen >= high_freq_wavelen) & (wavelen <= low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
-    return inv_freq_llama
-```
-
----
-
-## Weight Mapping Analysis
-
-### Qwen3 packed_modules_mapping
-```python
-packed_modules_mapping = {
-    "q_proj": ("qkv_proj", "q"),
-    "k_proj": ("qkv_proj", "k"),
-    "v_proj": ("qkv_proj", "v"),
-    "gate_proj": ("gate_up_proj", 0),
-    "up_proj": ("gate_up_proj", 1),
-}
-```
-
-### Llama Weight Names (from safetensors)
-预期 Llama 权重命名与 Qwen3 类似：
-- `model.layers.{i}.self_attn.q_proj.weight`
-- `model.layers.{i}.self_attn.k_proj.weight`
-- `model.layers.{i}.self_attn.v_proj.weight`
-- `model.layers.{i}.self_attn.o_proj.weight`
-- `model.layers.{i}.mlp.gate_proj.weight`
-- `model.layers.{i}.mlp.up_proj.weight`
-- `model.layers.{i}.mlp.down_proj.weight`
-- `model.layers.{i}.input_layernorm.weight`
-- `model.layers.{i}.post_attention_layernorm.weight`
-
-**结论**: Llama 的 `packed_modules_mapping` 与 Qwen3 相同，可以复用。
-
----
-
-## Shared Components (Can Reuse)
-
-| Component | File | Notes |
-|-----------|------|-------|
-| `RMSNorm` | `layers/layernorm.py` | 通用 |
-| `SiluAndMul` | `layers/activation.py` | 通用 |
-| `Attention` | `layers/attention.py` | FlashAttention wrapper |
-| `QKVParallelLinear` | `layers/linear.py` | 支持 bias=False |
-| `RowParallelLinear` | `layers/linear.py` | 通用 |
-| `MergedColumnParallelLinear` | `layers/linear.py` | 通用 |
-| `VocabParallelEmbedding` | `layers/embed_head.py` | 通用 |
-| `ParallelLMHead` | `layers/embed_head.py` | 通用 |
-| `load_model` | `utils/loader.py` | 通用 |
-
----
-
-## Llama vs Qwen3 Implementation Diff
-
-### Attention
-| Feature | Qwen3Attention | LlamaAttention |
-|---------|----------------|----------------|
-| QKV bias | 可配置 (attention_bias) | 始终 False |
-| q_norm | 有 (when bias=False) | 无 |
-| k_norm | 有 (when bias=False) | 无 |
-| RoPE | Standard | Llama3 scaled |
-
-### MLP
-| Feature | Qwen3MLP | LlamaMLP |
-|---------|----------|----------|
-| gate/up bias | False | False |
-| down bias | False | False |
-| hidden_act | silu | silu |
-
-**结论**: Llama MLP 与 Qwen3 MLP 几乎相同，可以直接复用或简化。
-
----
-
-## Risk Assessment
-
-| Risk | Impact | Mitigation |
-|------|--------|------------|
-| RoPE 实现错误 | 高 - 导致错误输出 | 参考 transformers 实现，单元测试 |
-| 权重映射错误 | 高 - 模型无法加载 | 检查 safetensors 键名 |
-| 注册表循环导入 | 中 - 启动失败 | 延迟导入 |

nanovllm/kvcache/sparse/__init__.py

@@ -61,8 +61,6 @@ def create_sparse_policy(policy_type: SparsePolicyType, **kwargs) -> SparsePolic
             block_size=kwargs.get("block_size", 128),
             samples_per_chunk=kwargs.get("samples_per_chunk", 128),
             threshold=kwargs.get("threshold", 0.9),
-            use_triton=kwargs.get("use_triton", True),
-            stride=kwargs.get("stride", 8),
         )
 
     else:

nanovllm/kvcache/sparse/full_policy.py

@@ -5,8 +5,11 @@ This serves as a baseline and default policy when sparse
 attention is not needed.
 """
 
-from typing import List
+import torch
+from typing import List, Optional
+
 from .policy import SparsePolicy, PolicyContext
+from nanovllm.utils.context import get_context
 
 
 class FullAttentionPolicy(SparsePolicy):
@@ -34,5 +37,129 @@ class FullAttentionPolicy(SparsePolicy):
         """Return all blocks - no sparsity."""
         return available_blocks
 
+    def compute_prefill_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_id: int,
+        softmax_scale: float,
+        offload_engine,
+        current_chunk_idx: int,
+        seq,
+    ) -> torch.Tensor:
+        """
+        Compute full attention for chunked prefill.
+
+        This method handles the complete chunked prefill flow:
+        1. Load historical blocks from CPU
+        2. Compute attention to historical chunks
+        3. Compute attention to current chunk
+        4. Merge all results
+
+        Args:
+            q: Query tensor [seq_len, num_heads, head_dim]
+            k: Key tensor [seq_len, num_kv_heads, head_dim] (unused, from prefill buffer)
+            v: Value tensor [seq_len, num_kv_heads, head_dim] (unused, from prefill buffer)
+            layer_id: Current layer index
+            softmax_scale: Softmax scaling factor
+            offload_engine: OffloadEngine for loading blocks
+            current_chunk_idx: Current chunk index
+            seq: ChunkedSequence
+
+        Returns:
+            Attention output [seq_len, num_heads, head_dim]
+        """
+        from nanovllm.kvcache.chunked_attention import flash_attn_with_lse, merge_attention_outputs
+
+        q_batched = q.unsqueeze(0)  # [1, seq_len, num_heads, head_dim]
+        num_tokens = q.shape[0]
+        o_acc = None
+        lse_acc = None
+        compute_stream = offload_engine.compute_stream
+
+        # Step 1: Get and load historical blocks
+        cpu_block_table = seq.kvcache_manager.get_prefilled_cpu_blocks(seq)
+
+        if cpu_block_table:
+            load_slots = list(range(offload_engine.num_ring_slots))
+            num_blocks = len(cpu_block_table)
+
+            if len(load_slots) == 1:
+                # Only 1 slot - 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):
+                        prev_k, prev_v = offload_engine.get_kv_for_slot(slot)
+                        prev_o, prev_lse = flash_attn_with_lse(
+                            q_batched, prev_k, prev_v,
+                            softmax_scale=softmax_scale,
+                            causal=False,
+                        )
+                        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)
+                        offload_engine.record_slot_compute_done(slot)
+            else:
+                # Multiple slots - use pipeline
+                num_slots = len(load_slots)
+                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])
+
+                for block_idx in range(num_blocks):
+                    current_slot = load_slots[block_idx % num_slots]
+                    cpu_block_id = cpu_block_table[block_idx]
+
+                    offload_engine.wait_slot_layer(current_slot)
+
+                    with torch.cuda.stream(compute_stream):
+                        prev_k, prev_v = offload_engine.get_kv_for_slot(current_slot)
+                        prev_o, prev_lse = flash_attn_with_lse(
+                            q_batched, prev_k, prev_v,
+                            softmax_scale=softmax_scale,
+                            causal=False,
+                        )
+                        offload_engine.record_slot_compute_done(current_slot)
+
+                        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)
+
+                    # Issue next transfer
+                    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)
+        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:
+                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 __repr__(self) -> str:
         return "FullAttentionPolicy()"

nanovllm/kvcache/sparse/xattn_bsa.py

@@ -1,15 +1,13 @@
 """
 XAttention Block Sparse Attention (BSA) Policy for nano-vllm.
 
-This module implements XAttention-inspired block sparse attention for chunked prefill,
+This module implements XAttention-inspired block sparse attention for chunked prefill.
-using block-level estimation to select important KV blocks for computation.
+Current implementation loads all historical blocks (FULL strategy).
 
-Reference: COMPASS/compass/src/Xattention.py
+Sparse selection to be implemented in next phase.
 """
 
-import math
 import torch
-import torch.nn.functional as F
 from typing import List, Optional, Tuple
 
 from nanovllm.kvcache.sparse.policy import SparsePolicy, PolicyContext
@@ -23,18 +21,11 @@ class XAttentionBSAPolicy(SparsePolicy):
     This policy uses block-level estimation to determine which KV blocks
     are important for the current chunk's queries, enabling sparse computation.
 
-    Key features:
+    Note: Current implementation loads all historical chunks (FULL strategy).
-    - Double-loading design: estimate phase loads samples, compute phase loads selected blocks
+    Sparse selection to be implemented in next phase.
-    - 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_prefill = False  # Uses standard select_blocks interface
     supports_decode = False  # BSA is prefill-only
     requires_block_selection = False  # Selection happens at chunk level, not block level
 
@@ -43,8 +34,6 @@ class XAttentionBSAPolicy(SparsePolicy):
         block_size: int = 128,
         samples_per_chunk: int = 128,
         threshold: float = 0.9,
-        use_triton: bool = True,
-        stride: int = 8,
     ):
         """
         Initialize XAttention BSA policy.
@@ -53,457 +42,29 @@ class XAttentionBSAPolicy(SparsePolicy):
             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).
+        Select blocks to load from CPU.
 
-        For prefill, BSA handles chunk-level selection internally.
+        Current implementation returns all blocks (FULL strategy).
+        Sparse selection to be implemented in next phase.
+
+        Args:
+            available_blocks: List of all available CPU block IDs
+            ctx: Policy context with query info, chunk index, etc.
+
+        Returns:
+            List of selected block IDs to load
         """
-        # For prefill, we return all blocks - selection happens in sparse_prefill_attention
+        # Current: Return all blocks (FULL strategy)
+        # TODO: Implement sparse selection based on query attention estimation
         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

nanovllm/layers/attention.py

@@ -210,22 +210,7 @@ class Attention(nn.Module):
             # Apply sparse policy if enabled
             sparse_policy = kvcache_manager.sparse_policy
 
-            # === XAttention BSA: Policy handles entire sparse prefill ===
+            # === All sparse policies use select_blocks interface ===
-            # Check if policy has sparse_prefill_attention method (XAttention BSA)
-            if (sparse_policy is not None and
-                hasattr(sparse_policy, 'sparse_prefill_attention') and
-                getattr(sparse_policy, 'supports_prefill', False)):
-                # Use policy's sparse_prefill_attention method
-                # Pass softmax_scale from attention layer
-                # IMPORTANT: Don't return early - we still need to do KV offload below!
-                o = sparse_policy.sparse_prefill_attention(q, k, v, self.layer_id, self.scale)
-                # Convert back to batched format for consistency with standard flow
-                o_acc = o.unsqueeze(0)  # [seq_len, heads, dim] -> [1, seq_len, heads, dim]
-                lse_acc = None  # sparse_prefill_attention returns final output, not intermediate LSE
-                # Skip standard flow processing since we already computed attention
-                cpu_block_table = None  # Signal to skip historical chunk processing
-
-            # === Standard sparse policy (Quest, etc.) ===
             if cpu_block_table and sparse_policy is not None:
                 num_chunks = getattr(context, 'num_chunks', current_chunk_idx + 1)
                 policy_ctx = PolicyContext(
@@ -262,8 +247,7 @@ class Attention(nn.Module):
         compute_stream = offload_engine.compute_stream if offload_engine is not None else None
 
         # Compute attention against current chunk's KV from prefill buffer (with causal mask)
-        # Skip this if XAttention BSA already computed full attention (o_acc is set, lse_acc is None)
+        needs_current_chunk_attention = True
-        needs_current_chunk_attention = (lse_acc is not None or o_acc is None)
 
         if needs_current_chunk_attention:
             if compute_stream is not None:
@@ -294,12 +278,10 @@ class Attention(nn.Module):
 
         # Merge with accumulated (all on compute_stream for consistency)
         if o_acc is None:
-            # No accumulated attention (standard flow or XAttention BSA with no historical chunks)
+            # No accumulated attention (no historical chunks processed)
-            final_o = current_o if needs_current_chunk_attention else o_acc
+            final_o = current_o
         else:
-            # Has accumulated attention (XAttention BSA with historical chunks)
+            # Has accumulated attention (historical chunks processed)
-            if needs_current_chunk_attention:
-                # Need to merge historical (from XAttention BSA) with current chunk
             if compute_stream is not None:
                 with torch.cuda.stream(compute_stream):
                     torch.cuda.nvtx.range_push(f"MergeAttn: L{self.layer_id}")
@@ -309,9 +291,6 @@ class Attention(nn.Module):
                 torch.cuda.nvtx.range_push(f"MergeAttn: L{self.layer_id}")
                 final_o, _ = merge_attention_outputs(o_acc, lse_acc, current_o, current_lse)
                 torch.cuda.nvtx.range_pop()
-            else:
-                # XAttention BSA already computed everything
-                final_o = o_acc
 
         torch.cuda.nvtx.range_pop()  # ChunkedPrefill
 

progress.md

@@ -1,76 +0,0 @@
-# Progress Log: Multi-Model Support
-
-## Session: 2026-01-10
-
-### Initial Analysis Complete
-
-**Time**: Session start
-
-**Actions:**
-1. Read `nanovllm/engine/model_runner.py` - 确认硬编码位置 (line 35)
-2. Read `nanovllm/models/qwen3.py` - 理解 Qwen3 模型结构
-3. Read `nanovllm/utils/loader.py` - 理解权重加载机制
-4. Read `nanovllm/layers/rotary_embedding.py` - 发现 RoPE scaling 限制
-5. Read `/home/zijie/models/Llama-3.1-8B-Instruct/config.json` - 理解 Llama 配置
-
-**Key Findings:**
-- 模型加载在 `model_runner.py:35` 硬编码为 Qwen3
-- RoPE 目前不支持 scaling (`assert rope_scaling is None`)
-- Llama 3.1 需要 "llama3" 类型的 RoPE scaling
-- Llama 无 q_norm/k_norm，无 attention bias
-
-**Created:**
-- `task_plan.md` - 6 阶段实施计划
-- `findings.md` - 技术分析和发现
-
----
-
-### Phase Status
-
-| Phase | Status | Notes |
-|-------|--------|-------|
-| 1. Model Registry | **COMPLETED** | `registry.py`, `__init__.py` |
-| 2. Llama3 RoPE | **COMPLETED** | `rotary_embedding.py` |
-| 3. Llama Model | **COMPLETED** | `llama.py` |
-| 4. ModelRunner | **COMPLETED** | Dynamic loading |
-| 5. Qwen3 Register | **COMPLETED** | `@register_model` decorator |
-| 6. Testing | **COMPLETED** | Both Llama & Qwen3 pass |
-
----
-
-## Test Results
-
-### Llama 3.1-8B-Instruct (32K needle, GPU 0, offload)
-```
-Input: 32768 tokens
-Expected: 7492
-Output: 7492
-Status: PASSED
-Prefill: 1644 tok/s
-```
-
-### Qwen3-4B (8K needle, GPU 1, offload) - Regression Test
-```
-Input: 8192 tokens
-Expected: 7492
-Output: 7492
-Status: PASSED
-Prefill: 3295 tok/s
-```
-
----
-
-## Files Modified This Session
-
-| File | Action | Description |
-|------|--------|-------------|
-| `nanovllm/models/registry.py` | created | Model registry with `@register_model` decorator |
-| `nanovllm/models/__init__.py` | created | Export registry functions, import models |
-| `nanovllm/models/llama.py` | created | Llama model implementation |
-| `nanovllm/models/qwen3.py` | modified | Added `@register_model` decorator |
-| `nanovllm/layers/rotary_embedding.py` | modified | Added Llama3 RoPE scaling |
-| `nanovllm/engine/model_runner.py` | modified | Dynamic model loading via registry |
-| `.claude/rules/gpu-testing.md` | created | GPU testing rules |
-| `task_plan.md` | created | Implementation plan |
-| `findings.md` | created | Technical findings |
-| `progress.md` | created | Progress tracking |

task_plan.md

@@ -1,144 +1,353 @@
-# Task Plan: Multi-Model Support for nanovllm
+# Task Plan: Sparse Policy 架构重构 v3
 
 ## Goal
-扩展 nanovllm 框架以支持多种模型（当前只支持 Qwen3），特别是添加 Llama-3.1-8B-Instruct 支持，并建立可扩展的模型添加范式。
 
-## Current State Analysis
+将 chunked prefill 的 attention 计算逻辑完全从 `attention.py` 移到 `SparsePolicy` 内部。attention.py 只负责调用 policy，不包含任何计算逻辑。
 
-### 硬编码问题位置
+## 核心设计原则（强制要求）
-- `nanovllm/engine/model_runner.py:35`: 直接实例化 `Qwen3ForCausalLM(hf_config)`
-- `nanovllm/engine/model_runner.py:9`: 硬编码导入 `from nanovllm.models.qwen3 import Qwen3ForCausalLM`
 
-### Qwen3 vs Llama 3.1 架构差异
+1. **Policy 内部完成所有计算**：包括 attention 计算和结果合并
+2. **select_blocks 传入 offload_engine**：policy 通过 offload_engine 加载 blocks
+3. **强制实现计算函数**：所有 policy 必须实现 `compute_block_attention` 和 `merge_attention_outputs`
+4. **chunked_prefill 强制 policy 存在**：没有 policy 则报错
+5. **外部默认 FULL policy**：model_runner.py 默认创建 FullPolicy
+6. **attention.py 零计算逻辑**：_chunked_prefill_attention 只调用 policy，不直接调用 flashattn 或 merge
 
-| Feature | Qwen3 | Llama 3.1 |
+## 目标架构
-|---------|-------|-----------|
-| Config Class | Qwen3Config | LlamaConfig |
-| attention_bias | True (可配置) | False |
-| q_norm/k_norm | 有 (when bias=False) | 无 |
-| mlp_bias | N/A | False |
-| RoPE Scaling | None (目前) | llama3 类型 |
-| RoPE theta | 1000000 | 500000 |
-| hidden_act | silu | silu |
-| tie_word_embeddings | True | False |
 
-### 关键限制
+```
-- `rotary_embedding.py:59`: `assert rope_scaling is None` - 不支持 RoPE scaling
+model_runner.py:
+  默认创建 FullPolicy（如果没有指定 sparse policy）
 
----
+attention.py (_chunked_prefill_attention):
+  检查 sparse_policy 是否存在
+    ↓
+  调用 sparse_policy.compute_prefill_attention(q, k, v, ...)
+    ↓
+  返回最终输出（不包含任何计算逻辑）
+
+SparsePolicy.compute_prefill_attention():
+  1. select_blocks(blocks, offload_engine, ctx) → 筛选 blocks
+  2. 加载 blocks（通过 offload_engine）
+  3. 遍历 blocks：
+     - 调用 self.compute_block_attention(q, k, v, ...)
+     - 调用 self.merge_attention_outputs(...)
+  4. 计算当前 chunk attention
+  5. 合并最终结果
+  6. 返回 final_output
+```
+
+## 关键设计决策
+
+| 决策 | 说明 |
+|------|------|
+| **决策 1** | `compute_block_attention` 是抽象方法，所有 policy 必须实现 |
+| **决策 2** | `merge_attention_outputs` 是抽象方法，所有 policy 必须实现 |
+| **决策 3** | `compute_prefill_attention` 是抽象方法，定义完整的 prefill 流程 |
+| **决策 4** | `select_blocks` 接收 `offload_engine` 参数（为未来准备） |
+| **决策 5** | chunked_prefill 检查 policy 是否存在，不存在则抛出错误 |
+| **决策 6** | model_runner 默认创建 FullPolicy 作为兜底 |
+| **决策 7** | attention.py 的 _chunked_prefill_attention 不包含任何 flashattn 或 merge 调用 |
 
 ## Phases
 
-### Phase 1: Create Model Registry Pattern [pending]
+- [ ] Phase 1: 分析当前架构，理解所有计算逻辑的位置
-**Files to modify:**
+- [ ] Phase 2: 在 SparsePolicy 基类中添加三个抽象方法
-- `nanovllm/models/__init__.py` (new)
+- [ ] Phase 3: 修改 FullPolicy，实现三个抽象方法
-- `nanovllm/models/registry.py` (new)
+- [ ] Phase 4: 修改 QuestPolicy，实现三个抽象方法
+- [ ] Phase 5: 修改 XAttentionBSAPolicy，实现三个抽象方法
+- [ ] Phase 6: 修改 model_runner.py，默认创建 FullPolicy
+- [ ] Phase 7: 修改 attention.py，移除所有计算逻辑，只调用 policy
+- [ ] Phase 8: 测试验证
 
-**Tasks:**
+## Phase 1: 分析当前架构，理解所有计算逻辑的位置
-1. 创建模型注册表机制
+
-2. 定义模型注册装饰器 `@register_model`
+### 当前 attention.py 中包含的计算逻辑
-3. 实现 `get_model_class(hf_config)` 函数，根据 `architectures` 字段自动选择模型
+
+1. `_ring_buffer_pipeline_load` 方法：
+   - 调用 `offload_engine.load_to_slot_layer()`
+   - 调用 `offload_engine.wait_slot_layer()`
+   - 调用 `offload_engine.get_kv_for_slot()`
+   - 调用 `flash_attn_with_lse()` ← **直接调用**
+   - 调用 `merge_attention_outputs()` ← **直接调用**
+
+2. `_sync_load_previous_chunks` 方法：
+   - 同上，直接调用 flashattn 和 merge
+
+3. `_chunked_prefill_attention` 方法：
+   - 调用 `_ring_buffer_pipeline_load` 或 `_sync_load_previous_chunks`
+   - 调用 `flash_attn_with_lse()` 计算当前 chunk
+   - 调用 `merge_attention_outputs()` 合并结果
+
+### 需要移动的计算逻辑
+
+所有 `flash_attn_with_lse` 和 `merge_attention_outputs` 调用都应该在 SparsePolicy 内部。
+
+## Phase 2: 在 SparsePolicy 基类中添加三个抽象方法
+
+### 2.1 compute_block_attention
 
-**Design:**
 ```python
-MODEL_REGISTRY: dict[str, type] = {}
+@abstractmethod
+def compute_block_attention(
+    self,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    layer_id: int,
+    softmax_scale: float,
+    causal: bool,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """
+    计算单个 block 的 attention。
 
-def register_model(*architectures):
+    Args:
-    """Decorator to register a model class for given architecture names."""
+        q: [1, seq_len, num_heads, head_dim] 或 [seq_len, num_heads, head_dim]
-    def decorator(cls):
+        k, v: 同上
-        for arch in architectures:
+        layer_id: 层索引
-            MODEL_REGISTRY[arch] = cls
+        softmax_scale: softmax 缩放因子
-        return cls
+        causal: 是否应用因果掩码
-    return decorator
 
-def get_model_class(hf_config) -> type:
+    Returns:
-    """Get model class based on HF config architectures."""
+        (o, lse) - attention 输出和 LSE
-    for arch in hf_config.architectures:
+    """
-        if arch in MODEL_REGISTRY:
+    pass
-            return MODEL_REGISTRY[arch]
-    raise ValueError(f"Unsupported architecture: {hf_config.architectures}")
 ```
 
-### Phase 2: Add Llama3 RoPE Scaling Support [pending]
+### 2.2 merge_attention_outputs
-**Files to modify:**
-- `nanovllm/layers/rotary_embedding.py`
 
-**Tasks:**
-1. 实现 `Llama3RotaryEmbedding` 类，支持 llama3 rope_type
-2. 修改 `get_rope()` 函数，根据 rope_scaling 类型选择实现
-3. 保持向后兼容（rope_scaling=None 使用原实现）
-
-**Llama3 RoPE Scaling Formula:**
 ```python
-# From transformers:
+@abstractmethod
-# low_freq_factor, high_freq_factor, original_max_position_embeddings
+def merge_attention_outputs(
-# Adjust frequencies based on wavelength thresholds
+    self,
+    o_acc: torch.Tensor,
+    lse_acc: Optional[torch.Tensor],
+    o_new: torch.Tensor,
+    lse_new: Optional[torch.Tensor],
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """
+    合并两个 attention 输出。
+
+    Args:
+        o_acc: 累积的 attention 输出 [1, seq_len, num_heads, head_dim]
+        lse_acc: 累积的 LSE
+        o_new: 新的 attention 输出
+        lse_new: 新的 LSE
+
+    Returns:
+        (merged_o, merged_lse)
+    """
+    pass
 ```
 
-### Phase 3: Implement Llama Model [pending]
+### 2.3 compute_chunked_attention
-**Files to create:**
-- `nanovllm/models/llama.py`
 
-**Tasks:**
-1. 创建 `LlamaAttention` 类（无 q_norm/k_norm，无 QKV bias）
-2. 创建 `LlamaMLP` 类（与 Qwen3MLP 类似，无 bias）
-3. 创建 `LlamaDecoderLayer` 类
-4. 创建 `LlamaModel` 和 `LlamaForCausalLM` 类
-5. 添加 `packed_modules_mapping` 以支持权重加载
-6. 使用 `@register_model("LlamaForCausalLM")` 注册
-
-### Phase 4: Modify ModelRunner for Dynamic Loading [pending]
-**Files to modify:**
-- `nanovllm/engine/model_runner.py`
-
-**Tasks:**
-1. 移除硬编码 `from nanovllm.models.qwen3 import Qwen3ForCausalLM`
-2. 导入 `from nanovllm.models import get_model_class`
-3. 替换 `self.model = Qwen3ForCausalLM(hf_config)` 为:
 ```python
-   model_class = get_model_class(hf_config)
+@abstractmethod
-   self.model = model_class(hf_config)
+def compute_chunked_attention(
+    self,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    layer_id: int,
+    softmax_scale: float,
+    offload_engine: OffloadEngine,
+    current_chunk_idx: int,
+    seq: ChunkedSequence,
+    num_tokens: int,
+) -> torch.Tensor:
+    """
+    计算 chunked prefill attention（完整流程）。
+
+    这是 policy 的主入口，定义完整的 prefill 计算流程：
+    1. 获取历史 blocks
+    2. 筛选 blocks（调用 select_blocks）
+    3. 加载和计算历史 blocks
+    4. 计算当前 chunk attention
+    5. 合并所有结果
+
+    Args:
+        q, k, v: 当前 chunk 的 QKV
+        layer_id: 层索引
+        softmax_scale: softmax 缩放因子
+        offload_engine: offload engine
+        current_chunk_idx: 当前 chunk 索引
+        seq: chunked 序列
+        num_tokens: 当前 chunk 的 token 数
+
+    Returns:
+        [seq_len, num_heads, head_dim] 最终 attention输出
+    """
+    pass
 ```
 
-### Phase 5: Register Qwen3 Model [pending]
+### 2.4 修改 select_blocks 接口
-**Files to modify:**
-- `nanovllm/models/qwen3.py`
 
-**Tasks:**
+```python
-1. 导入 `from nanovllm.models.registry import register_model`
+def select_blocks(
-2. 添加 `@register_model("Qwen3ForCausalLM", "Qwen2ForCausalLM")` 装饰器
+    self,
+    available_blocks: List[int],
+    offload_engine: OffloadEngine,
+    ctx: PolicyContext,
+) -> List[int]:
+    """
+    选择要加载的 blocks。
 
-### Phase 6: Test with Llama-3.1-8B-Instruct [pending]
+    Args:
-**Files:**
+        available_blocks: 所有可用的 block IDs
-- `tests/test_needle.py` (existing, use for validation)
+        offload_engine: offload engine（为未来准备，当前可能不使用）
+        ctx: policy context
 
-**Tasks:**
+    Returns:
-1. 运行 needle 测试: `python tests/test_needle.py --model ~/models/Llama-3.1-8B-Instruct`
+        选择的 block IDs
-2. 验证模型加载正确
+    """
-3. 验证推理输出正确
+    pass
+```
 
----
+## Phase 3: 修改 FullPolicy，实现三个抽象方法
+
+### 3.1 FullPolicy.compute_block_attention
+
+直接调用 `flash_attn_with_lse`，处理 3D 输入。
+
+### 3.2 FullPolicy.merge_attention_outputs
+
+调用 `chunked_attention.merge_attention_outputs`。
+
+### 3.3 FullPolicy.compute_prefill_attention
+
+实现完整的 prefill 流程：
+1. 获取 `cpu_block_table = kvcache_manager.get_prefilled_cpu_blocks(seq)`
+2. 调用 `select_blocks(cpu_block_table, offload_engine, ctx)`
+3. 遍历 blocks：
+   - `offload_engine.load_to_slot_layer(slot, layer_id, cpu_block_id)`
+   - `offload_engine.wait_slot_layer(slot)`
+   - `k, v = offload_engine.get_kv_for_slot(slot)`
+   - 调用 `self.compute_block_attention(q, k, v, layer_id, scale, causal=False)`
+   - 调用 `self.merge_attention_outputs(o_acc, lse_acc, prev_o, prev_lse)`
+4. 计算当前 chunk attention
+5. 合并最终结果
+
+### 需要移动的代码
+
+从 `attention.py` 的 `_ring_buffer_pipeline_load` 和 `_sync_load_previous_chunks` 移动逻辑：
+- slot 遍历逻辑
+- offload_engine 调用
+- 计算和合并逻辑
+
+从 `attention.py` 的 `_chunked_prefill_attention` 移动逻辑：
+- 当前 chunk 的 attention 计算
+- 最终合并逻辑
+
+## Phase 4: 修改 QuestPolicy
+
+QuestPolicy 实现与 FullPolicy 类似，区别在于：
+- `select_blocks` 返回 Top-K blocks
+- 其他计算逻辑相同
+
+## Phase 5: 修改 XAttentionBSAPolicy
+
+当前 XAttentionBSAPolicy 只返回所有 blocks，修改后：
+- `select_blocks` 当前返回所有 blocks
+- `compute_block_attention` 与 FullPolicy 相同
+- `merge_attention_outputs` 与 FullPolicy 相同
+- `compute_prefill_attention` 与 FullPolicy 相同
+
+未来可以实现稀疏计算。
+
+## Phase 6: 修改 model_runner.py，默认创建 FullPolicy
+
+### 6.1 当前创建 sparse policy 的逻辑
+
+```python
+# 当前：只有指定 sparse_policy_type 时才创建
+if sparse_policy_type is not None:
+    sparse_policy = create_sparse_policy(sparse_policy_type, **kwargs)
+```
+
+### 6.2 修改后
+
+```python
+# 默认创建 FullPolicy
+if sparse_policy_type is None:
+    sparse_policy_type = SparsePolicyType.FULL
+
+sparse_policy = create_sparse_policy(sparse_policy_type, **kwargs)
+```
+
+### 6.3 位置
+
+`model_runner.py` 中的 `allocate_kv_cache` 方法。
+
+## Phase 7: 修改 attention.py，移除所有计算逻辑
+
+### 7.1 _chunked_prefill_attention 简化
+
+**当前（伪代码）**：
+```python
+# 获取 cpu_block_table
+# 调用 select_blocks
+# 调用 _ring_buffer_pipeline_load（包含计算逻辑）
+# 计算当前 chunk（flash_attn）
+# 合并结果（merge）
+```
+
+**修改后**：
+```python
+sparse_policy = kvcache_manager.sparse_policy
+if sparse_policy is None:
+    raise RuntimeError("sparse_policy is required for chunked prefill")
+
+o = sparse_policy.compute_prefill_attention(
+    q, k, v, self.layer_id, self.scale,
+    offload_engine, current_chunk_idx, seq, num_tokens
+)
+
+# 直接返回，不需要合并（policy 内部已完成所有计算）
+return o
+```
+
+### 7.2 删除的方法
+
+删除以下方法（逻辑移到 policy 中）：
+- `_ring_buffer_pipeline_load` - 逻辑移到 FullPolicy.compute_prefill_attention
+- `_sync_load_previous_chunks` - 逻辑移到 FullPolicy.compute_prefill_attention
+
+### 7.3 保留的方法
+
+- `_decode_with_layer_pipeline` - decode 逻辑保持不变
+- `_decode_ring_buffer_pipeline` - decode 逻辑保持不变
+
+## Phase 8: 测试验证
+
+- [ ] 运行 `test_needle.py --enable-offload` (FULL policy)
+- [ ] 验证输出正确 (needle value: 7492)
+- [ ] 验证性能无明显下降
+
+## 关键文件清单
+
+| 文件 | 修改内容 |
+|------|----------|
+| `nanovllm/kvcache/sparse/policy.py` | 添加三个抽象方法，修改 select_blocks 签名 |
+| `nanovllm/kvcache/sparse/full_policy.py` | 实现三个抽象方法，移动计算逻辑 |
+| `nanovllm/kvcache/sparse/quest.py` | 实现三个抽象方法 |
+| `nanovllm/kvcache/sparse/xattn_bsa.py` | 实现三个抽象方法 |
+| `nanovllm/engine/model_runner.py` | 默认创建 FullPolicy |
+| `nanovllm/layers/attention.py` | 简化 _chunked_prefill_attention，删除计算方法 |
+
+## Decisions Made
+
+- **决策 1**: 三个方法都是抽象方法，强制所有 policy 实现
+- **决策 2**: compute_prefill_attention 定义完整的 prefill 流程，是 policy 的主入口
+- **决策 3**: attention.py 只调用 policy.compute_prefill_attention，零计算逻辑
+- **决策 4**: chunked_prefill 检查 policy 是否存在，不存在则抛出错误
+- **决策 5**: model_runner 默认创建 FullPolicy 作为兜底
+- **决策 6**: _ring_buffer_pipeline_load 和 _sync_load_previous_chunks 删除，逻辑移到 policy
 
 ## Errors Encountered
-| Error | Attempt | Resolution |
-|-------|---------|------------|
-| (none yet) | | |
 
----
+（待记录）
 
-## Success Criteria
+## Status
-- [x] 分析完成：理解当前架构和需要的改动
-- [ ] Phase 1: 模型注册表实现
-- [ ] Phase 2: Llama3 RoPE scaling 支持
-- [ ] Phase 3: Llama 模型实现
-- [ ] Phase 4: ModelRunner 动态加载
-- [ ] Phase 5: Qwen3 模型注册
-- [ ] Phase 6: Llama needle 测试通过
 
----
+**Currently in Phase 1** - 分析当前架构，理解所有计算逻辑的位置
-
-## Notes
-- 保持现有 Qwen3 功能不变
-- 遵循现有代码风格
-- 复用现有 layers 组件（Linear, RMSNorm, Embedding 等）
-- 只添加必要的代码，不过度工程化