zijie-tian/nano-vllm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[claudesquad] update from 'fix-bug-2' on 09 Jan 26 16:05 CST

Browse Source
This commit is contained in:
Zijie Tian
2026-01-09 16:05:36 +08:00
parent ccf04d3917
commit 1425510a2e
3 changed files with 267 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

170
nanovllm/engine/model_runner.py
View File

@@ -45,14 +45,7 @@ class ModelRunner:
self.allocate_kv_cache()
if not self.enforce_eager:
if config.enable_cpu_offload:
# TODO: Implement capture_offload_cudagraph() for offload mode
# For now, offload mode uses eager execution
# The standard capture_cudagraph() cannot be used because:
# - It captures the PagedAttention decode path via Attention.forward()
# - In offload mode, Attention.k_cache/v_cache are empty (KV is in ring buffer)
# - The refactored offload decode now uses Attention.forward() with ring buffer
# - Need specialized graph capture that sets up ring buffer correctly
pass
self.capture_offload_cudagraph()
else:
self.capture_cudagraph()
torch.set_default_device("cpu")
@@ -74,7 +67,10 @@ class ModelRunner:
if self.rank == 0:
self.shm.unlink()
if not self.enforce_eager:
del self.graphs, self.graph_pool
if hasattr(self, 'graphs'):
del self.graphs, self.graph_pool
if hasattr(self, 'offload_graphs'):
del self.offload_graphs, self.offload_graph_pool
# torch.cuda.synchronize()
dist.destroy_process_group()
@@ -858,6 +854,7 @@ class ModelRunner:
- Uses standard Attention.forward() path (not bypassing)
- Per-layer decode buffer for accumulating new tokens
- Async block offload when decode buffer is full
- Uses CUDA graphs when available (not enforce_eager)
"""
assert len(seqs) == 1, "Layer-wise offload only supports single sequence"
seq = seqs[0]
@@ -867,9 +864,20 @@ class ModelRunner:
num_layers = len(self.model.model.layers)
num_buffers = offload_engine.num_kv_buffers
# Check if using CUDA graphs
use_cuda_graph = not self.enforce_eager and hasattr(self, 'offload_graphs')
# 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")
if use_cuda_graph:
# Use fixed-address tensors for graph replay
graph_vars = self.offload_graph_vars
graph_vars["input_ids"][0] = seq.last_token
graph_vars["positions"][0] = len(seq) - 1
input_ids = graph_vars["input_ids"]
positions = graph_vars["positions"]
else:
input_ids = torch.tensor([seq.last_token], dtype=torch.int64, device="cuda")
positions = torch.tensor([len(seq) - 1], dtype=torch.int64, device="cuda")
# Get prefilled CPU blocks and compute valid tokens per block
cpu_block_table = self.kvcache_manager.get_prefilled_cpu_blocks(seq)
@@ -898,8 +906,14 @@ class ModelRunner:
context_len = total_prefill_tokens + num_prev_decode_tokens
# Context setup for Attention.forward() - contiguous mode (no block tables)
slot_mapping = torch.tensor([context_len], dtype=torch.int32, device="cuda")
context_lens = torch.tensor([context_len + 1], dtype=torch.int32, device="cuda")
if use_cuda_graph:
graph_vars["slot_mapping"][0] = context_len
graph_vars["context_lens"][0] = context_len + 1
slot_mapping = graph_vars["slot_mapping"]
context_lens = graph_vars["context_lens"]
else:
slot_mapping = torch.tensor([context_len], dtype=torch.int32, device="cuda")
context_lens = torch.tensor([context_len + 1], dtype=torch.int32, device="cuda")
# Phase 1: Preload first N layers to ring buffer (fill pipeline)
num_preload = min(num_buffers, num_layers)
@@ -910,8 +924,14 @@ class ModelRunner:
# Step 1: Embedding (on compute stream)
with torch.cuda.stream(compute_stream):
hidden_states = self.model.model.embed_tokens(input_ids)
residual = None
if use_cuda_graph:
# Copy embedding output to graph's hidden_states
embedded = self.model.model.embed_tokens(input_ids)
graph_vars["hidden_states"].copy_(embedded)
graph_vars["residual"].zero_() # Reset residual for first layer
else:
hidden_states = self.model.model.embed_tokens(input_ids)
residual = None
# Phase 2: Layer-by-layer processing with ring buffer pipeline
for layer_id in range(num_layers):
@@ -947,12 +967,22 @@ class ModelRunner:
block_tables=None, # Contiguous mode, no block tables
)
# 2e. Forward through layer using standard path
# This calls Qwen3Attention.forward() -> Attention.forward()
# Attention.forward() will:
# - Store new K,V to ring buffer via store_kvcache
# - Compute attention via flash_attn_with_kvcache
hidden_states, residual = layer(positions, hidden_states, residual)
if use_cuda_graph:
# 2e. Replay CUDA graph for this layer
self.offload_graphs[layer_id].replay()
# Synchronize to ensure graph completes before next operation
torch.cuda.current_stream().synchronize()
# Copy outputs to inputs for next layer
if layer_id < num_layers - 1:
graph_vars["hidden_states"].copy_(graph_vars["layer_outputs"])
graph_vars["residual"].copy_(graph_vars["layer_residual"])
else:
# 2e. Forward through layer using standard path (eager mode)
# This calls Qwen3Attention.forward() -> Attention.forward()
# Attention.forward() will:
# - Store new K,V to ring buffer via store_kvcache
# - Compute attention via flash_attn_with_kvcache
hidden_states, residual = layer(positions, hidden_states, residual)
# 2f. Copy new token's KV from ring buffer to decode buffer (for persistence)
# The new token was stored at position context_len in ring buffer
@@ -972,7 +1002,12 @@ class ModelRunner:
)
# Step 3: Final norm
hidden_states, _ = self.model.model.norm(hidden_states, residual)
if use_cuda_graph:
hidden_states, _ = self.model.model.norm(
graph_vars["layer_outputs"], graph_vars["layer_residual"]
)
else:
hidden_states, _ = self.model.model.norm(hidden_states, residual)
# Step 4: Compute logits
logits = self.model.compute_logits(hidden_states)
@@ -1036,3 +1071,94 @@ class ModelRunner:
block_tables=block_tables,
outputs=outputs,
)
@torch.inference_mode()
def capture_offload_cudagraph(self):
"""
Capture CUDA graphs for offload decode using ring buffer.
Key design:
- Captures per-layer graphs (not full decode)
- Each layer's graph uses its corresponding ring buffer slot
- H2D transfers happen outside the graph
- Graph replays single layer forward pass
Ring buffer mapping: buffer_idx = layer_id % num_buffers
"""
offload_engine = self.kvcache_manager.offload_engine
num_layers = len(self.model.model.layers)
num_buffers = offload_engine.num_kv_buffers
hf_config = self.config.hf_config
logger.info(f"Capturing offload CUDA graphs: {num_layers} layers, {num_buffers} buffers")
# Fixed-address tensors for graph capture (batch_size=1 for offload)
input_ids = torch.zeros(1, dtype=torch.int64, device="cuda")
positions = torch.zeros(1, dtype=torch.int64, device="cuda")
slot_mapping = torch.zeros(1, dtype=torch.int32, device="cuda")
context_lens = torch.ones(1, dtype=torch.int32, device="cuda") # At least 1 for valid attention
hidden_states = torch.randn(1, hf_config.hidden_size, dtype=hf_config.torch_dtype, device="cuda")
residual = torch.randn(1, hf_config.hidden_size, dtype=hf_config.torch_dtype, device="cuda")
# Per-layer outputs (hidden_states after each layer)
layer_outputs = torch.zeros(1, hf_config.hidden_size, dtype=hf_config.torch_dtype, device="cuda")
layer_residual = torch.zeros(1, hf_config.hidden_size, dtype=hf_config.torch_dtype, device="cuda")
self.offload_graphs = {}
self.offload_graph_pool = None
# Capture per-layer graphs
for layer_id in range(num_layers):
buffer_idx = layer_id % num_buffers
layer = self.model.model.layers[layer_id]
attn_module = layer.self_attn.attn
# Set Attention cache to ring buffer (fixed address for this layer)
attn_module.k_cache = offload_engine.layer_k_cache[buffer_idx:buffer_idx+1]
attn_module.v_cache = offload_engine.layer_v_cache[buffer_idx:buffer_idx+1]
# Set context for contiguous mode (no block tables)
set_context(
is_prefill=False,
slot_mapping=slot_mapping,
context_lens=context_lens,
block_tables=None,
)
# Warmup run - execute layer and propagate state
out_h, out_r = layer(positions, hidden_states, residual)
layer_outputs.copy_(out_h)
layer_residual.copy_(out_r)
torch.cuda.synchronize()
# Capture graph - use same input/output tensors
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph, self.offload_graph_pool):
out_h, out_r = layer(positions, hidden_states, residual)
layer_outputs.copy_(out_h)
layer_residual.copy_(out_r)
if self.offload_graph_pool is None:
self.offload_graph_pool = graph.pool()
self.offload_graphs[layer_id] = graph
reset_context()
# Update hidden_states and residual for next layer's capture
# This ensures subsequent layers see realistic input distributions
hidden_states.copy_(layer_outputs)
residual.copy_(layer_residual)
# Store graph variables for replay
self.offload_graph_vars = dict(
input_ids=input_ids,
positions=positions,
slot_mapping=slot_mapping,
context_lens=context_lens,
hidden_states=hidden_states,
residual=residual,
layer_outputs=layer_outputs,
layer_residual=layer_residual,
)
logger.info(f"Captured {num_layers} offload CUDA graphs")

113
task_plan.md
View File

@@ -1,8 +1,25 @@
# Task Plan: Enable CUDA Graphs for CPU Offload Mode
## Current Status
## Current Status: ✅ COMPLETED
### Completed: Refactor Offload Decode to Use Standard Attention Path
### Phase 0 Completed: Refactor Offload Decode to Use Standard Attention Path
### Phases 1-3 Completed: CUDA Graph Support for Offload Mode
**Implementation**: Added per-layer CUDA graph capture and replay for offload decode path.
**Key Changes**:
1. `capture_offload_cudagraph()` captures one graph per transformer layer
2. Each graph uses the corresponding ring buffer slot based on `layer_id % num_buffers`
3. `run_layerwise_offload_decode()` replays graphs when `enforce_eager=False`
4. Synchronization added between graph replays to ensure correct data flow
**Test Results**:
- `test_needle.py --input-len 32768 --enable-offload --use-cuda-graph`: **PASSED**
---
### Previous Work: Refactor Offload Decode to Use Standard Attention Path
**Problem solved**: The original offload decode (`run_layerwise_offload_decode`) bypassed `Attention.forward()` by manually calling attention components. This was inconsistent with the standard execution path.
@@ -179,9 +196,9 @@ Instead of per-layer graphs, capture entire decode step:
| Phase | Description | Status |
|-------|-------------|--------|
| Phase 0 | Refactor offload decode to use Attention.forward() | ✅ Completed |
| Phase 1 | Implement `capture_offload_cudagraph()` with per-buffer graphs | ⬜ Pending |
| Phase 2 | Modify `run_layerwise_offload_decode()` to use graphs | ⬜ Pending |
| Phase 3 | Test and benchmark | ⬜ Pending |
| Phase 1 | Implement `capture_offload_cudagraph()` with per-layer graphs | ✅ Completed |
| Phase 2 | Modify `run_layerwise_offload_decode()` to use graphs | ✅ Completed |
| Phase 3 | Test and benchmark | ✅ Completed |
| Phase 4 | (Optional) Optimize to full-decode graph | ⬜ Future |
## Architecture After Refactoring
@@ -212,12 +229,86 @@ Instead of per-layer graphs, capture entire decode step:
| File | Changes |
|------|---------|
| `model_runner.py:46-57` | Conditional CUDA graph capture (skip for offload) |
| `model_runner.py:841-991` | Refactored `run_layerwise_offload_decode()` to use standard `layer.forward()` |
| `model_runner.py:46-50` | Conditional CUDA graph capture: calls `capture_offload_cudagraph()` for offload mode |
| `model_runner.py:69-73` | Updated `exit()` to clean up offload graph resources |
| `model_runner.py:844-1031` | Refactored `run_layerwise_offload_decode()` to use standard `layer.forward()` with optional CUDA graph |
| `model_runner.py:1075-1164` | New `capture_offload_cudagraph()` method for per-layer graph capture |
| `tests/test_needle.py` | Added `--use-cuda-graph` flag to test CUDA graph mode |
## Implementation Details
### `capture_offload_cudagraph()` (line 1075-1164)
Captures per-layer CUDA graphs for offload decode:
```python
def capture_offload_cudagraph(self):
# Fixed-address tensors for graph capture
hidden_states = torch.randn(1, hidden_size, ...)
residual = torch.randn(1, hidden_size, ...)
layer_outputs = torch.zeros(1, hidden_size, ...)
layer_residual = torch.zeros(1, hidden_size, ...)
for layer_id in range(num_layers):
buffer_idx = layer_id % num_buffers
# Set Attention cache to ring buffer
attn_module.k_cache = ring_buffer[buffer_idx:buffer_idx+1]
attn_module.v_cache = ring_buffer[buffer_idx:buffer_idx+1]
# Warmup and capture
with torch.cuda.graph(graph):
out_h, out_r = layer(positions, hidden_states, residual)
layer_outputs.copy_(out_h)
layer_residual.copy_(out_r)
# Update inputs for next layer
hidden_states.copy_(layer_outputs)
residual.copy_(layer_residual)
```
### `run_layerwise_offload_decode()` CUDA Graph Mode
When CUDA graphs are available:
```python
use_cuda_graph = not self.enforce_eager and hasattr(self, 'offload_graphs')
if use_cuda_graph:
# Use fixed-address tensors
graph_vars["positions"][0] = len(seq) - 1
graph_vars["slot_mapping"][0] = context_len
graph_vars["context_lens"][0] = context_len + 1
graph_vars["hidden_states"].copy_(embedding)
graph_vars["residual"].zero_()
for layer_id in range(num_layers):
# Set up ring buffer and context
...
# Replay graph
self.offload_graphs[layer_id].replay()
torch.cuda.current_stream().synchronize()
# Copy outputs to inputs for next layer
if layer_id < num_layers - 1:
graph_vars["hidden_states"].copy_(graph_vars["layer_outputs"])
graph_vars["residual"].copy_(graph_vars["layer_residual"])
```
## Test Results
| Test | Mode | CUDA Graph | Status |
|------|------|------------|--------|
| `test_needle.py --input-len 4096` | GPU-only | N/A | PASSED |
| `test_needle.py --input-len 4096 --enable-offload` | CPU offload | Disabled | PASSED |
| `test_needle.py --input-len 32768 --enable-offload` | CPU offload | Disabled | PASSED |
| `test_needle.py --input-len 32768 --enable-offload --use-cuda-graph` | CPU offload | Enabled | PASSED |
## Next Steps
1. Implement `capture_offload_cudagraph()` method
2. Modify `run_layerwise_offload_decode()` to optionally use captured graphs
3. Benchmark performance improvement from CUDA graphs
4. Consider full-decode graph optimization for maximum performance
1. ~~Implement `capture_offload_cudagraph()` method~~ ✅
2. ~~Modify `run_layerwise_offload_decode()` to optionally use captured graphs~~ ✅
3. ~~Test correctness with needle-in-haystack~~ ✅
4. Benchmark performance improvement from CUDA graphs (optional)
5. Consider full-decode graph optimization for maximum performance (future)

18
tests/test_needle.py
View File

@@ -38,6 +38,7 @@ def run_needle_test(
minference_vertical: int = 1000,
minference_slash: int = 6096,
gpu_utilization: float = 0.9,
enforce_eager: bool = True,
verbose: bool = True,
) -> bool:
"""
@@ -97,7 +98,7 @@ def run_needle_test(
# 1. Initialize LLM
llm_kwargs = {
"enforce_eager": True,
"enforce_eager": enforce_eager,
"max_model_len": max_model_len,
"max_num_batched_tokens": max_model_len,
"enable_cpu_offload": enable_cpu_offload,
@@ -259,11 +260,25 @@ if __name__ == "__main__":
default=0.9,
help="GPU memory utilization (default: 0.9)"
)
parser.add_argument(
"--enforce-eager",
action="store_true",
default=True,
help="Force eager execution (disable CUDA graphs)"
)
parser.add_argument(
"--use-cuda-graph",
action="store_true",
help="Enable CUDA graph (disable enforce_eager)"
)
args = parser.parse_args()
# Convert budget=0 to None for fixed mode
minference_budget = args.minference_budget if args.minference_budget > 0 else None
# Determine enforce_eager: use_cuda_graph overrides enforce_eager
enforce_eager = not args.use_cuda_graph
passed = run_needle_test(
model_path=args.model,
max_model_len=args.max_model_len,
@@ -282,6 +297,7 @@ if __name__ == "__main__":
minference_vertical=args.minference_vertical,
minference_slash=args.minference_slash,
gpu_utilization=args.gpu_utilization,
enforce_eager=enforce_eager,
verbose=True,
)