🙈 chore: exclude planning-with-files from git tracking

- Add planning files (task_plan.md, findings.md, progress.md) to .gitignore
- Remove existing planning files from git index (keep local)
- Update planning-with-files rule with git management policy

These temporary session files should not be version controlled.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Zijie Tian
2026-01-20 02:06:28 +08:00
parent e5a17c832c
commit 6080bf7554
4 changed files with 40 additions and 829 deletions

View File

@@ -1,5 +1,37 @@
# Planning with Files Rule # Planning with Files Rule
## Git 管理政策
**重要**Planning 文件已从 Git 管理中排除,不会被提交。
### 已配置的 .gitignore 规则
```gitignore
# Planning-with-files temporary files
task_plan.md
findings.md
progress.md
task_plan_*.md
findings_*.md
progress_*.md
```
### 为什么排除这些文件
1. **临时性质**:计划文件是会话级别的临时文件,不应进入版本控制
2. **避免冲突**:多实例并行开发时,不同任务的计划文件会产生冲突
3. **保持仓库整洁**:这些文件只对当前任务有用,不需要历史记录
### 如果不小心已经 commit 了
```bash
# 从 git 中移除(保留本地文件)
git rm --cached task_plan.md findings.md progress.md
git commit -m "chore: remove planning files from git tracking"
```
---
## 自动清理旧计划文件 ## 自动清理旧计划文件
**重要**:每次开始新的复杂任务使用 planning-with-files 时,先删除旧的计划文件。 **重要**:每次开始新的复杂任务使用 planning-with-files 时,先删除旧的计划文件。

8
.gitignore vendored
View File

@@ -230,3 +230,11 @@ tests/data/
# Serena MCP tool config # Serena MCP tool config
.serena/ .serena/
# Planning-with-files temporary files
task_plan.md
findings.md
progress.md
task_plan_*.md
findings_*.md
progress_*.md

View File

@@ -1,467 +0,0 @@
# Task Plan: Sparse Policy 架构重构 v4 (FullPolicy Only)
## Goal
将 chunked prefill 的 attention 计算逻辑完全从 `attention.py` 移到 `SparsePolicy` 内部。
### 验收标准(必须全部满足)
| # | 标准 | 说明 |
|---|------|------|
| **1** | `test_needle.py --enable-offload` 通过 | 功能正确性验证 |
| **2** | `attention.py` 中 chunked prefill 路径零计算调用 | 不直接调用 `flash_attn_*``merge_attention_outputs`,全部由 policy 完成 |
| **3** | 所有 KV 通信由 `offload_engine` 完成 | 不直接调用 `torch.copy_``.copy()` 进行 KV 数据传输 |
**范围**: 仅实现 FullPolicy暂不涉及 QuestPolicy 和 XAttentionBSAPolicy。Decode 阶段不处理。
## 当前代码状态(重要发现)
**`FullPolicy.compute_prefill_attention` 已经实现了完整的 prefill 流程!**
`attention.py` 没有调用它,而是:
- 调用 `sparse_policy.select_blocks()` 仅做 block 筛选
- 自己实现 `_ring_buffer_pipeline_load``_sync_load_previous_chunks`
- 自己调用 `flash_attn_with_lse``merge_attention_outputs`
**结论**:当前代码有冗余,同样的逻辑在两个地方实现。
### 当前 attention.py 中的违规调用(需要移除)
```python
# 直接计算调用(违反目标 2
flash_attn_with_lse(...)
merge_attention_outputs(...)
# 直接通信调用(违反目标 3
offload_engine.prefill_k_buffer[self.layer_id, :num_tokens].copy_(k)
offload_engine.prefill_v_buffer[self.layer_id, :num_tokens].copy_(v)
```
## 核心设计原则
1. **Policy 内部完成所有 prefill 计算**:包括 block 加载、attention 计算和结果合并
2. **select_blocks 传入 offload_engine**其他策略Quest/XAttn可能需要加载 KV 来判断
3. **统一方法命名**:使用 `compute_chunked_attention`(不是 `compute_prefill_attention`
4. **chunked_prefill 强制 policy 存在**:没有 policy 则报错
5. **attention.py 零计算逻辑**`_chunked_prefill_attention` 只调用 policy
6. **所有 KV 通信通过 offload_engine**:不直接调用 torch.copy
## 目标架构
```
attention.py (_chunked_prefill_attention):
检查 sparse_policy 是否存在
调用 sparse_policy.compute_chunked_attention(q, k, v, ...)
处理 async offload通过 offload_engine
返回最终输出(不包含任何计算逻辑,不包含任何直接 copy 调用)
SparsePolicy.compute_chunked_attention():
1. 获取 cpu_block_table
2. 调用 select_blocks(blocks, offload_engine, ctx) → 筛选 blocks
3. 通过 offload_engine 加载 blocks 并计算 attentionpipeline 或 sync
4. 通过 offload_engine 获取当前 chunk KV计算 attentioncausal
5. 合并所有结果
6. 返回 final_output
```
## 关键设计决策
| 决策 | 说明 |
|------|------|
| **决策 1** | `compute_chunked_attention` 是唯一的抽象方法,定义完整 prefill 流程 |
| **决策 2** | 不添加 `compute_block_attention``merge_attention_outputs` 抽象方法(过度设计) |
| **决策 3** | `select_blocks` 接收 `offload_engine` 参数(其他策略需要) |
| **决策 4** | attention.py 的 `_chunked_prefill_attention` 不包含任何 flashattn 或 merge 调用 |
| **决策 5** | Decode 阶段不处理,保持现有逻辑 |
| **决策 6** | async offload 逻辑保留在 attention.py通过 offload_engine 方法调用) |
| **决策 7** | Phase 4 需要添加 debug 输出验证执行路径 |
| **决策 8** | 所有 KV 通信必须通过 offload_engine 方法,不直接调用 torch.copy |
## Phases
- [x] Phase 1: 分析当前架构 ✅ 已完成
- [ ] Phase 2: 修改 SparsePolicy 基类
- [ ] Phase 3: 修改 FullPolicy
- [ ] Phase 4: 验证执行路径(添加 debug 输出)
- [ ] Phase 5: 修改 attention.py
- [ ] Phase 6: 测试验证
## Phase 1: 分析当前架构 ✅ 已完成
### 当前 attention.py 中包含的计算逻辑(需要移除)
1. `_ring_buffer_pipeline_load` 方法:直接调用 flashattn 和 merge
2. `_sync_load_previous_chunks` 方法:直接调用 flashattn 和 merge
3. `_chunked_prefill_attention` 方法:
- 调用上述两个方法
- 计算当前 chunkflash_attn
- 合并结果merge
### 当前 attention.py 中的直接 copy 调用(需要移除或封装)
```python
# attention.py:115-116 - 写入 prefill buffer
offload_engine.prefill_k_buffer[self.layer_id, :num_tokens].copy_(k)
offload_engine.prefill_v_buffer[self.layer_id, :num_tokens].copy_(v)
```
**处理方案**:在 offload_engine 中添加封装方法,或将此逻辑移入 policy。
### 当前 FullPolicy 已实现的功能
`full_policy.py:40-162``compute_prefill_attention` 已实现:
- ring buffer pipeline 加载
- sync 加载 fallback
- 当前 chunk attention 计算
- 结果合并
**只需重命名为 `compute_chunked_attention` 并微调接口。**
## Phase 2: 修改 SparsePolicy 基类
### 2.1 修改 select_blocks 接口
```python
@abstractmethod
def select_blocks(
self,
available_blocks: List[int],
offload_engine: "OffloadEngine", # 新增参数
ctx: PolicyContext,
) -> List[int]:
"""
选择要加载的 blocks。
Args:
available_blocks: 所有可用的 block IDs
offload_engine: offload engine其他策略可能需要加载 KV 来判断)
ctx: policy context
Returns:
选择的 block IDs
"""
pass
```
### 2.2 添加 compute_chunked_attention 抽象方法
```python
@abstractmethod
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. 通过 offload_engine 加载和计算历史 blocks
4. 通过 offload_engine 获取当前 chunk KV计算 attention
5. 合并所有结果
Args:
q: [seq_len, num_heads, head_dim] 当前 chunk 的 query
k, v: [seq_len, num_kv_heads, head_dim] 当前 chunk 的 KV已写入 prefill buffer
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 3: 修改 FullPolicy
### 3.1 重命名方法
`compute_prefill_attention` 重命名为 `compute_chunked_attention`
### 3.2 修改 select_blocks 签名
```python
def select_blocks(
self,
available_blocks: List[int],
offload_engine: "OffloadEngine", # 新增参数(不使用)
ctx: PolicyContext,
) -> List[int]:
"""Return all blocks - no sparsity."""
return available_blocks
```
### 3.3 验证 compute_chunked_attention 实现
当前 `compute_prefill_attention` 已实现完整逻辑,确认:
- [x] 获取 cpu_block_table
- [x] ring buffer pipeline 加载(通过 offload_engine
- [x] sync 加载 fallback通过 offload_engine
- [x] 当前 chunk attention 计算
- [x] 结果合并
**注意**:当前实现没有调用 `select_blocks`,需要添加。
### 3.4 确保所有 KV 通信通过 offload_engine
检查 `compute_chunked_attention` 内部:
- 历史 block 加载:已通过 `offload_engine.load_to_slot_layer()` 等方法 ✅
- 当前 chunk KV 获取:已通过 `offload_engine.get_prefill_buffer_slice()`
## Phase 4: 验证执行路径(添加 debug 输出)
### 4.1 验证目标
确认代码修改后,执行路径正确:
| 检查点 | 位置 | 预期行为 |
|--------|------|----------|
| **Policy 创建** | `kvcache/__init__.py` | FullAttentionPolicy 被创建 |
| **Policy 调用** | `attention.py` | `_chunked_prefill_attention` 调用 `sparse_policy.compute_chunked_attention` |
| **select_blocks 调用** | `full_policy.py` | `compute_chunked_attention` 内部调用 `select_blocks` |
| **旧方法未调用** | `attention.py` | `_ring_buffer_pipeline_load``_sync_load_previous_chunks` 不再被调用 |
| **无直接 copy 调用** | `attention.py` | chunked prefill 路径不直接调用 `.copy_()` |
### 4.2 添加 debug 输出位置
**位置 1: `kvcache/__init__.py` - policy 创建时**
```python
sparse_policy = create_sparse_policy(sparse_policy_type, **policy_kwargs)
logger.info(f"[DEBUG] Created sparse policy: {sparse_policy}")
```
**位置 2: `attention.py` - 调用 policy 时**
```python
# 在 _chunked_prefill_attention 中
logger.debug(f"[DEBUG] Calling sparse_policy.compute_chunked_attention, "
f"policy={sparse_policy}, layer={self.layer_id}, chunk={current_chunk_idx}")
```
**位置 3: `full_policy.py` - compute_chunked_attention 入口**
```python
def compute_chunked_attention(self, ...):
logger.debug(f"[DEBUG] FullPolicy.compute_chunked_attention called, "
f"layer={layer_id}, chunk={current_chunk_idx}, num_tokens={num_tokens}")
# ... 实现
```
**位置 4: `full_policy.py` - select_blocks 调用**
```python
# 在 compute_chunked_attention 内部
selected_blocks = self.select_blocks(cpu_block_table, offload_engine, policy_ctx)
logger.debug(f"[DEBUG] select_blocks: input={len(cpu_block_table)} blocks, "
f"output={len(selected_blocks)} blocks")
```
### 4.3 验证方法
运行测试并检查日志输出:
```bash
PYTHONPATH=/home/zijie/Code/nano-vllm:$PYTHONPATH \
python tests/test_needle.py --model <model_path> --enable-offload 2>&1 | grep DEBUG
```
预期输出:
```
[DEBUG] Created sparse policy: FullAttentionPolicy()
[DEBUG] Calling sparse_policy.compute_chunked_attention, policy=FullAttentionPolicy(), layer=0, chunk=0
[DEBUG] FullPolicy.compute_chunked_attention called, layer=0, chunk=0, num_tokens=...
[DEBUG] select_blocks: input=0 blocks, output=0 blocks
[DEBUG] Calling sparse_policy.compute_chunked_attention, policy=FullAttentionPolicy(), layer=0, chunk=1
[DEBUG] FullPolicy.compute_chunked_attention called, layer=0, chunk=1, num_tokens=...
[DEBUG] select_blocks: input=1 blocks, output=1 blocks
...
```
### 4.4 清理 debug 输出
验证完成后,将 debug 级别的日志改为更低级别(如 `logger.debug`),或通过环境变量控制:
```python
if os.environ.get('NANOVLLM_DEBUG_POLICY'):
logger.info(f"[DEBUG] ...")
```
## Phase 5: 修改 attention.py
### 5.1 简化 _chunked_prefill_attention
**修改后**
```python
def _chunked_prefill_attention(self, q, k, v, context):
kvcache_manager = context.kvcache_manager
seq = context.chunked_seq
offload_engine = kvcache_manager.offload_engine
current_chunk_idx = context.current_chunk_idx
num_tokens = k.shape[0]
# 获取 sparse policy
sparse_policy = kvcache_manager.sparse_policy
if sparse_policy is None:
raise RuntimeError("sparse_policy is required for chunked prefill")
# [DEBUG] 验证执行路径
logger.debug(f"[DEBUG] Calling sparse_policy.compute_chunked_attention, "
f"policy={sparse_policy}, layer={self.layer_id}, chunk={current_chunk_idx}")
# 调用 policy 计算 attention所有计算逻辑在 policy 内部)
# 注意:不直接调用 flash_attn 或 merge全部由 policy 完成
final_o = sparse_policy.compute_chunked_attention(
q, k, v,
self.layer_id,
self.scale,
offload_engine,
current_chunk_idx,
seq,
num_tokens,
)
# Per-layer ASYNC offload通过 offload_engine 方法,不直接 copy
if offload_engine is not None and seq is not None:
cpu_block_ids, _ = kvcache_manager.get_all_cpu_blocks(seq)
if current_chunk_idx < len(cpu_block_ids):
cpu_block_id = cpu_block_ids[current_chunk_idx]
offload_engine.offload_prefill_buffer_async(
self.layer_id, cpu_block_id, num_tokens
)
return final_o
```
### 5.2 处理 prefill buffer 写入
当前 `forward()` 方法中有直接 copy 调用:
```python
# 当前代码(违反目标 3
offload_engine.prefill_k_buffer[self.layer_id, :num_tokens].copy_(k)
offload_engine.prefill_v_buffer[self.layer_id, :num_tokens].copy_(v)
```
**方案 A**:在 offload_engine 中添加封装方法
```python
# offload_engine.py
def write_prefill_buffer(self, layer_id: int, k: Tensor, v: Tensor, num_tokens: int):
self.prefill_k_buffer[layer_id, :num_tokens].copy_(k)
self.prefill_v_buffer[layer_id, :num_tokens].copy_(v)
# attention.py
offload_engine.write_prefill_buffer(self.layer_id, k, v, num_tokens)
```
**方案 B**:将此逻辑移入 policy作为 compute_chunked_attention 的一部分)
**推荐方案 A**:保持 attention.py 调用 offload_engine 方法,但不直接操作 buffer。
### 5.3 删除的方法
删除以下方法(逻辑已移到 FullPolicy
- `_ring_buffer_pipeline_load`
- `_sync_load_previous_chunks`
### 5.4 保留的方法
Decode 相关方法保持不变:
- `_chunked_decode_attention`
- `_decode_with_layer_pipeline`
- `_decode_ring_buffer_pipeline`
## Phase 6: 测试验证
### 6.1 功能测试
- [ ] 运行 `test_needle.py --enable-offload` (FULL policy)
- [ ] 验证输出正确needle value 匹配)
- [ ] 检查 debug 日志确认执行路径正确
### 6.2 代码审查(验收标准检查)
- [ ] **标准 1**: test_needle.py 通过 ✓
- [ ] **标准 2**: `_chunked_prefill_attention` 方法内无 `flash_attn``merge_attention_outputs` 调用
- [ ] **标准 3**: `_chunked_prefill_attention` 方法内无直接 `.copy_()` 调用
**注意**:标准 2 和 3 仅适用于 chunked prefill 路径。Decode 路径和其他路径可以有 `flash_attn` 调用。
**验证方法**
**方法 1使用 cclsp LSP 工具验证调用链(推荐)**
使用 `mcp__cclsp__find_references` 查找计算函数的调用位置,确认 chunked prefill 路径无直接调用:
```
# 查找 flash_attn_with_lse 的所有调用
mcp__cclsp__find_references(file_path="nanovllm/layers/attention.py", symbol_name="flash_attn_with_lse")
# 查找 merge_attention_outputs 的所有调用
mcp__cclsp__find_references(file_path="nanovllm/layers/attention.py", symbol_name="merge_attention_outputs")
# 查找 _chunked_prefill_attention 的实现
mcp__cclsp__find_definition(file_path="nanovllm/layers/attention.py", symbol_name="_chunked_prefill_attention")
```
验证结果应显示:
- `flash_attn_with_lse` 调用仅出现在 decode 路径或 `full_policy.py`
- `_chunked_prefill_attention` 内部只调用 `sparse_policy.compute_chunked_attention`
**方法 2手动代码审查**
检查 `_chunked_prefill_attention` 方法实现,确认:
1. 只调用 `sparse_policy.compute_chunked_attention(...)`
2. 只调用 `offload_engine.offload_prefill_buffer_async(...)` 等 offload_engine 方法
3. 不直接调用 `flash_attn_*``merge_attention_outputs``.copy_()`
```bash
# 辅助检查:找出所有 flash_attn 调用位置
grep -n "flash_attn\|merge_attention_outputs" nanovllm/layers/attention.py
# 辅助检查:找出所有 copy 调用位置
grep -n "\.copy_\|\.copy(" nanovllm/layers/attention.py
```
### 6.3 回归测试
- [ ] 验证 decode 阶段不受影响
- [ ] 验证非 offload 模式不受影响(如果适用)
## 关键文件清单
| 文件 | 修改内容 |
|------|----------|
| `nanovllm/kvcache/sparse/policy.py` | 添加 `compute_chunked_attention` 抽象方法,修改 `select_blocks` 签名 |
| `nanovllm/kvcache/sparse/full_policy.py` | 重命名方法,修改 `select_blocks` 签名,添加 `select_blocks` 调用,添加 debug 输出 |
| `nanovllm/layers/attention.py` | 简化 `_chunked_prefill_attention`,删除 `_ring_buffer_pipeline_load``_sync_load_previous_chunks`,添加 debug 输出 |
| `nanovllm/kvcache/__init__.py` | 添加 policy 创建的 debug 输出 |
| `nanovllm/kvcache/offload_engine.py` | (可选)添加 `write_prefill_buffer` 方法封装 |
## Decisions Made
- **决策 1**: 只添加一个抽象方法 `compute_chunked_attention`(不添加 `compute_block_attention``merge_attention_outputs`
- **决策 2**: `select_blocks` 接收 `offload_engine` 参数
- **决策 3**: 统一使用 `compute_chunked_attention` 命名
- **决策 4**: Decode 阶段不处理
- **决策 5**: async offload 逻辑保留在 attention.py通过 offload_engine 方法调用)
- **决策 6**: Phase 4 添加 debug 输出验证执行路径,验证完成后可降级或移除
- **决策 7**: prefill buffer 写入通过 offload_engine 封装方法实现(方案 A
- **决策 8**: 所有 KV 通信必须通过 offload_engine 方法,不直接调用 torch.copy
## Errors Encountered
(待记录)
## Status
**Planning Complete** - v4 计划已完成,包含明确的验收标准和执行路径验证步骤

View File

@@ -1,362 +0,0 @@
# Task Plan: XAttention BSA 模块化集成
## Goal
将 XAttention BSA 策略按照统一接口集成到 nano-vllm 的 sparse policy 框架中,实现模块化设计。
**最终验证目标**: 运行 `tests/test_ruler.py` 测试 32K 数据的 10 个以内的 sample得到合理结果不一定全部 PASS但结果应在预期精度范围内
---
## 强制要求:使用 Hive-Mind 集群思考
**必须使用 Claude Flow MCP 的 hive-mind 集群进行深度推理,提高实现精度。**
### 启动 Hive-Mind 的方式
在每个复杂阶段开始前,必须执行以下步骤:
1. **初始化 Hive-Mind 集群**
```python
# 通过 MCP 调用
mcp__claude-flow_alpha__hive-mind_init(
topology="mesh", # 或 "hierarchical", "ring", "star"
maxAgents=5, # 集群大小
)
```
2. **生成专业代理Spawning Specialists**
```python
# 为不同任务类型创建代理
mcp__claude-flow_alpha__hive-mind_spawn(
count=3,
type="specialist", # researcher, coder, analyst
)
```
3. **广播思考任务**
```python
mcp__claude-flow_alpha__hive-mind_broadcast(
message="分析当前架构设计的潜在问题...",
priority="high"
)
```
4. **获取集群状态和共识**
```python
mcp__claude-flow_alpha__hive-mind_status(verbose=True)
mcp__claude-flow_alpha__hive-mind_consensus(
action="propose",
type="design",
value="模块化接口设计方案"
)
```
### 适用阶段
以下阶段**必须**使用 Hive-Mind 集群思考:
- ✅ Phase 1: SparsePolicy 基类接口确认
- ✅ Phase 2: XAttentionBSAPolicy 接口对齐
- ✅ Phase 3: OffloadEngine 辅助方法模块化
- ✅ Phase 5: attention.py 集成点验证
其他阶段Phase 4, 6, 7可以使用标准思考模式。
### 集群配置建议
```yaml
# 推荐配置
topology: mesh # 网状拓扑,适合并行推理
maxAgents: 5 # 5个专业代理
agentTypes:
- researcher # 架构分析
- coder # 代码实现
- analyst # 接口验证
- optimizer # 性能优化
- validator # 正确性验证
```
### 输出要求
使用 Hive-Mind 后,必须在计划中记录:
1. 集群产生的关键洞察
2. 多代理共识达成的决策
3. 发现的潜在问题和解决方案
---
## 当前架构分析
### SparsePolicy 基类接口
从 `nanovllm/kvcache/sparse/policy.py` 需要确认基类定义:
```python
class SparsePolicy:
# 能力标记
supports_prefill: bool
supports_decode: bool
requires_block_selection: bool
# 核心方法
def select_blocks(self, available_blocks: List[int], ctx: PolicyContext) -> List[int]
# 可选方法prefill 专用)
def sparse_prefill_attention(self, q, k, v, layer_id) -> torch.Tensor
# 初始化
def initialize(self, num_layers, num_kv_heads, head_dim, num_cpu_blocks, dtype, device)
def reset(self)
```
### 当前 XAttentionBSAPolicy 实现
已实现但需要确认模块化集成的部分:
- `xattn_bsa.py` - 策略类实现
- `config.py` - 枚举和参数
- `sparse/__init__.py` - 策略工厂
- `offload_engine.py` - 辅助方法
- `attention.py` - 集成点
## 详细实现计划
### Phase 1: 确保 SparsePolicy 基类接口统一
**任务**: 验证 `SparsePolicy` 基类定义是否包含所有必需的方法
**步骤**:
1. 读取 `nanovllm/kvcache/sparse/policy.py`
2. 确认基类定义包含:
- `supports_prefill`, `supports_decode`, `requires_block_selection` 类属性
- `select_blocks()` 方法
- `sparse_prefill_attention()` 方法(可选)
- `initialize()`, `reset()` 方法
3. 如果缺失,补充到基类定义中
**预期结果**: 基类定义完整,所有策略类可以遵循统一接口
---
### Phase 2: XAttentionBSAPolicy 接口对齐
**任务**: 确保 XAttentionBSAPolicy 完全符合 SparsePolicy 接口
**步骤**:
1. 确认 `xattn_bsa.py` 中的类属性正确:
```python
class XAttentionBSAPolicy(SparsePolicy):
supports_prefill = True
supports_decode = False
requires_block_selection = False # 注意BSA 内部处理选择
```
2. 确保方法签名与基类一致:
- `select_blocks(available_blocks, ctx) -> List[int]`
- `sparse_prefill_attention(q, k, v, layer_id) -> Tensor`
- `initialize(...)`
- `reset()`
3. 添加文档说明BSA 在 prefill 阶段内部处理 block 选择,因此 `select_blocks` 返回所有可用块
**预期结果**: XAttentionBSAPolicy 完全符合 SparsePolicy 统一接口
---
### Phase 3: OffloadEngine 辅助方法模块化
**任务**: 确保 OffloadEngine 的辅助方法正确定义且模块化
**步骤**:
1. 确认 `offload_engine.py` 中的辅助方法位置:
```python
# 在 OffloadEngine 类中添加这两个方法
def load_block_sample_from_cpu(self, cpu_block_id, layer_id, num_samples):
"""加载采样 tokens 用于估算阶段"""
...
def load_block_full_from_cpu(self, cpu_block_id, layer_id):
"""加载完整 block 用于计算阶段"""
...
```
2. 确保方法签名与 `xattn_bsa.py` 中的调用一致
3. 添加适当的文档说明这两个方法的用途和使用场景
**预期结果**: OffloadEngine 提供统一的 block 加载接口
---
### Phase 4: 模块化集成到工厂模式
**任务**: 确保策略创建通过统一的工厂模式
**步骤**:
1. 检查 `nanovllm/kvcache/__init__.py` 中的 `create_kvcache_manager` 函数
2. 确认策略创建逻辑清晰:
```python
# 根据策略类型构建相应的 kwargs
if sparse_policy_type == SparsePolicyType.XATTN_BSA:
policy_kwargs = {
'block_size': getattr(config, 'sparse_block_size', 128),
'samples_per_chunk': getattr(config, 'sparse_samples_per_chunk', 128),
'threshold': getattr(config, 'sparse_threshold', 0.9),
'use_triton': getattr(config, 'sparse_use_triton', True),
'stride': getattr(config, sparse_stride', 8),
}
```
3. 确认所有策略类型都有相应的 kwargs 构建逻辑
**预期结果**: 通过 `create_sparse_policy()` 创建所有策略
---
### Phase 5: attention.py 集成点验证
**任务**: 确保 attention.py 中的集成点正确调用策略接口
**步骤**:
1. 检查 `nanovllm/layers/attention.py` 中的 `_chunked_prefill_attention` 方法
2. 确认集成逻辑:
```python
# 检测策略是否有 sparse_prefill_attention 方法
if sparse_policy is not None and hasattr(sparse_policy, 'sparse_prefill_attention'):
if sparse_policy.supports_prefill:
# 使用策略的 sparse_prefill_attention 方法
o = sparse_policy.sparse_prefill_attention(q, k, v, self.layer_id)
# 处理异步 offload
return o
# 否则使用标准流程Quest, etc.
# ...
```
3. 确保没有绕过策略接口直接调用其他逻辑
**预期结果**: attention.py 通过统一的策略接口调用 BSA
---
### Phase 6: 配置参数模块化
**任务**: 确保配置参数结构清晰,易于使用
**步骤**:
1. 检查 `nanovllm/config.py` 中的配置结构
2. 确认 XAttention BSA 参数组织清晰:
```python
# 通用 sparse 参数
sparse_policy: SparsePolicyType = SparsePolicyType.FULL
sparse_topk_blocks: int = 8 # Quest
sparse_threshold_blocks: int = 4 # Quest
# XATTN_BSA 专用参数
sparse_block_size: int = 128
sparse_samples_per_chunk: int = 128
sparse_threshold: float = 0.9
sparse_use_triton: bool = True
sparse_stride: int = 8
```
3. 考虑是否需要参数分组或嵌套配置
**预期结果**: 配置参数清晰,易于理解和使用
---
### Phase 7: 模块化验证测试
**任务**: 创建简单的验证脚本确保模块化集成正确
**步骤**:
1. 创建 `tests/test_xattn_bsa_integration.py` 测试脚本
2. 验证以下功能:
- XAttentionBSAPolicy 可以通过 `create_sparse_policy()` 创建
- 策略正确响应 `supports_prefill`, `supports_decode` 查询
- `select_blocks()` 方法返回正确结果
- OffloadEngine 辅助方法可以正常调用
- 在模拟环境中策略可以被正确调用
3. 测试用例:
```python
# Test 1: 策略创建
from nanovllm.config import Config, SparsePolicyType
from nanovllm.kvcache.sparse import create_sparse_policy
policy = create_sparse_policy(SparsePolicyType.XATTN_BSA)
assert hasattr(policy, 'sparse_prefill_attention')
assert policy.supports_prefill == True
assert policy.supports_decode == False
# Test 2: 接口一致性
# 验证方法签名
# ...
# Test 3: OffloadEngine 辅助方法
# ...
```
**预期结果**: 所有测试通过,模块化集成验证成功
---
## 关键设计原则
### 1. 接口统一性
- 所有策略通过 `SparsePolicy` 基类提供统一接口
- 工厂模式创建策略实例
- 策略切换透明,不影响其他模块
### 2. 模块化独立性
- 每个策略类独立实现
- OffloadEngine 提供通用辅助方法
- attention.py 通过策略接口调用,不依赖具体实现
### 3. 可扩展性
- 添加新策略只需:
1. 创建新的策略类继承 `SparsePolicy`
2. 添加到 `SparsePolicyType` 枚举
3. 在工厂函数中添加创建逻辑
4. 添加相应的配置参数
---
## 文件修改清单
### 必须修改的文件
1. `nanovllm/kvcache/sparse/policy.py` - 确保基类定义完整
2. `nanovllm/kvcache/sparse/xattn_bsa.py` - 确保接口对齐
3. `nanovllm/kvcache/offload_engine.py` - 添加辅助方法
4. `nanovllm/layers/attention.py` - 验证集成点
5. `nanovllm/config.py` - 确认参数结构
6. `nanovllm/kvcache/__init__.py` - 确认工厂模式
7. `nanovllm/kvcache/sparse/__init__.py` - 确认注册逻辑
### 可选创建的文件
- `tests/test_xattn_bsa_integration.py` - 集成验证测试
---
## 实现状态
- [ ] Phase 1: SparsePolicy 基类接口确认
- [ ] Phase 2: XAttentionBSAPolicy 接口对齐
- [ ] Phase 3: OffloadEngine 辅助方法模块化
- [ ] Phase 4: 工厂模式集成验证
- [ ] Phase 5: attention.py 集成点验证
- [ ] Phase 6: 配置参数模块化
- [ ] Phase 7: 模块化验证测试
---
## 备注
- 此计划专注于模块化集成,不涉及算法优化
- 所有修改都遵循现有框架的设计模式
- 重点在于接口统一和模块解耦
- 测试阶段使用简单脚本验证即可,不需要完整的端到端测试