✅ test: add comprehensive RULER benchmark test suite

- Add test_ruler.py supporting all 13 RULER tasks (NIAH, QA, CWE, FWE, VT) - Implement RULER official evaluation metrics (string_match_all/part) - Fix max_model_len to 32896 to prevent decode OOM on long inputs - Add ruler_benchmark_report.md with full test results (92.1% accuracy) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
[claudesquad] update from 'multi-request-2' on 13 Jan 26 02:01 CST
2026-01-14 00:51:30 +08:00 · 2026-01-13 02:01:07 +08:00 · 2026-01-12 21:08:35 +08:00 · 2026-01-12 21:08:27 +08:00 · 2026-01-12 16:27:25 +08:00 · 2026-01-12 15:16:39 +08:00
50 changed files with 8559 additions and 2481 deletions
--- a/.claude/commands/commit.md
+++ b/.claude/commands/commit.md
@@ -0,0 +1,166 @@
 ---
 allowed-tools: Bash(git add:*), Bash(git status:*), Bash(git commit:*), Bash(git diff:*), Bash(git log:*)
 argument-hint: [message] | --no-verify | --amend
 description: Create well-formatted commits with conventional commit format and emoji
 ---
 # Smart Git Commit
 Create well-formatted commit: $ARGUMENTS
 ## Current Repository State
 - Git status: !`git status --porcelain`
 - Current branch: !`git branch --show-current`
 - Staged changes: !`git diff --cached --stat`
 - Unstaged changes: !`git diff --stat`
 - Recent commits: !`git log --oneline -5`
 ## What This Command Does
 1. Unless specified with `--no-verify`, automatically runs pre-commit checks:
   - `pnpm lint` to ensure code quality
   - `pnpm build` to verify the build succeeds
   - `pnpm generate:docs` to update documentation
 2. Checks which files are staged with `git status`
 3. If 0 files are staged, automatically adds all modified and new files with `git add`
 4. Performs a `git diff` to understand what changes are being committed
 5. Analyzes the diff to determine if multiple distinct logical changes are present
 6. If multiple distinct changes are detected, suggests breaking the commit into multiple smaller commits
 7. For each commit (or the single commit if not split), creates a commit message using emoji conventional commit format
 ## Best Practices for Commits
 - **Verify before committing**: Ensure code is linted, builds correctly, and documentation is updated
 - **Atomic commits**: Each commit should contain related changes that serve a single purpose
 - **Split large changes**: If changes touch multiple concerns, split them into separate commits
 - **Conventional commit format**: Use the format `<type>: <description>` where type is one of:
  - `feat`: A new feature
  - `fix`: A bug fix
  - `docs`: Documentation changes
  - `style`: Code style changes (formatting, etc)
  - `refactor`: Code changes that neither fix bugs nor add features
  - `perf`: Performance improvements
  - `test`: Adding or fixing tests
  - `chore`: Changes to the build process, tools, etc.
 - **Present tense, imperative mood**: Write commit messages as commands (e.g., "add feature" not "added feature")
 - **Concise first line**: Keep the first line under 72 characters
 - **Emoji**: Each commit type is paired with an appropriate emoji:
  - ✨ `feat`: New feature
  - 🐛 `fix`: Bug fix
  - 📝 `docs`: Documentation
  - 💄 `style`: Formatting/style
  - ♻️ `refactor`: Code refactoring
  - ⚡️ `perf`: Performance improvements
  - ✅ `test`: Tests
  - 🔧 `chore`: Tooling, configuration
  - 🚀 `ci`: CI/CD improvements
  - 🗑️ `revert`: Reverting changes
  - 🧪 `test`: Add a failing test
  - 🚨 `fix`: Fix compiler/linter warnings
  - 🔒️ `fix`: Fix security issues
  - 👥 `chore`: Add or update contributors
  - 🚚 `refactor`: Move or rename resources
  - 🏗️ `refactor`: Make architectural changes
  - 🔀 `chore`: Merge branches
  - 📦️ `chore`: Add or update compiled files or packages
  - ➕ `chore`: Add a dependency
  - ➖ `chore`: Remove a dependency
  - 🌱 `chore`: Add or update seed files
  - 🧑‍💻 `chore`: Improve developer experience
  - 🧵 `feat`: Add or update code related to multithreading or concurrency
  - 🔍️ `feat`: Improve SEO
  - 🏷️ `feat`: Add or update types
  - 💬 `feat`: Add or update text and literals
  - 🌐 `feat`: Internationalization and localization
  - 👔 `feat`: Add or update business logic
  - 📱 `feat`: Work on responsive design
  - 🚸 `feat`: Improve user experience / usability
  - 🩹 `fix`: Simple fix for a non-critical issue
  - 🥅 `fix`: Catch errors
  - 👽️ `fix`: Update code due to external API changes
  - 🔥 `fix`: Remove code or files
  - 🎨 `style`: Improve structure/format of the code
  - 🚑️ `fix`: Critical hotfix
  - 🎉 `chore`: Begin a project
  - 🔖 `chore`: Release/Version tags
  - 🚧 `wip`: Work in progress
  - 💚 `fix`: Fix CI build
  - 📌 `chore`: Pin dependencies to specific versions
  - 👷 `ci`: Add or update CI build system
  - 📈 `feat`: Add or update analytics or tracking code
  - ✏️ `fix`: Fix typos
  - ⏪️ `revert`: Revert changes
  - 📄 `chore`: Add or update license
  - 💥 `feat`: Introduce breaking changes
  - 🍱 `assets`: Add or update assets
  - ♿️ `feat`: Improve accessibility
  - 💡 `docs`: Add or update comments in source code
  - 🗃️ `db`: Perform database related changes
  - 🔊 `feat`: Add or update logs
  - 🔇 `fix`: Remove logs
  - 🤡 `test`: Mock things
  - 🥚 `feat`: Add or update an easter egg
  - 🙈 `chore`: Add or update .gitignore file
  - 📸 `test`: Add or update snapshots
  - ⚗️ `experiment`: Perform experiments
  - 🚩 `feat`: Add, update, or remove feature flags
  - 💫 `ui`: Add or update animations and transitions
  - ⚰️ `refactor`: Remove dead code
  - 🦺 `feat`: Add or update code related to validation
  - ✈️ `feat`: Improve offline support
 ## Guidelines for Splitting Commits
 When analyzing the diff, consider splitting commits based on these criteria:
 1. **Different concerns**: Changes to unrelated parts of the codebase
 2. **Different types of changes**: Mixing features, fixes, refactoring, etc.
 3. **File patterns**: Changes to different types of files (e.g., source code vs documentation)
 4. **Logical grouping**: Changes that would be easier to understand or review separately
 5. **Size**: Very large changes that would be clearer if broken down
 ## Examples
 Good commit messages:
 - ✨ feat: add user authentication system
 - 🐛 fix: resolve memory leak in rendering process
 - 📝 docs: update API documentation with new endpoints
 - ♻️ refactor: simplify error handling logic in parser
 - 🚨 fix: resolve linter warnings in component files
 - 🧑‍💻 chore: improve developer tooling setup process
 - 👔 feat: implement business logic for transaction validation
 - 🩹 fix: address minor styling inconsistency in header
 - 🚑️ fix: patch critical security vulnerability in auth flow
 - 🎨 style: reorganize component structure for better readability
 - 🔥 fix: remove deprecated legacy code
 - 🦺 feat: add input validation for user registration form
 - 💚 fix: resolve failing CI pipeline tests
 - 📈 feat: implement analytics tracking for user engagement
 - 🔒️ fix: strengthen authentication password requirements
 - ♿️ feat: improve form accessibility for screen readers
 Example of splitting commits:
 - First commit: ✨ feat: add new solc version type definitions
 - Second commit: 📝 docs: update documentation for new solc versions
 - Third commit: 🔧 chore: update package.json dependencies
 - Fourth commit: 🏷️ feat: add type definitions for new API endpoints
 - Fifth commit: 🧵 feat: improve concurrency handling in worker threads
 - Sixth commit: 🚨 fix: resolve linting issues in new code
 - Seventh commit: ✅ test: add unit tests for new solc version features
 - Eighth commit: 🔒️ fix: update dependencies with security vulnerabilities
 ## Command Options
 - `--no-verify`: Skip running the pre-commit checks (lint, build, generate:docs)
 ## Important Notes
 - By default, pre-commit checks (`pnpm lint`, `pnpm build`, `pnpm generate:docs`) will run to ensure code quality
 - If these checks fail, you'll be asked if you want to proceed with the commit anyway or fix the issues first
 - If specific files are already staged, the command will only commit those files
 - If no files are staged, it will automatically stage all modified and new files
 - The commit message will be constructed based on the changes detected
 - Before committing, the command will review the diff to identify if multiple commits would be more appropriate
 - If suggesting multiple commits, it will help you stage and commit the changes separately
 - Always reviews the commit diff to ensure the message matches the changes
--- a/.claude/commands/create-architecture-documentation.md
+++ b/.claude/commands/create-architecture-documentation.md
@@ -0,0 +1,94 @@
 ---
 allowed-tools: Read, Write, Edit, Bash
 argument-hint: "[framework] | --c4-model | --arc42 | --adr | --plantuml | --full-suite"
 description: Generate comprehensive architecture documentation with diagrams, ADRs, and interactive visualization
 ---
 # Architecture Documentation Generator
 Generate comprehensive architecture documentation: $ARGUMENTS
 ## Current Architecture Context
 - Project structure: !`find . -type f -name "*.json" -o -name "*.yaml" -o -name "*.toml" | head -5`
 - Documentation exists: @docs/ or @README.md (if exists)
 - Architecture files: !`find . -name "*architecture*" -o -name "*design*" -o -name "*.puml" | head -3`
 - Services/containers: @docker-compose.yml or @k8s/ (if exists)
 - API definitions: !`find . -name "*api*" -o -name "*openapi*" -o -name "*swagger*" | head -3`
 ## Task
 Generate comprehensive architecture documentation with modern tooling and best practices:
 1. **Architecture Analysis and Discovery**
   - Analyze current system architecture and component relationships
   - Identify key architectural patterns and design decisions
   - Document system boundaries, interfaces, and dependencies
   - Assess data flow and communication patterns
   - Identify architectural debt and improvement opportunities
 2. **Architecture Documentation Framework**
   - Choose appropriate documentation framework and tools:
     - **C4 Model**: Context, Containers, Components, Code diagrams
     - **Arc42**: Comprehensive architecture documentation template
     - **Architecture Decision Records (ADRs)**: Decision documentation
     - **PlantUML/Mermaid**: Diagram-as-code documentation
     - **Structurizr**: C4 model tooling and visualization
     - **Draw.io/Lucidchart**: Visual diagramming tools
 3. **System Context Documentation**
   - Create high-level system context diagrams
   - Document external systems and integrations
   - Define system boundaries and responsibilities
   - Document user personas and stakeholders
   - Create system landscape and ecosystem overview
 4. **Container and Service Architecture**
   - Document container/service architecture and deployment view
   - Create service dependency maps and communication patterns
   - Document deployment architecture and infrastructure
   - Define service boundaries and API contracts
   - Document data persistence and storage architecture
 5. **Component and Module Documentation**
   - Create detailed component architecture diagrams
   - Document internal module structure and relationships
   - Define component responsibilities and interfaces
   - Document design patterns and architectural styles
   - Create code organization and package structure documentation
 6. **Data Architecture Documentation**
   - Document data models and database schemas
   - Create data flow diagrams and processing pipelines
   - Document data storage strategies and technologies
   - Define data governance and lifecycle management
   - Create data integration and synchronization documentation
 7. **Security and Compliance Architecture**
   - Document security architecture and threat model
   - Create authentication and authorization flow diagrams
   - Document compliance requirements and controls
   - Define security boundaries and trust zones
   - Create incident response and security monitoring documentation
 8. **Quality Attributes and Cross-Cutting Concerns**
   - Document performance characteristics and scalability patterns
   - Create reliability and availability architecture documentation
   - Document monitoring and observability architecture
   - Define maintainability and evolution strategies
   - Create disaster recovery and business continuity documentation
 9. **Architecture Decision Records (ADRs)**
   - Create comprehensive ADR template and process
   - Document historical architectural decisions and rationale
   - Create decision tracking and review process
   - Document trade-offs and alternatives considered
   - Set up ADR maintenance and evolution procedures
 10. **Documentation Automation and Maintenance**
    - Set up automated diagram generation from code annotations
    - Configure documentation pipeline and publishing automation
    - Set up documentation validation and consistency checking
    - Create documentation review and approval process
    - Train team on architecture documentation practices and tools
    - Set up documentation versioning and change management
--- a/.claude/commands/ultra-think.md
+++ b/.claude/commands/ultra-think.md
@@ -0,0 +1,158 @@
 ---
 description: Deep analysis and problem solving with multi-dimensional thinking
 argument-hint: [problem or question to analyze]
 ---
 # Deep Analysis and Problem Solving Mode
 Deep analysis and problem solving mode
 ## Instructions
 1. **Initialize Ultra Think Mode**
   - Acknowledge the request for enhanced analytical thinking
   - Set context for deep, systematic reasoning
   - Prepare to explore the problem space comprehensively
 2. **Parse the Problem or Question**
   - Extract the core challenge from: $ARGUMENTS
   - Identify all stakeholders and constraints
   - Recognize implicit requirements and hidden complexities
   - Question assumptions and surface unknowns
 3. **Multi-Dimensional Analysis**
   Approach the problem from multiple angles:
   ### Technical Perspective
   - Analyze technical feasibility and constraints
   - Consider scalability, performance, and maintainability
   - Evaluate security implications
   - Assess technical debt and future-proofing
   ### Business Perspective
   - Understand business value and ROI
   - Consider time-to-market pressures
   - Evaluate competitive advantages
   - Assess risk vs. reward trade-offs
   ### User Perspective
   - Analyze user needs and pain points
   - Consider usability and accessibility
   - Evaluate user experience implications
   - Think about edge cases and user journeys
   ### System Perspective
   - Consider system-wide impacts
   - Analyze integration points
   - Evaluate dependencies and coupling
   - Think about emergent behaviors
 4. **Generate Multiple Solutions**
   - Brainstorm at least 3-5 different approaches
   - For each approach, consider:
     - Pros and cons
     - Implementation complexity
     - Resource requirements
     - Potential risks
     - Long-term implications
   - Include both conventional and creative solutions
   - Consider hybrid approaches
 5. **Deep Dive Analysis**
   For the most promising solutions:
   - Create detailed implementation plans
   - Identify potential pitfalls and mitigation strategies
   - Consider phased approaches and MVPs
   - Analyze second and third-order effects
   - Think through failure modes and recovery
 6. **Cross-Domain Thinking**
   - Draw parallels from other industries or domains
   - Apply design patterns from different contexts
   - Consider biological or natural system analogies
   - Look for innovative combinations of existing solutions
 7. **Challenge and Refine**
   - Play devil's advocate with each solution
   - Identify weaknesses and blind spots
   - Consider "what if" scenarios
   - Stress-test assumptions
   - Look for unintended consequences
 8. **Synthesize Insights**
   - Combine insights from all perspectives
   - Identify key decision factors
   - Highlight critical trade-offs
   - Summarize innovative discoveries
   - Present a nuanced view of the problem space
 9. **Provide Structured Recommendations**
   Present findings in a clear structure:
   ```
   ## Problem Analysis
   - Core challenge
   - Key constraints
   - Critical success factors
   ## Solution Options
   ### Option 1: [Name]
   - Description
   - Pros/Cons
   - Implementation approach
   - Risk assessment
   ### Option 2: [Name]
   [Similar structure]
   ## Recommendation
   - Recommended approach
   - Rationale
   - Implementation roadmap
   - Success metrics
   - Risk mitigation plan
   ## Alternative Perspectives
   - Contrarian view
   - Future considerations
   - Areas for further research
   ```
 10. **Meta-Analysis**
    - Reflect on the thinking process itself
    - Identify areas of uncertainty
    - Acknowledge biases or limitations
    - Suggest additional expertise needed
    - Provide confidence levels for recommendations
 ## Usage Examples
 ```bash
 # Architectural decision
 /ultra-think Should we migrate to microservices or improve our monolith?
 # Complex problem solving
 /ultra-think How do we scale our system to handle 10x traffic while reducing costs?
 # Strategic planning
 /ultra-think What technology stack should we choose for our next-gen platform?
 # Design challenge
 /ultra-think How can we improve our API to be more developer-friendly while maintaining backward compatibility?
 ```
 ## Key Principles
 - **First Principles Thinking**: Break down to fundamental truths
 - **Systems Thinking**: Consider interconnections and feedback loops
 - **Probabilistic Thinking**: Work with uncertainties and ranges
 - **Inversion**: Consider what to avoid, not just what to do
 - **Second-Order Thinking**: Consider consequences of consequences
 ## Output Expectations
 - Comprehensive analysis (typically 2-4 pages of insights)
 - Multiple viable solutions with trade-offs
 - Clear reasoning chains
 - Acknowledgment of uncertainties
 - Actionable recommendations
 - Novel insights or perspectives
--- a/.claude/rules/commands.md
+++ b/.claude/rules/commands.md
@@ -1,20 +1,16 @@
 # Commands
-## Installation
+## Running (with PYTHONPATH)
-```bash
+For multi-instance development, use PYTHONPATH instead of pip install:
 pip install -e .
 ```
 ## Running
 ```bash
 # Run example
-python example.py
+PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python example.py
 # Run benchmarks
-python bench.py                    # Standard benchmark
+PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python bench.py
-python bench_offload.py            # CPU offload benchmark
+PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python bench_offload.py
 ```
 ## Config Defaults
--- a/.claude/rules/doc-management.md
+++ b/.claude/rules/doc-management.md
@@ -0,0 +1,105 @@
 # Documentation Management
 ## CLAUDE.md Content Policy
 **CLAUDE.md should only contain operational requirements:**
 - Environment setup (PYTHONPATH, GPU mutex)
 - Execution requirements (how to run tests/benchmarks)
 - Quick configuration reference
 - Documentation index (links to detailed docs)
 **Technical details should go to docs/:**
 - Architecture and design explanations
 - Implementation details and code flows
 - Debugging techniques
 - Memory analysis and profiling
 - Algorithm explanations
 ## When Adding New Technical Content
 Follow this workflow:
 ### Step 1: Analyze and Document
 If doing technical analysis (e.g., memory profiling):
 1. Calculate theoretical values using formulas
 2. Run actual tests to measure real values
 3. Compare theoretical vs actual (expect < 10% error for valid models)
 4. Document findings with both theory and empirical validation
 ### Step 2: Create/Update docs/
 Create a new doc or update existing one in `docs/`:
 ```
 docs/
 ├── architecture_guide.md      # Core components, design, flows
 ├── sparse_attention_guide.md  # Sparse attention methods
 ├── layerwise_offload_memory_analysis.md  # Memory analysis
 ├── debugging_guide.md         # Debugging techniques
 └── <new_topic>_guide.md       # New technical topic
 ```
 ### Step 3: Update CLAUDE.md Documentation Index
 Add entry to the Documentation Index table:
 ```markdown
 | Document | Purpose |
 |----------|---------|
 | [`docs/new_doc.md`](docs/new_doc.md) | Brief description |
 ```
 ### Step 4: Refactor if Needed
 If CLAUDE.md grows too large (> 150 lines), refactor:
 1. Identify technical details that can be moved
 2. Create appropriate doc in docs/
 3. Replace detailed content with reference link
 4. Keep only operational essentials in CLAUDE.md
 ## Documentation Structure Template
 For new technical docs:
 ```markdown
 # Topic Guide
 Brief overview of what this document covers.
 ## Section 1: Concepts
 - Key concepts and terminology
 ## Section 2: Implementation
 - Code locations
 - Key methods/functions
 ## Section 3: Details
 - Detailed explanations
 - Code examples
 ## Section 4: Validation (if applicable)
 - Theoretical analysis
 - Empirical measurements
 - Comparison table
 ```
 ## Memory Analysis Template
 When documenting memory behavior:
 ```markdown
 ## Theoretical Calculation
 | Component | Formula | Size |
 |-----------|---------|------|
 | Buffer X | `param1 × param2 × dtype_size` | X MB |
 ## Empirical Validation
 | Metric | Theoretical | Actual | Error |
 |--------|-------------|--------|-------|
 | Peak memory | X GB | Y GB | Z% |
 ## Key Findings
 1. Finding 1
 2. Finding 2
 ```
--- a/.claude/rules/no-extra-docs.md
+++ b/.claude/rules/no-extra-docs.md
@@ -2,39 +2,47 @@
 ## Do Not Create Unnecessary Documentation
-**IMPORTANT**: Do NOT create extra markdown documentation files unless explicitly requested by the user.
+**IMPORTANT**: Do NOT create extra markdown documentation files proactively unless:
 1. User explicitly requests documentation
 2. Refactoring CLAUDE.md to move technical details to docs/ (see `doc-management.md`)
 ### What NOT to do:
- ❌ Do NOT create README files proactively
+- Do NOT create README files proactively
- ❌ Do NOT create analysis documents (*.md) after completing tasks
+- Do NOT create standalone analysis documents after completing tasks
- ❌ Do NOT create tutorial/guide documents
+- Do NOT create summary documents without request
 - ❌ Do NOT create summary documents
 ### What TO do:
- ✅ Only create documentation when user explicitly asks for it
+- Provide information directly in conversation by default
- ✅ Provide information directly in conversation instead
+- When user requests documentation, follow `doc-management.md` workflow
- ✅ Update existing documentation if changes require it
+- Update existing docs in `docs/` when code changes affect them
- ✅ Add inline code comments where necessary
+- Keep CLAUDE.md concise (< 150 lines), move technical details to docs/
-### Exceptions:
+### Documentation Locations:
-Documentation is acceptable ONLY when:
+| Type | Location |
-1. User explicitly requests "create a README" or "write documentation"
+|------|----------|
-2. Updating existing documentation to reflect code changes
+| Operational requirements | CLAUDE.md |
-3. Adding inline comments/docstrings to code itself
+| Technical details | docs/*.md |
 | Code comments | Inline in source |
 ### Examples:
-**Bad** (Don't do this):
+**Proactive docs (Don't do)**:
 ```
 User: "Profile the code"
-Assistant: [Creates profiling_results.md after profiling]
+Assistant: [Creates profiling_results.md without being asked]
 ```
-**Good** (Do this instead):
+**On-request docs (Do this)**:
 ```
-User: "Profile the code"
+User: "Profile the code and document the findings"
-Assistant: [Runs profiling, shows results in conversation]
+Assistant: [Runs profiling, creates/updates docs/memory_analysis.md]
 ```
 **Refactoring (Do this)**:
 ```
 User: "CLAUDE.md is too long, refactor it"
 Assistant: [Moves technical sections to docs/, updates CLAUDE.md index]
 ```
--- a/.claude/rules/testing.md
+++ b/.claude/rules/testing.md
@@ -66,33 +66,27 @@ print("test_xxx: PASSED")
 ## Running Tests
 Use PYTHONPATH for multi-instance isolation (no pip install needed):
 ```bash
 # Run a specific test
-python tests/test_offload_engine.py
+PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python tests/test_offload_engine.py
 # Run with specific GPU
-CUDA_VISIBLE_DEVICES=0 python tests/test_ring_buffer.py
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python tests/test_ring_buffer.py
 ```
 ## Benchmarks
 ```bash
-# Standard GPU benchmark
+PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python bench.py
-python bench.py
+PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python bench_offload.py
-
+PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python bench_vllm.py
 # CPU offload benchmark
 python bench_offload.py
 # vLLM comparison benchmark
 python bench_vllm.py
 ```
 ## Quick Verification
 ```bash
 # Import test
-python -c "from nanovllm import LLM"
+PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python -c "from nanovllm import LLM"
 # Run offload benchmark (tests CPU-primary ring buffer mode)
 python bench_offload.py
 ```
--- a/.gitignore
+++ b/.gitignore
@@ -197,3 +197,36 @@ cython_debug/
 results/
 outputs/
 .local/
 # Claude Flow generated files
 .claude/settings.local.json
 .mcp.json
 claude-flow.config.json
 .swarm/
 .hive-mind/
 .claude-flow/
 memory/
 coordination/
 memory/claude-flow-data.json
 memory/sessions/*
 !memory/sessions/README.md
 memory/agents/*
 !memory/agents/README.md
 coordination/memory_bank/*
 coordination/subtasks/*
 coordination/orchestration/*
 *.db
 *.db-journal
 *.db-wal
 *.sqlite
 *.sqlite-journal
 *.sqlite-wal
 claude-flow
 # Removed Windows wrapper files per user request
 hive-mind-prompt-*.txt
 # Test data
 tests/data/
 # Serena MCP tool config
 .serena/
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -4,444 +4,73 @@ This file provides guidance to Claude Code when working with this repository.
 ## Overview
-Nano-vLLM is a lightweight vLLM implementation (~1,200 lines) for fast offline LLM inference. Supports Qwen3 models with CPU offload for long-context inference.
+Nano-vLLM is a lightweight vLLM implementation (~1,200 lines) for fast offline LLM inference. Supports multiple model architectures (Qwen3, Qwen2, Llama) with CPU offload for long-context inference.
 ## GPU Mutex for Multi-Instance Debugging
-**IMPORTANT**: When running multiple Claude instances for parallel debugging, only one GPU (cuda:0) is available. Before executing ANY command that uses the GPU (python scripts, benchmarks, tests), Claude MUST:
+**IMPORTANT**: When running multiple Claude instances for parallel debugging, different rules apply based on script type:
-1. **Check GPU availability** by running:
+### Benchmarks (`bench*.py`) - Exclusive GPU Access Required
   ```bash
   nvidia-smi --query-compute-apps=pid,name,used_memory --format=csv,noheader
   ```
-2. **If processes are running on GPU**:
+Before running any `bench*.py` script, Claude MUST wait for exclusive GPU access:
   - Wait and retry every 10 seconds until GPU is free
   - Use this polling loop:
     ```bash
     while [ -n "$(nvidia-smi --query-compute-apps=pid --format=csv,noheader)" ]; do
       echo "GPU busy, waiting 10s..."
       sleep 10
     done
     ```
 3. **Only proceed** when `nvidia-smi --query-compute-apps=pid --format=csv,noheader` returns empty output
 **Example workflow**:
 ```bash
 # First check if GPU is in use
 nvidia-smi --query-compute-apps=pid,name,used_memory --format=csv,noheader
 # If output is empty, proceed with your command
 python bench_offload.py
 # If output shows processes, wait until they finish
 ```
 **Note**: This applies to ALL GPU operations including:
 - Running tests (`python tests/test_*.py`)
 - Running benchmarks (`python bench*.py`)
 - Running examples (`python example.py`)
 - Any script that imports torch/cuda
 ## Local Package Installation for Multi-Instance
 **CRITICAL**: After ANY code modification in the `nanovllm/` directory, you MUST reinstall the package before running tests or benchmarks:
 ```bash
-pip install -e . --prefix=./.local --no-deps
+# Check and wait for GPU to be free
 while [ -n "$(nvidia-smi --query-compute-apps=pid --format=csv,noheader)" ]; do
  echo "GPU busy, waiting 10s..."
  sleep 10
 done
 ```
-Then run with PYTHONPATH:
+### Other Scripts (tests, examples) - No Special Requirements
 For non-benchmark scripts, exclusive GPU access is NOT required. Multiple nanovllm processes can run simultaneously on different GPUs - each process automatically selects a unique port for `torch.distributed` communication.
 ## Multi-Instance Development with PYTHONPATH
 **IMPORTANT**: When running multiple Claude instances on different worktrees, do NOT use `pip install -e .` globally as it will affect other instances.
 **Use PYTHONPATH directly** - no pip install needed:
 ```bash
-PYTHONPATH=./.local/lib/python3.10/site-packages:$PYTHONPATH python <script.py>
+# Set PYTHONPATH to point to the project root directory
 PYTHONPATH=/path/to/your/worktree:$PYTHONPATH python <script.py>
 # Example: running tests
 PYTHONPATH=/home/zijie/Code/nano-vllm:$PYTHONPATH python tests/test_needle.py
 ```
-**IMPORTANT**: When running multiple Claude instances on different worktrees, do NOT use `pip install -e .` globally as it will affect other instances. Instead, use local installation:
+**Benefits**:
-
+- No `pip install` required
-1. **Install to worktree-local directory**:
+- Code changes take effect immediately (no reinstall needed)
-   ```bash
+- Each worktree is completely isolated
-   pip install -e . --prefix=./.local --no-deps
+
-   ```
+## Documentation Index
-
+
-2. **Set PYTHONPATH before running any Python command**:
+| Document | Purpose |
-   ```bash
+|----------|---------|
-   export PYTHONPATH=./.local/lib/python3.10/site-packages:$PYTHONPATH
+| [`docs/architecture_guide.md`](docs/architecture_guide.md) | Core components, layer-wise CPU offload design, prefill/decode flows, implementation details |
-   ```
+| [`docs/multi_model_support.md`](docs/multi_model_support.md) | Model registry system, adding new models (Qwen3/Llama), architecture differences, RoPE scaling |
-
+| [`docs/cuda_graph_offload_guide.md`](docs/cuda_graph_offload_guide.md) | CUDA graph support for CPU offload decode path, 4x decode speedup |
-3. **Combined example**:
+| [`docs/sparse_attention_guide.md`](docs/sparse_attention_guide.md) | Block sparse attention methods (MInference, FlexPrefill, XAttention, Quest), computation flow |
-   ```bash
+| [`docs/sparse_prefill_integration_plan.md`](docs/sparse_prefill_integration_plan.md) | Integration plan for MInference/XAttention/FlexPrefill with unified BlockMask interface |
-   # One-liner for running tests with local package
+| [`docs/sparse_offload_integration.md`](docs/sparse_offload_integration.md) | Sparse policy integration with layerwise offload, `requires_block_selection` interface design |
-   PYTHONPATH=./.local/lib/python3.10/site-packages:$PYTHONPATH python tests/test_needle.py
+| [`docs/layerwise_offload_memory_analysis.md`](docs/layerwise_offload_memory_analysis.md) | Memory allocation analysis with theoretical formulas and empirical validation (< 5% error) |
-   ```
+| [`docs/debugging_guide.md`](docs/debugging_guide.md) | PyTorch hooks for debugging, tensor comparison, memory profiling |
-
+| [`docs/gpu_only_performance_issue.md`](docs/gpu_only_performance_issue.md) | GPU-only mode slower than offload due to PagedAttention scatter overhead, optimization proposals |
-**Note**: The Python version in the path (python3.10) should match your environment.
+| [`docs/offload_accuracy_issue.md`](docs/offload_accuracy_issue.md) | **BUG**: CPU offload mode 66% accuracy vs 100% non-offload on RULER NIAH benchmark |
 **CRITICAL**: After making code changes to `nanovllm/` source files, you MUST reinstall the package for changes to take effect:
 ```bash
 pip install -e . --prefix=./.local --no-deps
 ```
 Without reinstallation, Python will use the old cached version and your changes will NOT be reflected!
 ## Sparse Attention
 For sparse attention related content (block sparse attention, MInference, FlexPrefill, XAttention, AvgPool, etc.), refer to [`docs/sparse_attention_guide.md`](docs/sparse_attention_guide.md).
 ### Quest Sparse Policy
 **Files**: `nanovllm/kvcache/sparse/quest.py`, `nanovllm/kvcache/sparse/policy.py`
 Quest policy selects Top-K blocks based on query-key similarity bounds using min/max key metadata.
 **Scoring Mechanism**:
 ```python
 score_min = torch.einsum('hd,bhd->bh', q, key_min)  # [num_blocks, kv_heads]
 score_max = torch.einsum('hd,bhd->bh', q, key_max)  # [num_blocks, kv_heads]
 scores = torch.maximum(score_min, score_max).mean(dim=-1)  # [num_blocks] ← averaged!
 ```
 **Critical Limitation - No Per-Head Scheduling**:
 The `.mean(dim=-1)` averages scores across all heads, making a **unified** block selection for all heads:
 ```
 Block A: head0 needs (+4), head1 doesn't (-4) → avg = 0 → NOT selected
 Block B: head0 doesn't (-4), head1 needs (+4) → avg = 0 → NOT selected
 Block C: both heads moderately need (+2, +2) → avg = +2 → selected
 ```
 **Why Per-Head Scheduling is Infeasible**:
 1. **Memory Layout**: GPU cache stores all heads together `[block_size, kv_heads, head_dim]`
 2. **FlashAttention**: Requires complete heads - partial heads cause dimension mismatch
 3. **Block Granularity**: If any head needs a block, the entire block (all heads) must be loaded
 **Policy Types**:
 - `FullAttentionPolicy`: `supports_prefill=True, supports_decode=True` - loads all blocks
 - `QuestPolicy`: `supports_prefill=False, supports_decode=True` - decode-only Top-K selection
 ## Architecture
 ### Core Components
 - **LLMEngine** (`llm_engine.py`): Main entry, runs prefill-decode loop
 - **ModelRunner** (`model_runner.py`): Loads weights, allocates KV cache, CUDA graphs
 - **Scheduler** (`scheduler.py`): Two-phase scheduling (prefill → decode)
 - **BlockManager** (`block_manager.py`): Paged attention with prefix caching (xxhash), default block size 4096
 - **Attention** (`layers/attention.py`): FlashAttention with chunked methods for CPU offload
 ## PyTorch Hooks for Debugging
 ### Hook Positions in Qwen3
 ```
 decoder_layer
 ├── input_layernorm (RMSNorm)
 ├── self_attn (Qwen3Attention)          ← Hook here for attention I/O after o_proj
 │   ├── q_proj → q_norm → RoPE
 │   ├── k_proj → k_norm → RoPE
 │   ├── v_proj
 │   ├── attn (Attention)                ← Hook here for Q/K/V tensors
 │   │   └── FlashAttention / SDPA
 │   └── o_proj
 ├── post_attention_layernorm (RMSNorm)
 └── mlp (Qwen3MLP)
 ```
 ### Hook Types & Data Shapes
 | Hook Position | Type | Captured Data |
 |---------------|------|---------------|
 | `self_attn` | post | `[batch, seq_len, hidden_size]` - after o_proj |
 | `self_attn.attn` | pre | Q,K,V: `[seq_len, num_heads, head_dim]` - after RoPE |
 | `self_attn.attn` | post | `[seq_len, num_heads, head_dim]` - before o_proj |
 ### Example: Capture Attention Outputs
 ```python
 storage = {}
 def make_hook(layer_id: int, storage: dict):
    def hook(module, inputs, output):
        if isinstance(output, tuple):
            attn_output = output[0]
        else:
            attn_output = output
        # nanovllm shape: [num_tokens, hidden_size] -> add batch dim
        if attn_output.dim() == 2:
            attn_output = attn_output.unsqueeze(0)
        storage[layer_id] = attn_output.detach().clone()
    return hook
 # Register hooks
 hooks = []
 for layer_idx, layer in enumerate(model.model.layers):
    hooks.append(layer.self_attn.register_forward_hook(make_hook(layer_idx, storage)))
 # Run inference...
 # Cleanup
 for hook in hooks:
    hook.remove()
 ```
 ### Reference Implementation
 Key files:
 - `tests/modeling_qwen3.py`: Reference Qwen3 implementation (torch + transformers only)
 - `tests/test_needle_ref.py`: Reference needle test using custom Qwen3
 - `tests/test_needle.py`: Needle-in-haystack test for nanovllm
 ### Common Pitfalls
 1. **Shape mismatch**: nanovllm uses `[num_tokens, ...]` while torch uses `[batch, seq_len, ...]`
 2. **Hook position**: `self_attn` captures after o_proj, `self_attn.attn` captures before o_proj
 3. **Output format**: nanovllm returns tuple `(attn_output, None)`, handle with `output[0]`
 ## CPU Offload System
 ### Ring Buffer Design
 ```
 GPU Slots: [0]  [1]  [2]  [3]  ...  (unified ring buffer)
 Prefill: slot = chunk_idx % N
 Decode:  slot[0] = decode, slots[1:] = load previous chunks
 ```
 **Key Files**: `kvcache/offload_engine.py`, `kvcache/hybrid_manager.py`
 **Memory Layout**:
 - GPU: `[num_layers, num_gpu_blocks, block_size, kv_heads, head_dim]`
 - CPU: `[num_layers, num_cpu_blocks, ...]` (pinned memory)
 **Key Methods**:
 - `load_to_slot_layer(slot, layer, cpu_block)`: Async H2D load
 - `offload_slot_to_cpu(slot, cpu_block)`: Async D2H offload
 - Per-slot per-layer CUDA events for fine-grained synchronization
 **Pipeline**: N-way pipeline with dedicated streams for full compute-transfer overlap. Pipeline depth = N-1 (prefill), (N-1)/2 (decode).
 ### Stream Architecture
 ```
 Transfer Streams: [slot_0_stream] [slot_1_stream] ... [slot_N_stream]
                       ↓              ↓                    ↓
 GPU Slots:          [slot_0]      [slot_1]    ...     [slot_N]
                       ↓              ↓                    ↓
 Compute Stream:    ←←←←←←←←←←←← [dedicated compute stream] →→→→→→→→→→→→
 ```
 **Key Design Decisions**:
 - **Per-slot transfer streams**: Each GPU slot has its own CUDA stream for H2D transfers, enabling parallel loading
 - **Dedicated compute stream**: Created with `torch.cuda.Stream()` (NOT `current_stream()`) to avoid implicit synchronization with default stream
 - **CUDA Events**: `ring_slot_ready` (transfer complete), `ring_slot_compute_done` (safe to overwrite)
 ## Scatter-Gather DMA (sgDMA) - INTEGRATED ✓
 ### Problem & Solution
 **Problem**: Strided CPU cache access `k_cache_cpu[:, block_id]` caused slow Device→Pageable transfers at ~1.4 GB/s instead of optimal ~24 GB/s pinned memory bandwidth.
 **Solution**: Implemented `cudaMemcpy2D` via custom CUDA extension to handle strided layouts natively. **Integration complete** as of 2025-12-25.
 ### Quick Start
 ```python
 from nanovllm.comm import memcpy_2d_async
 # Transfer block_id across all layers
 spitch = num_blocks * features * dtype_size  # stride between layers
 dpitch = features * dtype_size               # contiguous destination
 width = features * dtype_size                # bytes per row
 height = num_layers                          # number of rows
 memcpy_2d_async(gpu_buf, cpu_cache[:, block_id], dpitch, spitch, width, height, "h2d", stream)
 ```
 ### Benchmark Performance (Synthetic, 256MB)
 | Method | Bandwidth | Speedup |
 |--------|-----------|---------|
 | **cudaMemcpy2D (sgDMA)** | **24.95 GB/s** | **Baseline** |
 | PyTorch strided | 4.25 GB/s | **5.87x slower** |
 | PyTorch contiguous | 24.92 GB/s | Same |
 ### Real-World Performance (A100, Attention Offload)
 **Measured from `test_attention_offload.py` profiling**:
 | Transfer Type | Count | Bandwidth | Previous | Speedup |
 |---------------|-------|-----------|----------|---------|
 | **Device→Pinned (D2H)** | 416 | **21.49 GB/s** | 1.40 GB/s | **15.35x** |
 | **Pinned→Device (H2D)** | 24,960 | **23.39 GB/s** | N/A | N/A |
 | Device→Pageable (D2H) | **0** | N/A | ~40 transfers | **Eliminated** |
 **Verification**: All slow Device→Pageable transfers eliminated. System achieves near-optimal PCIe Gen3 x16 bandwidth.
 **Build**: `python setup.py build_ext --inplace`
 **Files**:
 - `csrc/sgdma_kernel.cu`, `csrc/sgdma.cpp`: CUDA extension
 - `nanovllm/comm/sgdma.py`: Python API
 - `kvcache/offload_engine.py`: Integration (4 methods updated)
 ### Integration Details
 **Modified methods in `offload_engine.py`**:
 - `load_to_slot_all_layers()`: H2D ring buffer load
 - `offload_slot_to_cpu()`: D2H ring buffer offload
 - `offload_decode_slot()`: D2H decode slot offload
 - `load_cpu_blocks_to_gpu_slots_all_layers()`: Batch H2D load
 **Example replacement**:
 ```python
 # Before (slow, Device→Pageable fallback)
 self.k_cache_gpu[:, slot].copy_(self.k_cache_cpu[:, cpu_block], non_blocking=True)
 # After (fast, Device→Pinned via sgDMA)
 memcpy_2d_async(
    self.k_cache_gpu[:, slot], self.k_cache_cpu[:, cpu_block],
    self.gpu_pitch, self.cpu_pitch, self.width, self.height,
    "h2d", stream=self.transfer_stream_main
 )
 ```
 **Actual Impact**: 15.35x faster D2H transfers, eliminates memory transfer bottleneck. Expected 2-3x overall prefill throughput improvement.
 ## Online Softmax Merge - Triton Fused Kernel ✓
 ### Problem & Solution
 **Problem**: Original PyTorch implementation of `merge_attention_outputs()` launches 7 separate kernels per merge operation:
 1. `torch.maximum()` - max(lse1, lse2)
 2. `torch.exp()` (2x) - exp(lse1-max), exp(lse2-max)
 3. `transpose()` + `unsqueeze()` - reshape for broadcasting
 4. Accumulation (6x) - weighted sum operations
 5. Division - normalize output
 6. `torch.log()` - merge LSE
 7. `.to()` - type conversion
 **Profiling revealed**: In ChunkedPrefill with 8 layers, these operations consumed **698 ms** GPU time (vs FlashAttention 603 ms), becoming a major bottleneck.
 **Solution**: Implemented Triton fused kernels that combine all operations into 2 kernels. **Integration complete** as of 2025-12-25.
 ### Implementation
 **File**: `nanovllm/kvcache/chunked_attention.py:278-408`
 Two Triton kernels replace all PyTorch operations:
 ```python
@triton.jit
 def _merge_lse_kernel(...):
    """Fused: max + exp + log"""
    max_lse = tl.maximum(lse1, lse2)
    exp1 = tl.exp(lse1 - max_lse)
    exp2 = tl.exp(lse2 - max_lse)
    lse_merged = max_lse + tl.log(exp1 + exp2)
    tl.store(lse_out_ptr + offsets, lse_merged, mask=mask)
@triton.jit
 def _merge_output_kernel(...):
    """Fused: broadcast + weighted sum + division"""
    # Load LSE, compute scaling factors
    exp1 = tl.exp(lse1 - max_lse)
    exp2 = tl.exp(lse2 - max_lse)
    sum_exp = exp1 + exp2
    # Process headdim in chunks
    for d_offset in range(0, headdim, BLOCK_SIZE):
        o1_val = tl.load(o1_ptr + o_idx, mask=mask)
        o2_val = tl.load(o2_ptr + o_idx, mask=mask)
        o_merged = (o1_val * exp1 + o2_val * exp2) / sum_exp
        tl.store(o_out_ptr + o_idx, o_merged, mask=mask)
 ```
 ### Performance Results
 **From `test_attention_offload.py` profiling** (8 layers, 16K tokens, 16 chunks, 10 iterations):
 | Metric | PyTorch (7 kernels) | Triton (2 kernels) | Speedup |
 |--------|---------------------|---------------------|---------|
 | **GPU time (8 layers)** | 698 ms | 160.7 ms | **4.3x** |
 | **Per-layer time** | 87.3 ms | 20.1 ms | **4.3x** |
 | **Avg per merge** | 56 µs | 12.9 µs | **4.3x** |
 | **Kernel launches** | 10,920 | 3,120 | **71% reduction** |
 **Breakdown** (per-layer, 1,560 merges):
 - `_merge_output_kernel`: 126.9 ms / 8 = 15.9 ms/layer (avg 10.2 µs/call)
 - `_merge_lse_kernel`: 33.8 ms / 8 = 4.2 ms/layer (avg 2.7 µs/call)
 ### Overall ChunkedPrefill Impact
 **GPU time distribution** (test_attention_offload.py):
 | Component | Time (ms) | Percentage |
 |-----------|-----------|------------|
 | FlashAttention | 603.2 | 74.8% |
 | Triton Merge | 160.7 | 19.9% |
 | Other | 42.1 | 5.3% |
 | **Total** | **806.0** | **100%** |
 **If using PyTorch merge** (estimated):
 - Total GPU time: ~1,343 ms
 - **Overall speedup with Triton**: 1.67x
 ### Key Files
 - `nanovllm/kvcache/chunked_attention.py`: Triton kernels + merge function
 ## Known Issues and Fixes
 ### Partial Last Block Bug (FIXED ✓)
 **Problem**: When prefill token count is not an exact multiple of `block_size`, decode outputs garbage.
 **Root Cause**: `_chunked_decode_attention` calculated `last_block_valid_tokens` using `len(seq) - 1`, which increases during decode. But CPU blocks are fixed after prefill!
 ```python
 # BUG: len(seq) increases each decode step
 total_prefill_tokens = len(seq) - 1  # Wrong!
 last_block_valid_tokens = total_prefill_tokens % block_size  # Reads garbage from CPU
 ```
 **Fix**: Cache original prefill length in `HybridKVCacheManager.get_prefill_len()`:
 ```python
 # CORRECT: Use cached prefill length
 total_prefill_tokens = kvcache_manager.get_prefill_len(seq)  # Fixed value
 ```
 **Files Modified**:
 - `nanovllm/kvcache/hybrid_manager.py`: Added `_prefill_len` dict and `get_prefill_len()` method
 - `nanovllm/layers/attention.py`: Use `get_prefill_len()` instead of `len(seq) - 1`
 ### Block Size 4096 Race Condition (FIXED ✓)
 **Problem**: `block_size=4096` with multiple chunks produced `index_copy_(): index out of bounds` CUDA error during Chunk 2 processing.
 **Root Cause**: Race condition between default stream and compute stream. In `_prepare_chunked_offload_chunk()`, `slot_mapping` tensor was created with `non_blocking=True` H2D transfer on the default stream. However, `store_kvcache` runs on `compute_stream`. Without synchronization, `compute_stream` could use `slot_mapping` before its transfer completed, causing corrupted indices.
 **Fix** (in `attention.py`):
 ```python
 if is_chunked_offload:
    compute_stream = context.kvcache_manager.offload_engine.compute_stream
    if k_cache.numel() and v_cache.numel():
        # CRITICAL: Wait for default stream to ensure slot_mapping tensor transfer is complete
        compute_stream.wait_stream(torch.cuda.default_stream())
        with torch.cuda.stream(compute_stream):
            store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
 ```
 **Tested block sizes**: 512, 1024, 4096, 8192 - all pass.
 ## Configuration
 | Parameter | Default | Notes |
 |-----------|---------|-------|
-| `kvcache_block_size` | 1024 | Tokens per block (4096 now works after race condition fix) |
+| `kvcache_block_size` | 4096 | Tokens per block |
 | `max_num_batched_tokens` | 16384 | Set = max_model_len for long context |
 | `gpu_memory_utilization` | 0.9 | GPU memory fraction |
 | `enable_cpu_offload` | False | Enable for long context |
 | `num_gpu_blocks` | 2 | GPU blocks for offload mode |
 | `num_kv_buffers` | 4 | Ring buffer size for decode pipeline |
 | `enforce_eager` | False | Set True to disable CUDA graphs |
 ## Benchmarking
@@ -455,58 +84,13 @@ if is_chunked_offload:
 **Model Limits**:
 - Qwen3-0.6B/4B: 40960 tokens
 - Qwen2.5-7B-Instruct-1M: 1048576 tokens
 - Llama-3.1-8B-Instruct: 131072 tokens
-**Performance (Qwen3-0.6B)**:
+**Performance (Qwen3-4B, CPU Offload)**:
- GPU: ~18k tok/s (prefill), ~100 tok/s (decode)
+- Prefill: ~5700-8000 tok/s (varies by context length)
- CPU Offload (16K): ~14k tok/s (prefill)
+- Decode with CUDA Graph: ~50 tok/s (TPOT ~19ms)
- CPU Offload (32K): ~13k tok/s (prefill)
+- Decode Eager Mode: ~12 tok/s (TPOT ~80ms)
-
+- **CUDA Graph speedup: 4x decode throughput**
 ## Performance Summary
 ### Completed Optimizations ✓
 1. **sgDMA Integration** (2025-12-25)
   - Eliminated Device→Pageable transfers
   - Achieved 21-23 GB/s bandwidth (near PCIe limit)
   - 15.35x speedup on memory transfers
 2. **Triton Fused Merge Kernel** (2025-12-25)
   - Reduced 7 PyTorch kernels → 2 Triton kernels
   - 4.3x speedup on merge operations
   - 1.67x overall ChunkedPrefill speedup
 3. **N-way Pipeline with Dedicated Streams** (2025-12-25)
   - Per-slot transfer streams for parallel H2D across slots
   - Dedicated compute stream (avoids CUDA default stream implicit sync)
   - N-way pipeline using all available slots (not just 2-slot double buffering)
   - **2.0x improvement**: 7.2k → 14.1k tok/s (16K tokens prefill)
 ### Current Performance Bottlenecks
 **From profiling** (`test_attention_offload.py`, 8 layers, 16K tokens):
 | Component | GPU Time | Percentage | Optimization Potential |
 |-----------|----------|------------|------------------------|
 | FlashAttention | 603 ms | 74.8% | ⚠️ Main bottleneck |
 | Triton Merge | 161 ms | 19.9% | ✓ Optimized |
 | Other | 42 ms | 5.3% | Minor |
 ### Future Optimization Directions
 1. **FlashAttention Optimization** (highest priority)
   - Current: 74.8% of GPU time
   - Potential: Custom FlashAttention kernel for chunked case
   - Expected: 1.5-2x additional speedup
 2. ~~**Pipeline Optimization**~~ ✓ COMPLETED
   - ~~Better overlap between compute and memory transfer~~
   - ~~Multi-stream execution~~
   - See: N-way Pipeline with Dedicated Streams above
 3. **Alternative to sgDMA** (lower priority, PyTorch-only)
   - Reorganize cache layout: `[num_cpu_blocks, num_layers, ...]` instead of `[num_layers, num_cpu_blocks, ...]`
   - Trade-off: Extensive refactoring vs minimal sgDMA approach
   - Same performance as sgDMA (~24 GB/s)
 ---
--- a/bench.py
+++ b/bench.py
@@ -2,6 +2,7 @@ import os
 import time
 from random import randint, seed
 from nanovllm import LLM, SamplingParams
 from nanovllm.config import SparsePolicyType
 def bench_decode(llm, num_seqs, input_len, output_len):
@@ -23,8 +24,8 @@ def bench_decode(llm, num_seqs, input_len, output_len):
    print(f"         Throughput: {decode_throughput:.2f} tok/s (includes prefill overhead)")
-def bench_prefill(llm, num_seqs, input_len):
+def bench_prefill(llm, num_seqs, input_len, label=""):
-    """Benchmark prefill performance"""
+    """Benchmark prefill performance. Returns throughput."""
    seed(0)
    # Fixed length input, minimal output to focus on prefill
    prompt_token_ids = [[randint(0, 10000) for _ in range(input_len)] for _ in range(num_seqs)]
@@ -35,7 +36,28 @@ def bench_prefill(llm, num_seqs, input_len):
    t = time.time() - t
    total_input_tokens = num_seqs * input_len
    throughput = total_input_tokens / t
-    print(f"[Prefill] Input: {total_input_tokens}tok ({num_seqs}x{input_len}), Time: {t:.2f}s, Throughput: {throughput:.2f}tok/s")
+    label_str = f" ({label})" if label else ""
    print(f"[Prefill{label_str}] Input: {total_input_tokens}tok ({num_seqs}x{input_len}), Time: {t:.2f}s, Throughput: {throughput:.2f}tok/s")
    return throughput
 def create_llm(path, max_len, enable_minference=False, minference_budget=0.3,
               minference_vertical=1000, minference_slash=6096,
               gpu_utilization=0.8):
    """Create LLM with specified configuration."""
    kwargs = {
        "enforce_eager": True,  # MInference uses Triton, not compatible with CUDA graphs
        "max_model_len": max_len,
        "max_num_batched_tokens": max_len,
        "gpu_memory_utilization": gpu_utilization,
    }
    if enable_minference:
        kwargs["sparse_policy"] = SparsePolicyType.MINFERENCE
        kwargs["minference_adaptive_budget"] = minference_budget
        kwargs["minference_vertical_size"] = minference_vertical
        kwargs["minference_slash_size"] = minference_slash
    return LLM(path, **kwargs)
 def main():
@@ -46,24 +68,17 @@ def main():
    parser.add_argument("--max-len", type=int, default=32*1024, help="Max model length (default: 32K)")
    parser.add_argument("--bench-decode", action="store_true", help="Run decode benchmark (default: prefill only)")
    parser.add_argument("--bench-all", action="store_true", help="Run both prefill and decode benchmarks")
    parser.add_argument("--enable-minference", action="store_true", help="Enable MInference sparse prefill")
    parser.add_argument("--minference-budget", type=float, default=0.3, help="MInference adaptive budget (default: 0.3, use 0 for fixed mode)")
    parser.add_argument("--minference-vertical", type=int, default=1000, help="Fixed vertical_size (only used when budget=0)")
    parser.add_argument("--minference-slash", type=int, default=6096, help="Fixed slash_size (only used when budget=0)")
    parser.add_argument("--gpu-utilization", type=float, default=0.9, help="GPU memory utilization (default: 0.9)")
    parser.add_argument("--compare", action="store_true", help="Compare baseline vs MInference (runs both)")
    args = parser.parse_args()
    path = os.path.expanduser("~/models/Qwen3-4B-Instruct-2507/")
    max_len = args.max_len
    print(f"\n[nanovllm GPU] max_len={max_len}")
    llm = LLM(
        path,
        enforce_eager=False,
        max_model_len=max_len,
        max_num_batched_tokens=max_len,
    )
    # Warmup
    print("\nWarming up...")
    llm.generate(["Benchmark warmup: "], SamplingParams(max_tokens=10))
    # Default input lengths
    prefill_input_len = args.input_len if args.input_len else max_len - 1
    decode_input_len = args.input_len if args.input_len else max_len - args.output_len
@@ -72,17 +87,128 @@ def main():
    run_prefill = not args.bench_decode or args.bench_all
    run_decode = args.bench_decode or args.bench_all
-    if run_prefill:
+    # Convert budget=0 to None for fixed mode
-        print("\n" + "=" * 60)
+    minference_budget = args.minference_budget if args.minference_budget > 0 else None
        print("Prefill Benchmark (nanovllm GPU)")
        print("=" * 60)
        bench_prefill(llm, num_seqs=1, input_len=prefill_input_len)
-    if run_decode:
+    if args.compare:
-        print("\n" + "=" * 60)
+        # Compare baseline vs MInference using subprocesses to avoid NCCL issues
-        print("Decode Benchmark (nanovllm GPU)")
+        import subprocess
-        print("=" * 60)
+        import sys
-        bench_decode(llm, num_seqs=1, input_len=decode_input_len, output_len=args.output_len)
+
        print(f"\n{'='*60}")
        print(f"Baseline vs MInference Comparison")
        print(f"Input length: {prefill_input_len} tokens")
        if minference_budget is not None:
            print(f"MInference mode: adaptive (budget={minference_budget}, {minference_budget*100:.0f}% compute)")
        else:
            print(f"MInference mode: fixed (vertical={args.minference_vertical}, slash={args.minference_slash})")
        print(f"{'='*60}")
        # Get PYTHONPATH for subprocess
        pythonpath = os.environ.get("PYTHONPATH", "")
        # Run baseline in subprocess
        print(f"\n[1/2] Running baseline (FULL attention)...")
        cmd_baseline = [
            sys.executable, __file__,
            "--input-len", str(prefill_input_len),
            "--max-len", str(max_len),
            "--gpu-utilization", str(args.gpu_utilization),
        ]
        env = os.environ.copy()
        result = subprocess.run(cmd_baseline, capture_output=True, text=True, env=env)
        print(result.stdout)
        if result.returncode != 0:
            print(f"Error: {result.stderr}")
            return
        # Parse baseline throughput
        baseline_throughput = None
        for line in result.stdout.split('\n'):
            if "Throughput:" in line and "tok/s" in line:
                # Extract throughput value
                import re
                match = re.search(r'Throughput:\s*([\d.]+)tok/s', line)
                if match:
                    baseline_throughput = float(match.group(1))
        # Run MInference in subprocess
        if minference_budget is not None:
            print(f"\n[2/2] Running MInference (budget={minference_budget})...")
        else:
            print(f"\n[2/2] Running MInference (vertical={args.minference_vertical}, slash={args.minference_slash})...")
        cmd_minference = [
            sys.executable, __file__,
            "--input-len", str(prefill_input_len),
            "--max-len", str(max_len),
            "--gpu-utilization", str(args.gpu_utilization),
            "--enable-minference",
            "--minference-budget", str(args.minference_budget),
            "--minference-vertical", str(args.minference_vertical),
            "--minference-slash", str(args.minference_slash),
        ]
        result = subprocess.run(cmd_minference, capture_output=True, text=True, env=env)
        print(result.stdout)
        if result.returncode != 0:
            print(f"Error: {result.stderr}")
            return
        # Parse MInference throughput
        minference_throughput = None
        for line in result.stdout.split('\n'):
            if "Throughput:" in line and "tok/s" in line:
                import re
                match = re.search(r'Throughput:\s*([\d.]+)tok/s', line)
                if match:
                    minference_throughput = float(match.group(1))
        # Comparison
        if baseline_throughput and minference_throughput:
            print(f"\n{'='*60}")
            print(f"Results Summary")
            print(f"{'='*60}")
            print(f"Baseline:   {baseline_throughput:,.0f} tok/s")
            print(f"MInference: {minference_throughput:,.0f} tok/s")
            speedup = minference_throughput / baseline_throughput
            if speedup >= 1.0:
                print(f"Speedup:    {speedup:.2f}x faster")
            else:
                print(f"Slowdown:   {1/speedup:.2f}x slower")
            print(f"{'='*60}")
        else:
            print("Failed to parse throughput values")
    else:
        # Single run mode
        mode = "MInference" if args.enable_minference else "GPU"
        print(f"\n[nanovllm {mode}] max_len={max_len}")
        if args.enable_minference:
            if minference_budget is not None:
                print(f"MInference mode: adaptive (budget={minference_budget})")
            else:
                print(f"MInference mode: fixed (vertical={args.minference_vertical}, slash={args.minference_slash})")
        llm = create_llm(path, max_len, enable_minference=args.enable_minference,
                        minference_budget=minference_budget,
                        minference_vertical=args.minference_vertical,
                        minference_slash=args.minference_slash,
                        gpu_utilization=args.gpu_utilization)
        # Warmup
        print("\nWarming up...")
        llm.generate(["Benchmark warmup: "], SamplingParams(max_tokens=10))
        if run_prefill:
            print("\n" + "=" * 60)
            print(f"Prefill Benchmark (nanovllm {mode})")
            print("=" * 60)
            bench_prefill(llm, num_seqs=1, input_len=prefill_input_len)
        if run_decode:
            print("\n" + "=" * 60)
            print(f"Decode Benchmark (nanovllm {mode})")
            print("=" * 60)
            bench_decode(llm, num_seqs=1, input_len=decode_input_len, output_len=args.output_len)
 if __name__ == "__main__":
--- a/bench_vllm.py
+++ b/bench_vllm.py
@@ -1,4 +1,5 @@
 import os
 os.environ["VLLM_USE_V1"] = "1"
 import time
 from random import randint, seed
@@ -8,8 +9,12 @@ from vllm import LLM, SamplingParams
 def bench_decode(llm, num_seqs, input_len, output_len):
    """Benchmark decode performance"""
    seed(0)
-    prompt_token_ids = [[randint(0, 10000) for _ in range(input_len)] for _ in range(num_seqs)]
+    prompt_token_ids = [
-    sampling_params = SamplingParams(temperature=0.6, ignore_eos=True, max_tokens=output_len)
+        [randint(0, 10000) for _ in range(input_len)] for _ in range(num_seqs)
    ]
    sampling_params = SamplingParams(
        temperature=0.6, ignore_eos=True, max_tokens=output_len
    )
    prompt_token_ids = [dict(prompt_token_ids=p) for p in prompt_token_ids]
    t = time.time()
@@ -21,15 +26,21 @@ def bench_decode(llm, num_seqs, input_len, output_len):
    decode_tokens = num_seqs * output_len
    decode_throughput = decode_tokens / t
-    print(f"[Decode] Input: {num_seqs}x{input_len}tok, Output: {decode_tokens}tok, Time: {t:.2f}s")
+    print(
-    print(f"         Throughput: {decode_throughput:.2f} tok/s (includes prefill overhead)")
+        f"[Decode] Input: {num_seqs}x{input_len}tok, Output: {decode_tokens}tok, Time: {t:.2f}s"
    )
    print(
        f"         Throughput: {decode_throughput:.2f} tok/s (includes prefill overhead)"
    )
 def bench_prefill(llm, num_seqs, input_len):
    """Benchmark prefill performance"""
    seed(0)
    # Fixed length input, minimal output to focus on prefill
-    prompt_token_ids = [[randint(0, 10000) for _ in range(input_len)] for _ in range(num_seqs)]
+    prompt_token_ids = [
        [randint(0, 10000) for _ in range(input_len)] for _ in range(num_seqs)
    ]
    sampling_params = SamplingParams(temperature=0.6, ignore_eos=True, max_tokens=1)
    prompt_token_ids = [dict(prompt_token_ids=p) for p in prompt_token_ids]
@@ -38,17 +49,39 @@ def bench_prefill(llm, num_seqs, input_len):
    t = time.time() - t
    total_input_tokens = num_seqs * input_len
    throughput = total_input_tokens / t
-    print(f"[Prefill] Input: {total_input_tokens}tok ({num_seqs}x{input_len}), Time: {t:.2f}s, Throughput: {throughput:.2f}tok/s")
+    print(
        f"[Prefill] Input: {total_input_tokens}tok ({num_seqs}x{input_len}), Time: {t:.2f}s, Throughput: {throughput:.2f}tok/s"
    )
 def main():
    import argparse
-    parser = argparse.ArgumentParser(description="Benchmark vLLM performance (for comparison)")
+
-    parser.add_argument("--input-len", type=int, default=None, help="Input length in tokens")
+    parser = argparse.ArgumentParser(
-    parser.add_argument("--output-len", type=int, default=64, help="Output length for decode benchmark (default: 64)")
+        description="Benchmark vLLM performance (for comparison)"
-    parser.add_argument("--max-len", type=int, default=32*1024, help="Max model length (default: 32K)")
+    )
-    parser.add_argument("--bench-decode", action="store_true", help="Run decode benchmark (default: prefill only)")
+    parser.add_argument(
-    parser.add_argument("--bench-all", action="store_true", help="Run both prefill and decode benchmarks")
+        "--input-len", type=int, default=None, help="Input length in tokens"
    )
    parser.add_argument(
        "--output-len",
        type=int,
        default=64,
        help="Output length for decode benchmark (default: 64)",
    )
    parser.add_argument(
        "--max-len", type=int, default=32 * 1024, help="Max model length (default: 32K)"
    )
    parser.add_argument(
        "--bench-decode",
        action="store_true",
        help="Run decode benchmark (default: prefill only)",
    )
    parser.add_argument(
        "--bench-all",
        action="store_true",
        help="Run both prefill and decode benchmarks",
    )
    args = parser.parse_args()
    path = os.path.expanduser("~/models/Qwen3-4B-Instruct-2507/")
@@ -61,7 +94,7 @@ def main():
        enforce_eager=False,
        max_model_len=max_len,
        max_num_seqs=128,
-        gpu_memory_utilization=0.9,
+        gpu_memory_utilization=0.7,
    )
    # Warmup
@@ -86,7 +119,9 @@ def main():
        print("\n" + "=" * 60)
        print("Decode Benchmark (vLLM)")
        print("=" * 60)
-        bench_decode(llm, num_seqs=1, input_len=decode_input_len, output_len=args.output_len)
+        bench_decode(
            llm, num_seqs=1, input_len=decode_input_len, output_len=args.output_len
        )
 if __name__ == "__main__":
--- a/docs/architecture_guide.md
+++ b/docs/architecture_guide.md
@@ -0,0 +1,189 @@
 # Architecture Guide
 This document describes the core architecture and layer-wise CPU offload system of nano-vLLM.
 ## Core Components
 | Component | File | Purpose |
 |-----------|------|---------|
 | **LLMEngine** | `llm_engine.py` | Main entry, runs prefill-decode loop |
 | **ModelRunner** | `model_runner.py` | Loads weights, allocates KV cache, CUDA graphs, layer-wise offload |
 | **Scheduler** | `scheduler.py` | Two-phase scheduling (prefill → decode) |
 | **BlockManager** | `block_manager.py` | Paged attention with prefix caching (xxhash), default block size 4096 |
 | **Attention** | `layers/attention.py` | FlashAttention for standard inference |
 ## Layer-wise CPU Offload System
 ### Design Philosophy
 Unlike chunked prefill (which processes chunks across all layers), **layer-wise offload** processes the entire sequence through one layer at a time:
 ```
 Layer 0: [full sequence] → compute → offload K,V to CPU
 Layer 1: [full sequence] → compute → offload K,V to CPU
 ...
 Layer N: [full sequence] → compute → offload K,V to CPU
 ```
 **Benefits**:
 - Supports MInference sparse attention (requires full KV access per layer)
 - Simpler memory management (one layer's KV in GPU at a time)
 - Peak GPU memory = one layer's KV cache + attention workspace
 ### Key Files
 | File | Purpose |
 |------|---------|
 | `nanovllm/engine/model_runner.py` | Main implementation (`run_layerwise_offload_prefill`, `run_layerwise_offload_decode`) |
 | `nanovllm/kvcache/hybrid_manager.py` | CPU block management helpers |
 | `nanovllm/kvcache/offload_engine.py` | CPU/GPU cache storage, ring buffer, async transfers |
 ### Memory Layout
 **CPU Cache** (pinned memory):
 ```python
 k_cache_cpu: [num_layers, num_cpu_blocks, block_size, kv_heads, head_dim]
 v_cache_cpu: [num_layers, num_cpu_blocks, block_size, kv_heads, head_dim]
 ```
 **GPU Ring Buffer** (for decode H2D pipeline):
 ```python
 layer_k_cache: [num_kv_buffers, max_seq_len, kv_heads, head_dim]
 layer_v_cache: [num_kv_buffers, max_seq_len, kv_heads, head_dim]
 ```
 **Per-layer KV size** (Qwen3-4B: 8 kv_heads × 128 head_dim × 2 bytes × 2 for K+V = 4KB/token):
 | Context Length | KV per Layer |
 |----------------|--------------|
 | 128K tokens | 512 MB |
 | 256K tokens | 1 GB |
 | 512K tokens | 2 GB |
 | 1M tokens | 4 GB |
 ---
 ## Prefill Flow
 ```python
 def run_layerwise_offload_prefill(self, seqs: list[Sequence]) -> list[int]:
    # 1. Embedding
    hidden_states = self.model.model.embed_tokens(input_ids)
    # 2. Process each layer
    for layer_id in range(num_layers):
        # QKV projection + norms + RoPE
        q = apply_rotary_pos_emb(q_proj(hidden_states), cos, sin)
        k = apply_rotary_pos_emb(k_proj(hidden_states), cos, sin)
        v = v_proj(hidden_states)
        # Full FlashAttention (entire sequence)
        attn_out = flash_attn_varlen_func(q, k, v, cu_seqlens, max_seqlen, causal=True)
        # MLP
        hidden_states = mlp(attn_out + residual)
        # Synchronous offload to CPU (CRITICAL: must be sync to avoid memory reuse bugs)
        self._offload_layer_kv_to_cpu_sync(layer_id, k, v, cpu_block_ids, total_tokens)
    # 3. Final norm + sampling
    return sampled_tokens
 ```
 ---
 ## Decode Flow
 ```python
 def run_layerwise_offload_decode(self, seqs: list[Sequence]) -> list[int]:
    # Ring buffer pipeline: preload first N layers
    for i in range(num_buffers):
        offload_engine.load_layer_kv_to_buffer(i, i, cpu_block_table, valid_tokens)
    # For each layer:
    for layer_id in range(num_layers):
        current_buffer = layer_id % num_buffers
        # 1. Wait for buffer load to complete
        offload_engine.wait_buffer_load(current_buffer)
        # 2. Get prefilled KV from ring buffer
        k_prefill, v_prefill = offload_engine.get_buffer_kv(current_buffer, total_prefill_tokens)
        # 3. Compute new Q,K,V for current token
        q_new = apply_rotary_pos_emb(q_proj(hidden_states), cos, sin)
        k_new = apply_rotary_pos_emb(k_proj(hidden_states), cos, sin)
        v_new = v_proj(hidden_states)
        # 4. Concatenate and compute attention
        k_full = torch.cat([k_prefill, k_new], dim=0)
        v_full = torch.cat([v_prefill, v_new], dim=0)
        attn_out = flash_attn_varlen_func(q_new, k_full, v_full, ..., causal=False)
        # Note: causal=False because single query token should attend to ALL keys
        # 5. Mark buffer done, start loading next layer
        offload_engine.record_buffer_compute_done(current_buffer)
        if layer_id + num_buffers < num_layers:
            offload_engine.load_layer_kv_to_buffer(current_buffer, layer_id + num_buffers, ...)
 ```
 ---
 ## Critical Implementation Details
 ### 1. Synchronous Offload Required
 Async offload with `non_blocking=True` causes memory reuse bugs:
 ```python
 # BUG: PyTorch may reuse k,v GPU memory before async copy completes
 offload_engine.k_cache_cpu[layer_id, block_id].copy_(k[start:end], non_blocking=True)
 # CORRECT: Synchronous copy ensures data integrity
 offload_engine.k_cache_cpu[layer_id, block_id, :size].copy_(k[start:end])  # sync
 ```
 ### 2. Decode Attention: causal=False
 During decode, the single query token must attend to ALL keys (not just preceding ones):
 ```python
 # Prefill: causal=True (each token only attends to previous tokens)
 attn_out = flash_attn_varlen_func(..., causal=True)
 # Decode: causal=False (query at position N attends to all N-1 prefill + itself)
 attn_out = flash_attn_varlen_func(..., causal=False)
 ```
 ### 3. Ring Buffer Synchronization
 The ring buffer pipeline requires careful ordering:
 ```python
 # CORRECT order:
 offload_engine.store_decode_kv(layer_id, pos, k_new, v_new)  # Store new KV
 offload_engine.record_buffer_compute_done(current_buffer)     # Mark done FIRST
 offload_engine.load_layer_kv_to_buffer(...)                   # THEN start next load
 # BUG: Starting load before marking done causes race condition
 offload_engine.load_layer_kv_to_buffer(...)  # WRONG: buffer still in use!
 offload_engine.record_buffer_compute_done(current_buffer)
 ```
 ---
 ## Helper Methods in HybridKVCacheManager
 ```python
 # Get all CPU blocks for a sequence
 cpu_blocks = manager.get_all_cpu_blocks(seq)  # List[int]
 # Get only prefilled (offloaded) CPU blocks
 prefilled_blocks = manager.get_prefilled_cpu_blocks(seq)  # List[int]
 # Get cached prefill length (doesn't change during decode)
 prefill_len = manager.get_prefill_len(seq)  # int
 # Get decode start position
 decode_pos = manager.get_decode_start_pos(seq)  # int
 ```
--- a/docs/cuda_graph_offload_guide.md
+++ b/docs/cuda_graph_offload_guide.md
@@ -0,0 +1,196 @@
 # CUDA Graph Support for CPU Offload Mode
 This document describes the CUDA graph implementation for the CPU offload decode path, which provides significant performance improvements for decode throughput.
 ## Overview
 CUDA graphs capture a sequence of GPU operations and replay them with minimal CPU overhead. In offload mode, we capture per-layer graphs for the decode path, achieving **4x decode throughput improvement**.
 ## Performance Results
 | Metric | Eager Mode | CUDA Graph | Improvement |
 |--------|------------|------------|-------------|
 | Decode Throughput | ~12 tok/s | ~50 tok/s | **4.2x** |
 | TPOT (Time per output token) | ~80ms | ~19ms | **4.2x** |
 | Prefill Throughput | ~8000 tok/s | ~8000 tok/s | Same |
 ## Architecture
 ### Why Standard CUDA Graph Capture Doesn't Work
 The standard `capture_cudagraph()` captures the PagedAttention decode path:
 - Uses block tables for scattered KV cache access
 - `Attention.k_cache/v_cache` point to PagedAttention buffers
 In offload mode, the decode path is different:
 - Uses contiguous ring buffers for KV cache
 - `Attention.k_cache/v_cache` dynamically point to ring buffer slices
 - H2D transfers interleaved with compute
 ### Per-Layer Graph Design
 We capture one CUDA graph per transformer layer:
 ```
 ┌─────────────────────────────────────────────────────────────┐
 │                    Offload Decode with CUDA Graphs          │
 ├─────────────────────────────────────────────────────────────┤
 │                                                             │
 │  Initialization:                                            │
 │    capture_offload_cudagraph() captures 36 layer graphs     │
 │    Each graph: layer.forward() with ring buffer as cache    │
 │                                                             │
 │  Decode Step:                                               │
 │    1. Embedding (eager, outside graph)                      │
 │    2. For each layer:                                       │
 │       a. Wait for H2D load (outside graph)                  │
 │       b. Copy decode KV to ring buffer (outside graph)      │
 │       c. Set Attention.k_cache = ring_buffer[buffer_idx]    │
 │       d. Set context (slot_mapping, context_lens)           │
 │       e. graph.replay() - layer forward                     │
 │       f. synchronize()                                      │
 │       g. Copy layer_outputs -> hidden_states                │
 │       h. Copy new KV to decode buffer (outside graph)       │
 │       i. Start next layer H2D load                          │
 │    3. Final norm and logits (eager)                         │
 │                                                             │
 └─────────────────────────────────────────────────────────────┘
 ```
 ### Ring Buffer Mapping
 Each layer maps to a ring buffer slot:
 ```python
 buffer_idx = layer_id % num_kv_buffers
 ```
 With 4 buffers and 36 layers:
 - Layer 0, 4, 8, ... use buffer 0
 - Layer 1, 5, 9, ... use buffer 1
 - Layer 2, 6, 10, ... use buffer 2
 - Layer 3, 7, 11, ... use buffer 3
 ## Implementation Details
 ### Graph Capture (`capture_offload_cudagraph`)
 Location: `model_runner.py:1075-1164`
 ```python
 def capture_offload_cudagraph(self):
    # Fixed-address tensors for graph I/O
    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 slice
        attn_module.k_cache = ring_buffer[buffer_idx:buffer_idx+1]
        attn_module.v_cache = ring_buffer[buffer_idx:buffer_idx+1]
        # Set context for contiguous mode
        set_context(is_prefill=False, slot_mapping=...,
                    context_lens=..., block_tables=None)
        # Warmup and capture
        with torch.cuda.graph(graph, pool):
            out_h, out_r = layer(positions, hidden_states, residual)
            layer_outputs.copy_(out_h)
            layer_residual.copy_(out_r)
        # Propagate state for next layer's capture
        hidden_states.copy_(layer_outputs)
        residual.copy_(layer_residual)
 ```
 Key design decisions:
 1. **Fixed-address tensors**: Graph inputs/outputs use pre-allocated tensors
 2. **Include copy in graph**: `layer_outputs.copy_(out_h)` is captured
 3. **State propagation**: Update hidden_states between layer captures
 4. **Random initialization**: Use `randn` instead of zeros for realistic distributions
 ### Graph Replay (`run_layerwise_offload_decode`)
 Location: `model_runner.py:844-1031`
 ```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):
    # H2D and buffer setup (outside graph)
    offload_engine.wait_buffer_load(current_buffer)
    attn_module.k_cache = ring_buffer[current_buffer:current_buffer+1]
    set_context(...)
    if use_cuda_graph:
        # 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"])
    else:
        # Eager execution
        hidden_states, residual = layer(positions, hidden_states, residual)
 ```
 Key points:
 1. **Synchronization required**: `synchronize()` after each graph replay
 2. **Manual state propagation**: Copy layer_outputs to hidden_states between replays
 3. **H2D outside graph**: Ring buffer loads happen before graph replay
 ## Limitations and Future Work
 ### Current Limitations
 1. **Per-layer sync overhead**: Each layer requires synchronization
 2. **No kernel fusion across layers**: Each layer is a separate graph
 3. **Fixed batch size**: Only supports batch_size=1 for offload
 ### Future Optimization: Full-Decode Graph
 Potential improvement: Capture entire decode step as single graph
 - Complete all H2D loads before graph
 - Single graph covers all 36 layers
 - Better kernel fusion, less CPU overhead
 - More complex to implement (handle buffer rotation inside graph)
 ## Testing
 Run needle test with CUDA graph:
 ```bash
 PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python tests/test_needle.py \
    --input-len 32768 \
    --enable-offload \
    --use-cuda-graph
 ```
 Run benchmark:
 ```bash
 PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python bench_offload.py \
    --input-len 16384 \
    --bench-all
 ```
 ## Files Modified
 | File | Changes |
 |------|---------|
 | `model_runner.py:46-50` | Call `capture_offload_cudagraph()` for offload mode |
 | `model_runner.py:69-73` | Clean up offload graph resources in `exit()` |
 | `model_runner.py:844-1031` | Add CUDA graph support to `run_layerwise_offload_decode()` |
 | `model_runner.py:1075-1164` | New `capture_offload_cudagraph()` method |
 | `tests/test_needle.py` | Add `--use-cuda-graph` flag |
--- a/docs/debugging_guide.md
+++ b/docs/debugging_guide.md
@@ -0,0 +1,142 @@
 # Debugging Guide
 This document provides debugging techniques for nano-vLLM, including PyTorch hooks for capturing intermediate tensors.
 ## PyTorch Hooks for Debugging
 ### Hook Positions in Qwen3
 ```
 decoder_layer
 ├── input_layernorm (RMSNorm)
 ├── self_attn (Qwen3Attention)          ← Hook here for attention I/O after o_proj
 │   ├── q_proj → q_norm → RoPE
 │   ├── k_proj → k_norm → RoPE
 │   ├── v_proj
 │   ├── attn (Attention)                ← Hook here for Q/K/V tensors
 │   │   └── FlashAttention / SDPA
 │   └── o_proj
 ├── post_attention_layernorm (RMSNorm)
 └── mlp (Qwen3MLP)
 ```
 ### Hook Types & Data Shapes
 | Hook Position | Type | Captured Data |
 |---------------|------|---------------|
 | `self_attn` | post | `[batch, seq_len, hidden_size]` - after o_proj |
 | `self_attn.attn` | pre | Q,K,V: `[seq_len, num_heads, head_dim]` - after RoPE |
 | `self_attn.attn` | post | `[seq_len, num_heads, head_dim]` - before o_proj |
 ### Example: Capture Attention Outputs
 ```python
 storage = {}
 def make_hook(layer_id: int, storage: dict):
    def hook(module, inputs, output):
        if isinstance(output, tuple):
            attn_output = output[0]
        else:
            attn_output = output
        # nanovllm shape: [num_tokens, hidden_size] -> add batch dim
        if attn_output.dim() == 2:
            attn_output = attn_output.unsqueeze(0)
        storage[layer_id] = attn_output.detach().clone()
    return hook
 # Register hooks
 hooks = []
 for layer_idx, layer in enumerate(model.model.layers):
    hooks.append(layer.self_attn.register_forward_hook(make_hook(layer_idx, storage)))
 # Run inference...
 # Cleanup
 for hook in hooks:
    hook.remove()
 ```
 ### Reference Implementation
 Key files for comparison testing:
 | File | Purpose |
 |------|---------|
 | `tests/modeling_qwen3.py` | Reference Qwen3 implementation (torch + transformers only) |
 | `tests/test_needle_ref.py` | Reference needle test using custom Qwen3 |
 | `tests/test_needle.py` | Needle-in-haystack test for nanovllm |
 ### Common Pitfalls
 1. **Shape mismatch**: nanovllm uses `[num_tokens, ...]` while torch uses `[batch, seq_len, ...]`
 2. **Hook position**: `self_attn` captures after o_proj, `self_attn.attn` captures before o_proj
 3. **Output format**: nanovllm returns tuple `(attn_output, None)`, handle with `output[0]`
 ---
 ## Memory Debugging
 ### Track Peak GPU Memory
 ```python
 import torch
 # Reset stats before operation
 torch.cuda.reset_peak_memory_stats()
 torch.cuda.empty_cache()
 # Run operation
 outputs = llm.generate([prompt], sampling_params)
 # Check peak
 peak_gb = torch.cuda.max_memory_allocated() / 1024**3
 print(f"Peak GPU memory: {peak_gb:.2f} GB")
 ```
 ### Monitor Memory During Execution
 ```python
 import torch
 def memory_snapshot():
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    print(f"Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB")
 # Add snapshots at key points in your code
 ```
 ---
 ## Comparing Outputs
 ### Needle-in-Haystack Test
 ```bash
 # Test with CPU offload
 PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python tests/test_needle.py --enable-offload --input-len 8192
 # Test without CPU offload (GPU-only)
 PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python tests/test_needle.py --input-len 8192
 # Compare with reference implementation
 PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python tests/test_needle_ref.py --input-len 8192
 ```
 ### Tensor Comparison
 ```python
 def compare_tensors(a, b, name, rtol=1e-3, atol=1e-5):
    if a.shape != b.shape:
        print(f"{name}: Shape mismatch {a.shape} vs {b.shape}")
        return False
    diff = (a - b).abs()
    max_diff = diff.max().item()
    mean_diff = diff.mean().item()
    close = torch.allclose(a, b, rtol=rtol, atol=atol)
    print(f"{name}: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}, close={close}")
    return close
 ```
--- a/docs/gpu_only_performance_issue.md
+++ b/docs/gpu_only_performance_issue.md
@@ -0,0 +1,194 @@
 # GPU-only Performance Issue: PagedAttention Scatter Overhead
 ## Problem Summary
 GPU-only mode with MInference is **slower** than CPU offload mode for long-context single-sequence inference:
 | Mode | Prefill Speed (32K tokens, Qwen3-4B) |
 |------|--------------------------------------|
 | GPU-only + MInference | 3383 tok/s |
 | Offload + MInference | 5373 tok/s |
 This counterintuitive result is caused by **unnecessary `store_kvcache` overhead** in the GPU-only path.
 ## Root Cause Analysis
 ### GPU-only Execution Path
 ```python
 # attention.py line 86-110
 def forward(self, q, k, v):
    # ALWAYS store to cache first - OVERHEAD HERE
    if k_cache.numel() and v_cache.numel():
        store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)  # ← Always executed
    if context.is_prefill:
        if context.sparse_prefill_policy is not None:
            # MInference: uses k, v directly, NOT k_cache!
            o = sparse_prefill_attention(q, k, v, layer_id)
        else:
            # Full attention: also uses k, v directly
            o = flash_attn_varlen_func(q, k, v, ...)
 ```
 **Key observation**: Prefill attention **never reads from cache** - it uses the computed k, v directly. But `store_kvcache` is always called before attention.
 ### The `store_kvcache` Overhead
 ```python
 # attention.py line 8-59
 def store_kvcache(key, value, k_cache, v_cache, slot_mapping):
    # 1. Filter invalid slots (conditional logic)
    valid_mask = slot_mapping >= 0
    valid_slots = slot_mapping[valid_mask]
    valid_keys = key[valid_mask]
    # 2. Reshape for scatter operation
    k_cache_flat = k_cache.view(total_slots, D)
    valid_keys_flat = valid_keys.reshape(-1, D)
    # 3. Scatter write via index_copy_ - EXPENSIVE!
    k_cache_flat.index_copy_(0, valid_slots.long(), valid_keys_flat)
    v_cache_flat.index_copy_(0, valid_slots.long(), valid_values_flat)
 ```
 This scatter operation is called for **every layer** (28 layers for Qwen3-4B), writing **all tokens** (32K) to GPU cache.
 ### Offload Path (No Such Overhead)
 ```python
 # model_runner.py - run_layerwise_offload_prefill
 for layer_id in range(num_layers):
    # QKV projection + RoPE
    q, k = layer.self_attn.rotary_emb(positions, q, k)
    # Sparse attention - directly uses k, v
    attn_output = sparse_prefill_attention(q, k, v, layer_id)
    # Contiguous copy to CPU - no scatter!
    offload_engine.offload_layer_kv_sync(layer_id, k, v, cpu_block_ids, total_tokens)
 ```
 ## Memory Layout Comparison
 | Aspect | GPU-only (PagedAttention) | Offload (Contiguous) |
 |--------|---------------------------|----------------------|
 | **Layout** | `[num_blocks, block_size, heads, dim]` | `[seq_len, heads, dim]` |
 | **Write pattern** | Scatter via `index_copy_` | Contiguous `copy_()` |
 | **Indirection** | slot_mapping lookup | None |
 | **Memory efficiency** | High (shared block pool) | Low (reserved per seq) |
 | **Write performance** | Slow (memory-bound scatter) | Fast (simple DMA) |
 ### Why PagedAttention Uses Scatter
 PagedAttention is designed for:
 1. **Multi-sequence batching**: Different sequences share a block pool
 2. **Dynamic memory management**: No need to reserve max_len per sequence
 3. **Prefix caching**: Shared KV blocks across sequences
 But for **single-sequence long-context** inference, these benefits don't apply, and we only pay the scatter overhead.
 ## Why `store_kvcache` is Still Needed
 Even though prefill attention doesn't read from cache, **decode** does:
 ```python
 # attention.py line 111-114
 else:  # decode
    # Reads from cache!
    o = flash_attn_with_kvcache(q, k_cache, v_cache, block_table=...)
 ```
 So `store_kvcache` during prefill is preparing KV cache for future decode steps.
 ## Potential Optimizations
 ### Option 1: Async Store After Attention (Low Effort)
 Move `store_kvcache` after attention computation and make it async:
 ```python
 def forward(self, q, k, v):
    if context.is_prefill:
        # Compute attention first
        if context.sparse_prefill_policy is not None:
            o = sparse_prefill_attention(q, k, v, layer_id)
        else:
            o = flash_attn_varlen_func(q, k, v, ...)
        # Then store async (overlaps with next layer's QKV)
        if k_cache.numel():
            store_kvcache_async(k, v, k_cache, v_cache, slot_mapping)
    ...
 ```
 **Expected benefit**: Overlap store with compute, ~20-30% improvement.
 ### Option 2: Contiguous Layout for Single-Sequence Mode (Medium Effort)
 Add a "contiguous mode" for single-sequence long-context:
 ```python
 class ContiguousKVCache:
    """Simple contiguous KV cache for single-sequence mode."""
    def __init__(self, num_layers, max_seq_len, num_kv_heads, head_dim, dtype):
        self.k_cache = torch.zeros(num_layers, max_seq_len, num_kv_heads, head_dim, dtype=dtype)
        self.v_cache = torch.zeros(num_layers, max_seq_len, num_kv_heads, head_dim, dtype=dtype)
    def store(self, layer_id, k, v, start_pos):
        # Simple contiguous write - no scatter!
        seq_len = k.shape[0]
        self.k_cache[layer_id, start_pos:start_pos+seq_len] = k
        self.v_cache[layer_id, start_pos:start_pos+seq_len] = v
 ```
 **Expected benefit**: Match or exceed offload performance (~60% improvement).
 ### Option 3: Fused Store-Attention Kernel (High Effort)
 Create a fused Triton kernel that:
 1. Computes QKV projection
 2. Stores K, V to cache
 3. Computes attention
 This eliminates memory roundtrips entirely.
 **Expected benefit**: Best possible performance, but high implementation complexity.
 ## Recommended Action
 For **single-sequence long-context** workloads (the primary use case for MInference):
 1. **Short term**: Use offload mode - it's actually faster!
 2. **Medium term**: Implement Option 1 (async store) for quick win
 3. **Long term**: Consider Option 2 (contiguous layout) for GPU-only mode
 ## Performance Measurement
 To reproduce the benchmark:
 ```bash
 # GPU-only + MInference
 PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python tests/test_needle.py \
    --model ~/models/Qwen3-4B-Instruct-2507/ \
    --input-len 32768 \
    --enable-minference
 # Offload + MInference
 PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH python tests/test_needle.py \
    --model ~/models/Qwen3-4B-Instruct-2507/ \
    --input-len 32768 \
    --enable-offload \
    --enable-minference
 ```
 ## Related Files
 - `nanovllm/layers/attention.py`: `store_kvcache()` and `Attention.forward()`
 - `nanovllm/engine/model_runner.py`: `run_layerwise_offload_prefill()`
 - `nanovllm/kvcache/offload_engine.py`: `offload_layer_kv_sync()`
 ## References
 - [PagedAttention Paper](https://arxiv.org/abs/2309.06180) - vLLM's memory management
 - [MInference Paper](https://arxiv.org/abs/2407.02490) - Sparse prefill attention
--- a/docs/layerwise_offload_memory_analysis.md
+++ b/docs/layerwise_offload_memory_analysis.md
@@ -0,0 +1,547 @@
 # Layer-wise Offload Memory Analysis
 This document provides a detailed analysis of memory allocations in the layer-wise CPU offload system, distinguishing between pre-allocated (managed) memory and temporary (non-pre-allocated) memory.
 ## Variable Notation
 | Symbol | Description | Example (Qwen3-4B) |
 |--------|-------------|-------------------|
 | `seq_len` | Input sequence length | 131072 (128k) |
 | `hidden_size` | Model hidden dimension | 2560 |
 | `num_heads` | Number of attention heads | 20 |
 | `num_kv_heads` | Number of KV heads (GQA) | 8 |
 | `head_dim` | Dimension per head | 128 |
 | `intermediate_size` | MLP intermediate dimension | 13696 |
 | `num_layers` | Number of transformer layers | 36 |
 | `block_size` | KV cache block size | 1024 |
 | `num_kv_buffers` | Ring buffer count | 4 |
 | `num_cpu_blocks` | Number of CPU cache blocks | 128 |
 | `vocab_size` | Vocabulary size | 151936 |
 | `dtype_size` | Bytes per element (fp16/bf16) | 2 |
 Derived values:
 - `kv_dim = num_kv_heads × head_dim`
 - `q_size = num_heads × head_dim`
 - `kv_size = num_kv_heads × head_dim`
 - `qkv_size = q_size + 2 × kv_size`
 ---
 ## 1. Pre-allocated Memory (Managed by nanovllm)
 These tensors are allocated once during initialization and reused throughout inference.
 ### 1.1 OffloadEngine Managed Memory
 | Tensor | Shape | Size Formula | Location |
 |--------|-------|--------------|----------|
 | `layer_k_cache` | `[num_kv_buffers, seq_len, num_kv_heads, head_dim]` | `num_kv_buffers × seq_len × kv_dim × dtype_size` | GPU |
 | `layer_v_cache` | `[num_kv_buffers, seq_len, num_kv_heads, head_dim]` | `num_kv_buffers × seq_len × kv_dim × dtype_size` | GPU |
 | `decode_k_buffer` | `[num_layers, block_size, num_kv_heads, head_dim]` | `num_layers × block_size × kv_dim × dtype_size` | GPU |
 | `decode_v_buffer` | `[num_layers, block_size, num_kv_heads, head_dim]` | `num_layers × block_size × kv_dim × dtype_size` | GPU |
 | `k_cache_cpu` | `[num_layers, num_cpu_blocks, block_size, num_kv_heads, head_dim]` | `num_layers × num_cpu_blocks × block_size × kv_dim × dtype_size` | CPU (pinned) |
 | `v_cache_cpu` | `[num_layers, num_cpu_blocks, block_size, num_kv_heads, head_dim]` | `num_layers × num_cpu_blocks × block_size × kv_dim × dtype_size` | CPU (pinned) |
 **Total GPU (OffloadEngine)**: `2 × (num_kv_buffers × seq_len + num_layers × block_size) × kv_dim × dtype_size`
 **Total CPU (OffloadEngine)**: `2 × num_layers × num_cpu_blocks × block_size × kv_dim × dtype_size`
 ### 1.2 Model Weights
 | Component | Approximate Size |
 |-----------|-----------------|
 | Embedding | `vocab_size × hidden_size × dtype_size` |
 | Per-layer QKV proj | `hidden_size × qkv_size × dtype_size` |
 | Per-layer O proj | `q_size × hidden_size × dtype_size` |
 | Per-layer MLP | `hidden_size × 2 × intermediate_size × dtype_size + intermediate_size × hidden_size × dtype_size` |
 | Per-layer LayerNorm | `2 × hidden_size × dtype_size` |
 | LM Head | `hidden_size × vocab_size × dtype_size` |
 ### 1.3 RoPE Cache
 | Tensor | Shape | Size |
 |--------|-------|------|
 | `cos_sin_cache` | `[max_position, 1, head_dim]` | `max_position × head_dim × 4` (float32) |
 ---
 ## 2. Non-Pre-allocated Memory: Prefill Phase
 Location: `model_runner.py:run_layerwise_offload_prefill()`
 ### 2.1 Persistent Tensors (Live Throughout Prefill)
 | Variable | Line | Shape | Size | Notes |
 |----------|------|-------|------|-------|
 | `input_ids` | 488 | `[seq_len]` | `seq_len × 8` | int64 |
 | `positions` | 489 | `[seq_len]` | `seq_len × 8` | int64 |
 | `cu_seqlens` | 493 | `[2]` | negligible | int32 |
 | `hidden_states` | 497 | `[seq_len, hidden_size]` | `seq_len × hidden_size × dtype_size` | Embedding output |
 | `residual` | 506 | `[seq_len, hidden_size]` | `seq_len × hidden_size × dtype_size` | Residual connection |
 ### 2.2 Per-Layer Temporary Tensors
 These are allocated and deallocated within each layer iteration.
 #### 2.2.1 LayerNorm
 | Variable | Line | Shape | Size | Notes |
 |----------|------|-------|------|-------|
 | `hidden_ln` | 506-508 | `[seq_len, hidden_size]` | `seq_len × hidden_size × dtype_size` | Input layernorm output |
 **Inside RMSNorm** (`layernorm.py:add_rms_forward`):
 | Variable | Shape | Size | Notes |
 |----------|-------|------|-------|
 | `x.float()` | `[seq_len, hidden_size]` | `seq_len × hidden_size × 4` | Upcasted to float32 |
 | `var` | `[seq_len, 1]` | `seq_len × 4` | Variance |
 #### 2.2.2 QKV Projection
 | Variable | Line | Shape | Size | Notes |
 |----------|------|-------|------|-------|
 | `qkv` | 512 | `[seq_len, q_size + 2 × kv_size]` | `seq_len × qkv_size × dtype_size` | Merged QKV output |
 | `q` | 513-519 | `[seq_len, num_heads, head_dim]` | 0 (view) | View of qkv |
 | `k` | 513-520 | `[seq_len, num_kv_heads, head_dim]` | 0 (view) | View of qkv |
 | `v` | 513-521 | `[seq_len, num_kv_heads, head_dim]` | 0 (view) | View of qkv |
 #### 2.2.3 Q/K Norms (Qwen3 specific)
 | Variable | Line | Shape | Size | Notes |
 |----------|------|-------|------|-------|
 | `q.reshape()` | 526 | `[seq_len × num_heads, head_dim]` | 0 (view) | Reshape for norm |
 | `k.reshape()` | 528 | `[seq_len × num_kv_heads, head_dim]` | 0 (view) | Reshape for norm |
 | RMSNorm intermediates | - | see above | `seq_len × num_heads × head_dim × 4` | Float32 upcasting |
 #### 2.2.4 RoPE (Rotary Position Embedding)
 Location: `rotary_embedding.py:apply_rotary_emb()`
 | Variable | Line | Shape | Size | Notes |
 |----------|------|-------|------|-------|
 | `cos_sin` | 44 | `[seq_len, 1, head_dim]` | 0 (view) | View of cached cos_sin |
 | `cos` | 45 | `[seq_len, 1, head_dim/2]` | 0 (view) | Chunk view |
 | `sin` | 45 | `[seq_len, 1, head_dim/2]` | 0 (view) | Chunk view |
 **Inside `apply_rotary_emb` for Q** (`rotary_embedding.py:6-14`):
 | Variable | Shape | Size | Notes |
 |----------|-------|------|-------|
 | `x.float()` | `[seq_len, num_heads, head_dim]` | `seq_len × num_heads × head_dim × 4` | Upcast to float32 |
 | `x1` | `[seq_len, num_heads, head_dim/2]` | 0 (view) | Chunk view |
 | `x2` | `[seq_len, num_heads, head_dim/2]` | 0 (view) | Chunk view |
 | `y1 = x1*cos - x2*sin` | `[seq_len, num_heads, head_dim/2]` | `seq_len × num_heads × head_dim/2 × 4` | New tensor |
 | `y2 = x2*cos + x1*sin` | `[seq_len, num_heads, head_dim/2]` | `seq_len × num_heads × head_dim/2 × 4` | New tensor |
 | `torch.cat((y1, y2))` | `[seq_len, num_heads, head_dim]` | `seq_len × num_heads × head_dim × 4` | New tensor |
 | `.to(x.dtype)` | `[seq_len, num_heads, head_dim]` | `seq_len × num_heads × head_dim × dtype_size` | Downcast |
 **Inside `apply_rotary_emb` for K**:
 | Variable | Shape | Size | Notes |
 |----------|-------|------|-------|
 | Same pattern as Q | `[seq_len, num_kv_heads, head_dim]` | Similar, with `num_kv_heads` | |
 **Total RoPE temporary for Q+K**: ~`seq_len × (num_heads + num_kv_heads) × head_dim × 4 × 3` (float32 intermediates)
 #### 2.2.5 FlashAttention
 | Variable | Line | Shape | Size | Notes |
 |----------|------|-------|------|-------|
 | `attn_output` | 535 | `[seq_len, num_heads, head_dim]` | `seq_len × num_heads × head_dim × dtype_size` | Attention output |
 | Internal workspace | - | O(seq_len) | Variable | FlashAttention internal |
 #### 2.2.6 Output Projection
 | Variable | Line | Shape | Size | Notes |
 |----------|------|-------|------|-------|
 | `attn_output.view()` | 546 | `[seq_len, q_size]` | 0 (view) | Reshape for o_proj |
 | `o_proj(attn_output)` | 547 | `[seq_len, hidden_size]` | `seq_len × hidden_size × dtype_size` | O projection output |
 #### 2.2.7 Post-Attention LayerNorm
 Same as input layernorm (2.2.1).
 #### 2.2.8 MLP
 Location: `qwen3.py:Qwen3MLP.forward()`
 | Variable | Line | Shape | Size | Notes |
 |----------|------|-------|------|-------|
 | `gate_up` | 117 | `[seq_len, 2 × intermediate_size]` | `seq_len × 2 × intermediate_size × dtype_size` | **LARGEST TEMPORARY!** |
 | `x, y = chunk()` | activation.py:13 | `[seq_len, intermediate_size]` × 2 | 0 (views) | Chunk views |
 | `F.silu(x)` | activation.py:14 | `[seq_len, intermediate_size]` | `seq_len × intermediate_size × dtype_size` | SiLU activation |
 | `silu(x) * y` | activation.py:14 | `[seq_len, intermediate_size]` | `seq_len × intermediate_size × dtype_size` | Gated output |
 | `down_proj()` | 119 | `[seq_len, hidden_size]` | `seq_len × hidden_size × dtype_size` | MLP output |
 ### 2.3 Prefill Memory Summary
 **Peak per-layer temporary memory**:
 ```
 = qkv + RoPE_temps + attn_output + o_proj + layernorm + MLP_gate_up + MLP_activation
 ≈ seq_len × (qkv_size + (num_heads + num_kv_heads) × head_dim × 4 × 3
           + num_heads × head_dim + hidden_size × 2 + 2 × intermediate_size + intermediate_size) × dtype_size
 ```
 **Dominant term**: `seq_len × 2 × intermediate_size × dtype_size` (MLP gate_up)
 ---
 ## 3. Non-Pre-allocated Memory: Decode Phase
 Location: `model_runner.py:run_layerwise_offload_decode()`
 ### 3.1 Persistent Tensors
 | Variable | Line | Shape | Size | Notes |
 |----------|------|-------|------|-------|
 | `input_ids` | 604 | `[1]` | 8 bytes | Single token |
 | `positions` | 605 | `[1]` | 8 bytes | Single position |
 | `cu_seqlens_q` | 631 | `[2]` | 8 bytes | Fixed |
 | `valid_tokens_per_block` | 613-622 | Python list | negligible | |
 ### 3.2 Per-Layer Temporary Tensors
 #### 3.2.1 Views (Zero Additional Memory)
 | Variable | Line | Shape | Notes |
 |----------|------|-------|-------|
 | `k_prefill` | 682 | `[prefill_len, num_kv_heads, head_dim]` | View of ring buffer |
 | `v_prefill` | 682 | `[prefill_len, num_kv_heads, head_dim]` | View of ring buffer |
 | `k_decode_prev` | 686-687 | `[num_decode_tokens-1, num_kv_heads, head_dim]` | View of decode buffer |
 | `v_decode_prev` | 686-688 | `[num_decode_tokens-1, num_kv_heads, head_dim]` | View of decode buffer |
 #### 3.2.2 New Allocations
 | Variable | Line | Shape | Size | Notes |
 |----------|------|-------|------|-------|
 | `hidden_ln` | 654-657 | `[1, hidden_size]` | `hidden_size × dtype_size` | Tiny |
 | `qkv` | 660 | `[1, qkv_size]` | `qkv_size × dtype_size` | Tiny |
 | `q` | 667 | `[1, num_heads, head_dim]` | 0 (view) | |
 | `k_new` | 668 | `[1, num_kv_heads, head_dim]` | 0 (view) | |
 | `v_new` | 669 | `[1, num_kv_heads, head_dim]` | 0 (view) | |
 | **`k_full`** | 689/692 | `[prefill_len + num_decode_tokens, num_kv_heads, head_dim]` | `(prefill_len + num_decode_tokens) × kv_dim × dtype_size` | **torch.cat - NEW ALLOCATION** |
 | **`v_full`** | 690/693 | `[prefill_len + num_decode_tokens, num_kv_heads, head_dim]` | `(prefill_len + num_decode_tokens) × kv_dim × dtype_size` | **torch.cat - NEW ALLOCATION** |
 | `cu_seqlens_k` | 710 | `[2]` | 8 bytes | Created per layer |
 | `attn_output` | 712 | `[1, num_heads, head_dim]` | `num_heads × head_dim × dtype_size` | Tiny |
 | MLP temps | 728 | `[1, ...]` | negligible | Single token |
 ### 3.3 Decode Memory Summary
 **Peak per-layer temporary memory**:
 ```
 = k_full + v_full + small_tensors
 ≈ 2 × (prefill_len + num_decode_tokens) × num_kv_heads × head_dim × dtype_size
 ≈ 2 × seq_len × kv_dim × dtype_size
 ```
 **Dominant term**: `k_full` and `v_full` from `torch.cat()`
 ---
 ## 4. Memory Comparison Table
 For Qwen3-4B with 128k context:
 | Category | Memory | Notes |
 |----------|--------|-------|
 | **Pre-allocated GPU** | ~2.2 GB | Ring buffer + decode buffer |
 | **Pre-allocated CPU** | ~18.4 GB | Pinned memory |
 | **Model Weights** | ~8 GB | |
 | **Prefill Peak Temp** | ~10-12 GB | MLP gate_up dominant |
 | **Decode Peak Temp** | ~512 MB | k_full + v_full |
 ---
 ## 5. Optimization Opportunities
 ### 5.1 Decode: Pre-allocate k_full/v_full
 **Current** (L689-693):
 ```python
 k_full = torch.cat([k_prefill, k_decode_prev, k_new], dim=0)  # New allocation each layer
 v_full = torch.cat([v_prefill, v_decode_prev, v_new], dim=0)  # New allocation each layer
 ```
 **Optimized**:
 ```python
 # Pre-allocate in OffloadEngine.__init__():
 self.k_full_buffer = torch.zeros(max_seq_len + block_size, num_kv_heads, head_dim, ...)
 self.v_full_buffer = torch.zeros(max_seq_len + block_size, num_kv_heads, head_dim, ...)
 # In decode loop:
 total_len = prefill_len + num_decode_tokens
 k_full = self.k_full_buffer[:total_len]
 k_full[:prefill_len].copy_(k_prefill)
 k_full[prefill_len:prefill_len+num_decode_prev].copy_(k_decode_prev)
 k_full[-1:].copy_(k_new)
 ```
 **Savings**: ~512 MB per decode step (for 128k)
 ### 5.2 Decode: Reuse cu_seqlens_k
 **Current** (L710):
 ```python
 cu_seqlens_k = torch.tensor([0, total_kv_tokens], dtype=torch.int32, device="cuda")
 ```
 **Optimized**:
 ```python
 # Pre-allocate once:
 self.cu_seqlens_k = torch.zeros(2, dtype=torch.int32, device="cuda")
 # In decode loop:
 self.cu_seqlens_k[1] = total_kv_tokens
 ```
 **Savings**: Negligible memory, but reduces allocation overhead.
 ### 5.3 RoPE: In-place or Pre-allocated Buffers
 The RoPE implementation creates multiple float32 intermediate tensors. Options:
 1. Pre-allocate buffers for Q and K rotary outputs
 2. Use in-place operations where possible
 3. Use fused RoPE kernel (e.g., from FlashAttention)
 **Potential savings**: ~1.5 GB during prefill per layer
 ### 5.4 MLP: Cannot Optimize Easily
 The MLP `gate_up` tensor is inherently required for the gated activation:
 ```python
 gate_up = gate_up_proj(x)  # [seq_len, 2 × intermediate_size]
 x, y = gate_up.chunk(2, -1)
 output = silu(x) * y
 ```
 This is a fundamental computation pattern. Potential optimizations:
 - Chunked MLP computation (process seq_len in chunks)
 - Fused kernels that avoid materializing full gate_up
 ---
 ## 6. Memory Flow Diagram
 ### Prefill (per layer):
 ```
 hidden_states ──┬──► LayerNorm ──► hidden_ln
                │
 residual ◄──────┘
 hidden_ln ──► QKV_proj ──► qkv ──┬──► q ──► Q_norm ──► RoPE ──► q_rotated
                                 ├──► k ──► K_norm ──► RoPE ──► k_rotated
                                 └──► v
 q_rotated, k_rotated, v ──► FlashAttention ──► attn_output
 attn_output ──► O_proj ──► hidden_states'
 hidden_states', residual ──► LayerNorm ──► hidden_ln', residual'
 hidden_ln' ──► MLP_gate_up ──► gate_up ──► SiLU×gate ──► MLP_down ──► hidden_states''
 k_rotated, v ──► CPU_offload (sync copy)
 ```
 ### Decode (per layer):
 ```
 [CPU] k_cache_cpu, v_cache_cpu
           │
           ▼ (H2D async to ring buffer)
 [GPU] layer_k_cache[buffer_idx], layer_v_cache[buffer_idx]
           │
           ▼ (view)
      k_prefill, v_prefill
           │
           ├──► torch.cat([k_prefill, k_decode_prev, k_new]) ──► k_full  ⚠️ NEW ALLOC
           │
           └──► torch.cat([v_prefill, v_decode_prev, v_new]) ──► v_full  ⚠️ NEW ALLOC
 q_new, k_full, v_full ──► FlashAttention ──► attn_output
 k_new, v_new ──► decode_k_buffer, decode_v_buffer (in-place store)
 ```
 ---
 ## 7. Appendix: Size Calculations
 ### Qwen3-4B Example (128k context)
 ```python
 # Model config
 seq_len = 131072
 hidden_size = 2560
 num_heads = 20
 num_kv_heads = 8
 head_dim = 128
 intermediate_size = 13696
 num_layers = 36
 block_size = 1024
 num_kv_buffers = 4
 num_cpu_blocks = 128
 dtype_size = 2  # fp16/bf16
 # Derived
 kv_dim = num_kv_heads * head_dim  # 1024
 q_size = num_heads * head_dim     # 2560
 qkv_size = q_size + 2 * kv_dim    # 4608
 # Pre-allocated GPU (OffloadEngine)
 ring_buffer = 2 * num_kv_buffers * seq_len * kv_dim * dtype_size
 # = 2 * 4 * 131072 * 1024 * 2 = 2,147,483,648 bytes = 2048 MB
 decode_buffer = 2 * num_layers * block_size * kv_dim * dtype_size
 # = 2 * 36 * 1024 * 1024 * 2 = 150,994,944 bytes = 144 MB
 # Pre-allocated CPU
 cpu_cache = 2 * num_layers * num_cpu_blocks * block_size * kv_dim * dtype_size
 # = 2 * 36 * 128 * 1024 * 1024 * 2 = 19,327,352,832 bytes = 18432 MB
 # Prefill temporaries (per layer peak)
 mlp_gate_up = seq_len * 2 * intermediate_size * dtype_size
 # = 131072 * 2 * 13696 * 2 = 7,180,648,448 bytes = 6848 MB
 # Decode temporaries (per layer)
 k_full = seq_len * kv_dim * dtype_size
 # = 131072 * 1024 * 2 = 268,435,456 bytes = 256 MB
 v_full = k_full  # = 256 MB
 # Total: 512 MB
 ```
 ---
 ## 8. Empirical Validation
 This section validates the theoretical memory analysis against actual measurements.
 ### 8.1 Test Configuration
 ```bash
 python tests/test_needle.py --enable-offload --input-len 100000 --block-size 1024
 ```
 **Parameters:**
 - Model: Qwen3-4B-Instruct
 - `seq_len = 100000` (actual tokens: 99925)
 - `block_size = 1024`
 - `max_model_len = 131072`
 - `num_kv_buffers = 4`
 ### 8.2 Theoretical Peak Memory Calculation
 #### Step 1: Model Load Memory
 | Component | Formula | Size |
 |-----------|---------|------|
 | Model weights | ~4B params × 2 bytes | ~8 GB |
 | Ring buffer | 2 × 4 × 131072 × 1024 × 2 | 2048 MB |
 | Decode buffer | 2 × 36 × 1024 × 1024 × 2 | 144 MB |
 | **Subtotal** | | **~10.2 GB** |
 #### Step 2: Prefill Activation Peak (per-layer)
 | Component | Formula | Size |
 |-----------|---------|------|
 | hidden_states | 100000 × 2560 × 2 | 512 MB |
 | residual | 100000 × 2560 × 2 | 512 MB |
 | MLP gate_up | 100000 × 27392 × 2 | **5478 MB** |
 | MLP silu×gate | 100000 × 13696 × 2 | 2739 MB |
 | Other intermediates (qkv, RoPE, attn) | ~1-2 GB | ~1500 MB |
 | **Subtotal** | | **~10 GB** |
 #### Step 3: Total Peak
 ```
 Total Peak = Model Load + Activation Peak
           = 10.2 GB + 10 GB
           = ~20.2 GB
 ```
 ### 8.3 Actual Measurement Results
 ```python
 import torch
 torch.cuda.reset_peak_memory_stats()
 # ... run inference ...
 peak = torch.cuda.max_memory_allocated()
 ```
 | Metric | Value |
 |--------|-------|
 | After model load | 9.82 GB |
 | Peak during inference | **20.02 GB** |
 | Activation peak (delta) | 10.20 GB |
 ### 8.4 Comparison: Theory vs Actual
 | Component | Theoretical | Actual | Error |
 |-----------|-------------|--------|-------|
 | Model load memory | ~10.2 GB | 9.82 GB | -3.7% |
 | Activation peak | ~10 GB | 10.20 GB | +2.0% |
 | **Total peak** | **~20.2 GB** | **20.02 GB** | **< 1%** |
 ### 8.5 Key Findings
 1. **Theoretical model is accurate**: < 5% error in all components.
 2. **MLP gate_up is the dominant temporary**:
   - Size: 5.35 GB (for 100k tokens)
   - Accounts for ~50% of activation peak
   - Formula: `seq_len × 2 × intermediate_size × dtype_size`
 3. **Memory scaling with sequence length**:
   | seq_len | Model Load | Activation Peak | Total Peak |
   |---------|------------|-----------------|------------|
   | 8k | ~10 GB | ~0.8 GB | ~11 GB |
   | 32k | ~10 GB | ~3.2 GB | ~13 GB |
   | 64k | ~10 GB | ~6.4 GB | ~16 GB |
   | 100k | ~10 GB | ~10 GB | ~20 GB |
   | 128k | ~10 GB | ~13 GB | ~23 GB |
 4. **Decode memory is much smaller**:
   - Per-step: ~512 MB for k_full + v_full (at 100k context)
   - Does not grow with decode steps (constant per layer)
 ### 8.6 Memory Profiling Script
 To reproduce the measurement:
 ```python
 import os
 os.environ["NANOVLLM_LOG_LEVEL"] = "INFO"
 import torch
 from nanovllm import LLM, SamplingParams
 from tests.utils import generate_needle_prompt
 # Reset memory stats
 torch.cuda.reset_peak_memory_stats()
 torch.cuda.empty_cache()
 # Initialize LLM
 llm = LLM(
    "path/to/model",
    enforce_eager=True,
    max_model_len=131072,
    max_num_batched_tokens=131072,
    enable_cpu_offload=True,
    kvcache_block_size=1024,
    num_gpu_blocks=2,
 )
 after_load = torch.cuda.memory_allocated()
 print(f"After model load: {after_load / 1024**3:.2f} GB")
 # Generate prompt and run inference
 prompt, expected = generate_needle_prompt(
    tokenizer=llm.tokenizer,
    target_length=100000,
    needle_position=0.5,
 )
 torch.cuda.reset_peak_memory_stats()
 outputs = llm.generate([prompt], SamplingParams(max_tokens=32))
 peak = torch.cuda.max_memory_allocated()
 print(f"Peak during inference: {peak / 1024**3:.2f} GB")
 ```
--- a/docs/multi_model_support.md
+++ b/docs/multi_model_support.md
@@ -0,0 +1,233 @@
 # Multi-Model Support
 本文档描述 nanovllm 的多模型支持架构，以及如何添加新模型。
 ## 概述
 nanovllm 通过模型注册表 (Model Registry) 机制支持多种模型架构。系统根据 HuggingFace config 中的 `architectures` 字段自动选择对应的模型实现。
 ### 当前支持的模型
 | 架构 | 模型示例 | 文件 |
 |------|---------|------|
 | `Qwen3ForCausalLM` | Qwen3-0.6B, Qwen3-4B | `nanovllm/models/qwen3.py` |
 | `Qwen2ForCausalLM` | Qwen2.5-7B | `nanovllm/models/qwen3.py` |
 | `LlamaForCausalLM` | Llama-3.1-8B-Instruct | `nanovllm/models/llama.py` |
 ## 架构设计
 ### 模型注册表
 ```
 nanovllm/models/
 ├── __init__.py      # 导出 get_model_class, 导入所有模型
 ├── registry.py      # 注册表核心: MODEL_REGISTRY, @register_model
 ├── qwen3.py         # Qwen3/Qwen2 实现
 └── llama.py         # Llama 实现
 ```
 ### 动态模型加载流程
 ```
 LLM(model_path)
  → Config.__post_init__()
    → hf_config = AutoConfig.from_pretrained(model_path)
  → ModelRunner.__init__()
    → model_class = get_model_class(hf_config)  # 根据 architectures 选择
    → model = model_class(hf_config)
    → load_model(model, model_path)
 ```
 ## 添加新模型
 ### 步骤 1: 创建模型文件
 在 `nanovllm/models/` 下创建新文件，例如 `mistral.py`:
 ```python
 import torch
 from torch import nn
 import torch.distributed as dist
 from nanovllm.layers.activation import SiluAndMul
 from nanovllm.layers.attention import Attention
 from nanovllm.layers.layernorm import RMSNorm
 from nanovllm.layers.linear import QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear
 from nanovllm.layers.rotary_embedding import get_rope
 from nanovllm.layers.embed_head import VocabParallelEmbedding, ParallelLMHead
 from nanovllm.models.registry import register_model
 class MistralAttention(nn.Module):
    def __init__(self, ...):
        # 实现注意力层
        pass
 class MistralMLP(nn.Module):
    def __init__(self, ...):
        # 实现 MLP 层
        pass
 class MistralDecoderLayer(nn.Module):
    def __init__(self, config):
        # 组合 Attention + MLP
        pass
 class MistralModel(nn.Module):
    def __init__(self, config):
        # Embedding + Layers + Norm
        pass
@register_model("MistralForCausalLM")
 class MistralForCausalLM(nn.Module):
    # 权重映射 (HF 权重名 -> nanovllm 权重名)
    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),
    }
    def __init__(self, config):
        super().__init__()
        self.model = MistralModel(config)
        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
    def forward(self, input_ids, positions):
        return self.model(input_ids, positions)
    def compute_logits(self, hidden_states):
        return self.lm_head(hidden_states)
 ```
 ### 步骤 2: 注册模型
 在 `nanovllm/models/__init__.py` 中导入新模型:
 ```python
 from nanovllm.models import mistral  # 添加这行
 ```
 ### 步骤 3: 处理特殊配置
 如果模型有特殊的 RoPE scaling 或其他配置，需要在相应的 layer 中添加支持。
 ## 模型架构差异
 ### Qwen3 vs Llama
 | 特性 | Qwen3 | Llama |
 |------|-------|-------|
 | QKV Bias | 可配置 (`attention_bias`) | 无 |
 | Q/K Norm | 有 (RMSNorm, 当 bias=False) | 无 |
 | MLP Bias | 无 | 无 |
 | RoPE Scaling | 无 | llama3 类型 |
 | RoPE Theta | 1,000,000 | 500,000 |
 ### RoPE Scaling 支持
 目前支持的 RoPE 类型:
 | `rope_type` | 说明 | 模型 |
 |-------------|------|------|
 | `None` | 标准 RoPE | Qwen3 |
 | `llama3` | Llama 3 频率缩放 | Llama 3.1 |
 Llama3 RoPE 特点:
 - 低频分量 (长距离依赖): 缩放 1/factor
 - 高频分量 (短距离依赖): 保持不变
 - 中频分量: 平滑插值
 ## 权重加载
 ### packed_modules_mapping
 nanovllm 将多个 HuggingFace 权重合并到单个张量中以提高效率:
 ```python
 packed_modules_mapping = {
    # HF 权重名: (nanovllm 权重名, shard_id)
    "q_proj": ("qkv_proj", "q"),  # Q 投影 -> QKV 合并
    "k_proj": ("qkv_proj", "k"),  # K 投影 -> QKV 合并
    "v_proj": ("qkv_proj", "v"),  # V 投影 -> QKV 合并
    "gate_proj": ("gate_up_proj", 0),  # Gate -> Gate+Up 合并
    "up_proj": ("gate_up_proj", 1),    # Up -> Gate+Up 合并
 }
 ```
 ### 权重加载流程
 ```python
 # nanovllm/utils/loader.py
 def load_model(model, path):
    for file in glob(path + "/*.safetensors"):
        with safe_open(file) as f:
            for weight_name in f.keys():
                # 检查是否需要映射
                if weight_name in packed_modules_mapping:
                    # 使用自定义 weight_loader
                    param.weight_loader(param, tensor, shard_id)
                else:
                    # 直接复制
                    param.data.copy_(tensor)
 ```
 ## 测试验证
 ### Needle-in-Haystack 测试
 ```bash
 # Llama 3.1 (32K, offload 模式)
 CUDA_VISIBLE_DEVICES=0 python tests/test_needle.py \
    --model ~/models/Llama-3.1-8B-Instruct \
    --max-model-len 40960 \
    --input-len 32768 \
    --block-size 1024 \
    --num-gpu-blocks 4 \
    --enable-offload
 # Qwen3 (8K, offload 模式)
 CUDA_VISIBLE_DEVICES=0 python tests/test_needle.py \
    --model ~/models/Qwen3-4B-Instruct-2507 \
    --max-model-len 40960 \
    --input-len 8192 \
    --enable-offload
 ```
 ### 测试结果
 | 模型 | 输入长度 | Needle 位置 | 结果 |
 |------|---------|-------------|------|
 | Llama-3.1-8B | 32K | 50% | ✅ PASSED |
 | Llama-3.1-8B | 32K | 90% | ✅ PASSED |
 | Llama-3.1-8B | 32K | 10% | ❌ FAILED (Lost in Middle) |
 | Qwen3-4B | 8K | 50% | ✅ PASSED |
 ## 文件结构
 ```
 nanovllm/
 ├── models/
 │   ├── __init__.py         # 模型导出和导入
 │   ├── registry.py         # 注册表实现
 │   ├── qwen3.py           # Qwen3/Qwen2 模型
 │   └── llama.py           # Llama 模型
 ├── layers/
 │   ├── rotary_embedding.py # RoPE (含 Llama3 scaling)
 │   ├── attention.py        # FlashAttention wrapper
 │   ├── linear.py          # 并行 Linear 层
 │   └── ...
 └── engine/
    └── model_runner.py     # 动态模型加载
 ```
 ## 注意事项
 1. **Tokenizer 差异**: 不同模型的 tokenizer 分词策略不同，例如 Llama 将 "7492" 分为 2 tokens，Qwen3 分为 4 tokens。
 2. **RoPE Scaling**: 如果模型使用非标准 RoPE，需要在 `rotary_embedding.py` 中添加支持。
 3. **CPU Offload**: 在 3090 等显存有限的 GPU 上，使用 `--enable-offload` 进行长上下文测试。
 4. **Lost in Middle**: LLM 对开头信息的记忆能力较弱，这是模型本身的限制，不是实现问题。
--- a/docs/offload_accuracy_issue.md
+++ b/docs/offload_accuracy_issue.md
@@ -0,0 +1,306 @@
 # CPU Offload Accuracy Issue Investigation
 ## Problem Summary
 **UPDATE (2026-01-12)**: Single request inference works correctly! The issue is with batch/sequential request handling.
 | Mode | Testing Method | Accuracy |
 |------|----------------|----------|
 | **CPU Offload** | **Independent** (1 request per process) | **100%** ✓ |
 | **CPU Offload** | Batch (multiple requests per process) | 66% ✗ |
 | **Non-Offload** | Batch | 100% ✓ |
 **Conclusion**: The offload implementation is correct for single requests. The bug is in state cleanup between sequential requests within the same process.
 ## Test Environment
 - **Model**: Llama-3.1-8B-Instruct
 - **Task**: RULER NIAH (Needle-In-A-Haystack) 32K context
 - **GPU**: NVIDIA A100-SXM4-80GB
 - **Data**: `tests/data/ruler_niah/niah_single_1_32k.jsonl` (100 samples)
 ## Reproduction Commands
 ### Non-Offload Mode (100% accuracy)
 ```bash
 CUDA_VISIBLE_DEVICES=0 PYTHONPATH=.:$PYTHONPATH python tests/test_ruler_niah.py \
    --model ~/models/Llama-3.1-8B-Instruct \
    --gpu-utilization 0.7 \
    --quiet
 ```
 **Configuration**:
 - KV Cache: GPU only, 51 blocks (6528 MB)
 - Block size: 1024 tokens
 ### Offload Mode (66% accuracy)
 ```bash
 CUDA_VISIBLE_DEVICES=0 PYTHONPATH=.:$PYTHONPATH python tests/test_ruler_niah.py \
    --model ~/models/Llama-3.1-8B-Instruct \
    --enable-offload \
    --quiet
 ```
 **Configuration**:
 - KV Cache: GPU 4 blocks (512 MB) + CPU 32 blocks (4096 MB)
 - Ring buffer: 4 buffers × 33280 tokens (520 MB)
 - Per-layer decode buffer: 128 MB
 - Block size: 1024 tokens
 ## Observed Failure Patterns
 From the 5-sample verbose test:
 | Sample | Expected | Offload Output | Status |
 |--------|----------|----------------|--------|
 | 0 | 8930103 | `: 8930103.` | PASS |
 | 1 | 4194548 | `: 419 multiplication of 4548.` | **FAIL** |
 | 2 | 8231838 | `:ное 8231838.` | PASS |
 | 3 | 8835373 | `: 8835373.` | PASS |
 | 4 | 7754864 | `aster 7754864.` | PASS |
 **Failure pattern**: The model sometimes produces corrupted or split outputs (e.g., "419 multiplication of 4548" instead of "4194548").
 ## Architecture Overview
 ### Offload Mode Data Flow
 ```
 Prefill Phase:
 1. Input tokens → chunked into 2048-token chunks
 2. Each chunk processed layer by layer:
   - Load KV from CPU → GPU ring buffer
   - Compute attention
   - Store KV back to CPU
 3. Ring buffer holds recent KV for decode
 Decode Phase:
 1. For each new token:
   - Load all layer KV from CPU (one layer at a time)
   - Compute attention against full context
   - Generate next token
 ```
 ### Key Components
 | File | Component | Description |
 |------|-----------|-------------|
 | `nanovllm/kvcache/offload_engine.py` | `OffloadEngine` | Manages CPU↔GPU KV cache transfers |
 | `nanovllm/kvcache/offload_engine.py` | `RingKVBuffer` | GPU ring buffer for recent KV |
 | `nanovllm/engine/model_runner.py` | `run_chunked_offload_prefill()` | Chunked prefill with offload |
 | `nanovllm/engine/model_runner.py` | `run_offload_decode()` | Layer-wise decode with offload |
 | `nanovllm/kvcache/hybrid_manager.py` | `HybridBlockManager` | CPU block allocation |
 ## Potential Root Causes
 ### 1. Ring Buffer Index/Position Issues
 **Location**: `nanovllm/kvcache/offload_engine.py`
 The ring buffer uses modular indexing. Potential issues:
 - Position calculation errors during prefill/decode transition
 - Off-by-one errors in KV storage/retrieval
 - Incorrect handling when sequence length approaches `max_seq_len`
 **Recent fix applied**: `max_seq_len = max_model_len + 512` to prevent overflow, but there may be other indexing issues.
 ### 2. Chunked Prefill KV Storage
 **Location**: `nanovllm/engine/model_runner.py:run_chunked_offload_prefill()`
 During chunked prefill:
 - KV computed for chunk N must be correctly stored before processing chunk N+1
 - Position IDs must be correctly accumulated across chunks
 - CPU block allocation must be contiguous and correctly tracked
 **Suspect areas**:
 ```python
 # Check if positions are correctly tracked across chunks
 # Check if KV is correctly copied to CPU after each chunk
 # Check if ring buffer indices align with CPU block indices
 ```
 ### 3. Decode Phase KV Loading
 **Location**: `nanovllm/engine/model_runner.py:run_offload_decode()`
 During decode:
 - Must load KV for ALL previous tokens (both prefill and decode)
 - Layer-by-layer loading must be synchronized correctly
 - Attention computation must use correct sequence length
 **Suspect areas**:
 ```python
 # Check if decode loads KV for full context length
 # Check if new decode KV is stored correctly
 # Check if attention mask/positions are correct
 ```
 ### 4. CPU↔GPU Transfer Synchronization
 **Location**: `nanovllm/kvcache/offload_engine.py`
 CUDA streams and synchronization:
 - Async copies may complete out of order
 - Missing synchronization points could cause stale data
 - Stream priorities may affect correctness
 ### 5. Numerical Precision
 - CPU tensors use float16/bfloat16
 - GPU computation precision
 - Potential precision loss during transfers
 ## Debugging Strategy
 ### Step 1: Identify Failing Samples
 ```bash
 # Run verbose mode to see which samples fail
 CUDA_VISIBLE_DEVICES=0 PYTHONPATH=.:$PYTHONPATH python tests/test_ruler_niah.py \
    --model ~/models/Llama-3.1-8B-Instruct \
    --enable-offload \
    --verbose 2>&1 | tee offload_verbose.log
 ```
 ### Step 2: Compare Token-by-Token
 Create a debug script to compare token generation between offload and non-offload modes for a failing sample:
 ```python
 # Compare logits at each decode step
 # Check if divergence starts at a specific position
 # Log KV cache contents at divergence point
 ```
 ### Step 3: Verify KV Cache Contents
 Add debugging to `OffloadEngine`:
 ```python
 # In store_kv(): Log what's being stored
 # In load_kv(): Log what's being loaded
 # Compare loaded KV with expected values
 ```
 ### Step 4: Check Position/Index Calculations
 ```python
 # Log ring buffer write/read positions
 # Log CPU block indices
 # Verify position IDs match actual token positions
 ```
 ### Step 5: Isolate the Bug
 1. Test with shorter sequences (16K, 8K) to see if issue is length-dependent
 2. Test with single chunk (no chunking) to isolate chunked prefill
 3. Test prefill-only (no decode) to isolate decode phase
 ## Quick Debugging Commands
 ```bash
 # Test single failing sample with verbose output
 CUDA_VISIBLE_DEVICES=0 PYTHONPATH=.:$PYTHONPATH python tests/test_ruler_niah.py \
    --model ~/models/Llama-3.1-8B-Instruct \
    --enable-offload \
    --sample-indices 1 \
    --verbose
 # Test with different context lengths
 CUDA_VISIBLE_DEVICES=0 PYTHONPATH=.:$PYTHONPATH python tests/test_ruler_niah.py \
    --model ~/models/Llama-3.1-8B-Instruct \
    --enable-offload \
    --max-model-len 16384 \
    --verbose
 ```
 ## Related Documentation
 - [`docs/ruler_niah_standalone_test.md`](ruler_niah_standalone_test.md) - Test setup and background
 - [`docs/layerwise_offload_memory_analysis.md`](layerwise_offload_memory_analysis.md) - Memory analysis (if exists)
 ## Test Results Log
 ### 2026-01-12 (Updated - Independent Testing)
 **Key Finding**: When each sample is tested independently (separate Python process per sample), CPU offload achieves **100% accuracy**.
 | Test | Mode | Testing Method | Samples | Passed | Accuracy |
 |------|------|----------------|---------|--------|----------|
 | RULER NIAH 32K | CPU Offload | **Independent** (separate process) | 100 | 100 | **100%** |
 | RULER NIAH 32K | CPU Offload | Batch (single process) | 100 | 66 | 66% |
 | RULER NIAH 32K | Non-Offload | Batch (single process) | 100 | 100 | 100% |
 **Test Configuration (Independent Mode)**:
 - GPUs: 4x RTX 3090 (parallel testing)
 - Each sample: Fresh Python process with new LLM instance
 - Port: Each GPU uses unique port (2333+gpu_id)
 - Duration: 17.9 minutes for 100 samples
 - Throughput: 5.58 samples/min
 ### 2025-01-12 (Original - Batch Testing)
 | Test | Mode | Samples | Passed | Accuracy |
 |------|------|---------|--------|----------|
 | RULER NIAH 32K | Non-Offload | 100 | 100 | 100% |
 | RULER NIAH 32K | CPU Offload | 100 | 66 | 66% |
 ## Root Cause Analysis Update
 ### Confirmed: Single Request Inference is Correct
 The 100% accuracy in independent testing mode confirms that:
 1. **Single request inference works correctly** - The offload engine, ring buffer, and chunked prefill are functioning properly for individual requests
 2. **The bug is in batch/sequential request handling** - State accumulation or incomplete cleanup between requests causes failures
 ### Suspected Issue: State Accumulation Between Requests
 When multiple requests are processed in the same Python process:
 - The first request succeeds (e.g., Sample 0: PASS)
 - Subsequent requests may fail due to:
  - Residual state in ring buffer
  - Incomplete KV cache cleanup
  - Position tracking errors across requests
  - CPU block allocation fragmentation
 ### Evidence
 From batch mode testing (5 samples):
 | Sample | Expected | Output | Status |
 |--------|----------|--------|--------|
 | 0 | 8930103 | `: 8930103.` | PASS (first request) |
 | 1 | 4194548 | `: 419 multiplication of 4548.` | **FAIL** (second request) |
 | 2 | 8231838 | `:ное 8231838.` | PASS |
 | 3 | 8835373 | `: 8835373.` | PASS |
 | 4 | 7754864 | `aster 7754864.` | PASS |
 The corrupted output in Sample 1 suggests interference from Sample 0's state.
 ## Workaround
 Use independent testing mode (separate process per request) for production evaluation:
 ```bash
 # Using test_ruler_niah.sh for parallel independent testing
 ./tests/test_ruler_niah.sh --gpus "0,1,2,3" --total 100
 # Or manually run each sample in a separate process
 for i in $(seq 0 99); do
    CUDA_VISIBLE_DEVICES=0 python tests/test_ruler_niah.py \
        --enable-offload --sample-indices $i --quiet
 done
 ```
 ## Next Steps
 1. [x] ~~Identify pattern in failing samples~~ → Pattern: First sample usually passes, failures occur in subsequent samples
 2. [ ] **Investigate state cleanup between requests in offload mode**
   - Check `OffloadEngine` reset/cleanup logic
   - Check ring buffer state between requests
   - Check CPU block manager cleanup
 3. [ ] Add `reset()` method to `OffloadEngine` for explicit state cleanup
 4. [ ] Compare state between first and second request in batch mode
 5. [ ] Write unit test that reproduces the batch mode failure
--- a/docs/ruler_benchmark_report.md
+++ b/docs/ruler_benchmark_report.md
@@ -0,0 +1,99 @@
 # RULER Benchmark 测试报告
 **测试日期**: 2025-01-14
 **测试环境**: 6x RTX 3090, CPU Offload 模式
 **模型**: Llama-3.1-8B-Instruct
 **上下文长度**: 32K tokens
 ## 测试概述
 使用 RULER benchmark 对 nano-vllm 的 CPU offload 模式进行全面的长上下文能力测试。RULER 是 NVIDIA 开发的长上下文评测基准，包含 13 个任务类别。
 ## 测试结果
 ### 总体结果
 | 类别 | 数据集 | 正确/总数 | 准确率 | 平均分数 |
 |------|--------|-----------|--------|----------|
 | **NIAH Single** | niah_single_1 | 100/100 | 100.0% | 1.000 |
 | | niah_single_2 | 100/100 | 100.0% | 1.000 |
 | | niah_single_3 | 100/100 | 100.0% | 1.000 |
 | **NIAH MultiKey** | niah_multikey_1 | 100/100 | 100.0% | 1.000 |
 | | niah_multikey_2 | 90/100 | 90.0% | 0.900 |
 | | niah_multikey_3 | 93/100 | 93.0% | 0.930 |
 | **NIAH Other** | niah_multiquery | 100/100 | 100.0% | 1.000 |
 | | niah_multivalue | 100/100 | 100.0% | 1.000 |
 | **QA** | qa_1 | 79/100 | 79.0% | 0.790 |
 | | qa_2 | 51/100 | 51.0% | 0.510 |
 | **Aggregation** | cwe | 86/100 | 86.0% | 0.680 |
 | | fwe | 98/100 | 98.0% | 0.923 |
 | **Variable Tracking** | vt | 100/100 | 100.0% | 0.934 |
 | **总计** | **13 数据集** | **1197/1300** | **92.1%** | **0.897** |
 ### 分类性能分析
 | 任务类别 | 描述 | 准确率 | 评价 |
 |----------|------|--------|------|
 | NIAH Single | 单 needle 检索 | 100% | 优秀 |
 | NIAH MultiKey | 多 key 检索 | 94.3% | 良好 |
 | NIAH MultiQuery/Value | 复杂检索 | 100% | 优秀 |
 | QA | 问答理解 | 65% | 一般 |
 | Aggregation (CWE/FWE) | 信息聚合 | 92% | 良好 |
 | Variable Tracking | 变量追踪 | 100% | 优秀 |
 ## 发现的问题及修复
 ### 问题: FWE 测试崩溃
 **症状**: 第 63 个样本处触发 `AssertionError: No sequences scheduled`
 **根因分析**:
 1. Sample 63 的输入有 32760 tokens（接近 max_model_len=32768）
 2. Decode 到第 9 步时，需要第 33 个 KV block
 3. 但系统只配置了 32 个 blocks（32768/1024=32）
 4. 调度器尝试 preempt 但单序列模式下无法恢复
 **解决方案**:
 ```python
 # 修改前
 DEFAULT_MAX_MODEL_LEN = 32768
 # 修改后: 为 output tokens 预留空间
 DEFAULT_MAX_MODEL_LEN = 32896  # 32768 + 128
 ```
 **建议的代码改进**:
 1. 在 scheduler 中添加死锁检测和清晰错误信息
 2. 在配置验证时，如果 max_model_len 与 max_input 过于接近，发出警告
 ## 评估方法
 遵循 RULER 官方评估标准:
 - **NIAH/VT/CWE/FWE**: `string_match_all` - 召回率 (找到的参考数/总参考数)
 - **QA**: `string_match_part` - 任意参考匹配即满分
 参考: https://github.com/NVIDIA/RULER
 ## 测试配置
 ```python
 LLM(
    model_path="~/models/Llama-3.1-8B-Instruct",
    max_model_len=32896,
    max_num_batched_tokens=32896,
    enable_cpu_offload=True,
    num_gpu_blocks=4,
    kvcache_block_size=1024,
    enforce_eager=True,
 )
 ```
 ## 结论
 1. **长上下文检索能力**: nano-vllm CPU offload 模式在 32K 上下文下表现优秀，NIAH 类任务准确率接近 100%
 2. **复杂推理能力**: QA 任务准确率较低 (65%)，这是模型本身能力的体现，与 offload 机制无关
 3. **稳定性**: 修复 max_model_len 配置后，所有 1300 个样本测试均稳定完成
 4. **性能**: 单样本测试时间约 25-35 秒，主要受 CPU-GPU 数据传输影响
--- a/docs/ruler_niah_standalone_test.md
+++ b/docs/ruler_niah_standalone_test.md
@@ -0,0 +1,297 @@
 # RULER NIAH Standalone Test Plan
 ## Overview
 This document describes how to independently test nano-vllm's CPU offload functionality using RULER benchmark's NIAH (Needle-In-A-Haystack) task data.
 ## Background
 ### Problem Being Investigated
 When running 32K sequence length tests with CPU offload mode, the model outputs garbled text instead of finding the magic number. This issue was traced to:
 - **Root Cause**: Ring buffer `max_seq_len` was set equal to `max_model_len` (32768)
 - **Issue**: When prefill uses ~32K tokens, decode needs to store KV at position 32768+, but ring buffer only has indices 0-32767
 - **Fix Applied**: In `nanovllm/kvcache/__init__.py`, changed `max_seq_len = max_model_len + 512`
 ### Test Objective
 Verify that the fix works correctly by running a standalone test with actual RULER NIAH data.
 ## Step 1: Copy Test Data
 ### Source Location
 ```
 /home/zijie/Code/x-attention/eval/RULER/scripts/benchmark_root/full_fuse_16_llama3.1-8b-chat/synthetic/32768/data/niah_single_1/validation.jsonl
 ```
 ### Data Format
 Each line is a JSON object:
 ```json
 {
  "index": 0,
  "input": "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\nA special magic number is hidden within the following text...",
  "outputs": ["8930103"],
  "length": 32768
 }
 ```
 - `input`: Full prompt with Llama 3.1 chat template (~122K characters, ~30K tokens)
 - `outputs`: Expected answer (the magic number to find)
 - `length`: Target sequence length in tokens
 ### Copy Command
 ```bash
 mkdir -p /home/zijie/Code/nano-vllm/tests/data/ruler_niah
 cp /home/zijie/Code/x-attention/eval/RULER/scripts/benchmark_root/full_fuse_16_llama3.1-8b-chat/synthetic/32768/data/niah_single_1/validation.jsonl \
   /home/zijie/Code/nano-vllm/tests/data/ruler_niah/niah_single_1_32k.jsonl
 ```
 ## Step 2: Create Test Script
 Create `/home/zijie/Code/nano-vllm/tests/test_ruler_niah_32k.py`:
 ```python
 """
 Standalone test for RULER NIAH task with 32K context length.
 This test verifies that CPU offload mode correctly handles long sequences
 where prefill tokens approach max_model_len.
 Usage:
    python tests/test_ruler_niah_32k.py
 """
 import json
 import torch
 from pathlib import Path
 from nanovllm import LLM
 from nanovllm.config import SamplingParams
 # Configuration
 MODEL_PATH = "/data/models/Llama-3.1-8B-Instruct"
 DATA_FILE = Path(__file__).parent / "data/ruler_niah/niah_single_1_32k.jsonl"
 MAX_MODEL_LEN = 32768
 MAX_NEW_TOKENS = 50
 # CPU Offload Settings
 ENABLE_CPU_OFFLOAD = True
 NUM_GPU_BLOCKS = 4
 BLOCK_SIZE = 1024
 def load_test_sample(filepath: Path, index: int = 0) -> dict:
    """Load a single test sample from JSONL file."""
    with open(filepath) as f:
        for i, line in enumerate(f):
            if i == index:
                return json.loads(line)
    raise ValueError(f"Sample index {index} not found")
 def test_niah_single():
    """Test NIAH single needle task with 32K context."""
    print("=" * 60)
    print("RULER NIAH 32K Standalone Test")
    print("=" * 60)
    # Load test data
    sample = load_test_sample(DATA_FILE, index=0)
    prompt = sample["input"]
    expected = sample["outputs"][0]
    print(f"Prompt length: {len(prompt)} characters")
    print(f"Expected answer: {expected}")
    print()
    # Initialize model with CPU offload
    print("Initializing LLM with CPU offload...")
    llm = LLM(
        model=MODEL_PATH,
        max_model_len=MAX_MODEL_LEN,
        enable_cpu_offload=ENABLE_CPU_OFFLOAD,
        num_gpu_blocks=NUM_GPU_BLOCKS,
        kvcache_block_size=BLOCK_SIZE,
        enforce_eager=True,  # Disable CUDA graphs for debugging
    )
    # Generate
    print("Generating response...")
    sampling_params = SamplingParams(
        temperature=0.0,  # Greedy
        max_tokens=MAX_NEW_TOKENS,
    )
    outputs = llm.generate([prompt], sampling_params)
    generated_text = outputs[0].outputs[0].text
    print()
    print("=" * 60)
    print("Results")
    print("=" * 60)
    print(f"Expected: {expected}")
    print(f"Generated: {generated_text[:200]}...")
    print()
    # Check if expected number is in output
    if expected in generated_text:
        print("SUCCESS: Magic number found in output!")
        return True
    else:
        print("FAILED: Magic number NOT found in output")
        print(f"Full output: {generated_text}")
        return False
 def test_multiple_samples(num_samples: int = 5):
    """Test multiple NIAH samples."""
    print("=" * 60)
    print(f"Testing {num_samples} NIAH samples with 32K context")
    print("=" * 60)
    # Initialize model once
    llm = LLM(
        model=MODEL_PATH,
        max_model_len=MAX_MODEL_LEN,
        enable_cpu_offload=ENABLE_CPU_OFFLOAD,
        num_gpu_blocks=NUM_GPU_BLOCKS,
        kvcache_block_size=BLOCK_SIZE,
        enforce_eager=True,
    )
    sampling_params = SamplingParams(
        temperature=0.0,
        max_tokens=MAX_NEW_TOKENS,
    )
    correct = 0
    for i in range(num_samples):
        sample = load_test_sample(DATA_FILE, index=i)
        prompt = sample["input"]
        expected = sample["outputs"][0]
        outputs = llm.generate([prompt], sampling_params)
        generated_text = outputs[0].outputs[0].text
        if expected in generated_text:
            print(f"Sample {i}: PASS (found {expected})")
            correct += 1
        else:
            print(f"Sample {i}: FAIL (expected {expected}, got: {generated_text[:50]}...)")
    print()
    print(f"Accuracy: {correct}/{num_samples} ({100*correct/num_samples:.1f}%)")
    return correct == num_samples
 if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "--all":
        success = test_multiple_samples(5)
    else:
        success = test_niah_single()
    sys.exit(0 if success else 1)
 ```
 ## Step 3: Run Test
 ### Single Sample Test
 ```bash
 cd /home/zijie/Code/nano-vllm
 CUDA_VISIBLE_DEVICES=2,3,4,5 python tests/test_ruler_niah_32k.py
 ```
 ### All 5 Samples
 ```bash
 cd /home/zijie/Code/nano-vllm
 CUDA_VISIBLE_DEVICES=2,3,4,5 python tests/test_ruler_niah_32k.py --all
 ```
 ## Step 4: Expected Results
 ### Before Fix (Bug)
 - Output: Garbled text like "not only has been replaced by thesiums..."
 - Score: 0% (magic number not found)
 - Time: ~80 seconds per sample
 ### After Fix (Expected)
 - Output: The magic number (e.g., "8930103")
 - Score: ~100% (magic number found)
 - Time: ~80 seconds per sample (same, as the compute is unchanged)
 ## Debugging Tips
 ### Enable Verbose Logging
 ```python
 import logging
 logging.basicConfig(level=logging.DEBUG)
 ```
 ### Check Ring Buffer Size
 In the logs, verify:
 ```
 OffloadEngine initializing: num_layers=32, num_kv_buffers=4, max_seq_len=33280
 ```
 The `max_seq_len` should be `32768 + 512 = 33280` (not 32768).
 ### Monitor GPU Memory
 ```bash
 watch -n 1 nvidia-smi
 ```
 With CPU offload, GPU memory for KV cache should be ~640MB (ring buffer only).
 ## Related Files
 | File | Description |
 |------|-------------|
 | `nanovllm/kvcache/__init__.py` | Fix location: `max_seq_len = max_model_len + 512` |
 | `nanovllm/kvcache/offload_engine.py` | Ring buffer allocation |
 | `nanovllm/engine/model_runner.py` | Layer-wise offload prefill/decode |
 | `nanovllm/kvcache/hybrid_manager.py` | CPU block management |
 ## Test Data Details
 ### NIAH Task Description
 The NIAH (Needle-In-A-Haystack) task tests the model's ability to retrieve a specific piece of information (the "needle") from a large context (the "haystack").
 - **Needle**: A magic number associated with a keyword (e.g., "worried-purse")
 - **Haystack**: ~30K tokens of distractor text
 - **Task**: Extract the magic number when asked
 ### Sample Prompt Structure
 ```
 <|begin_of_text|><|start_header_id|>user<|end_header_id|>
 A special magic number is hidden within the following text. Make sure to memorize it. I will quiz you about the number afterwards.
 [... ~30K tokens of haystack text ...]
 The special magic number for worried-purse is 8930103.
 [... more haystack text ...]
 What is the special magic number for worried-purse mentioned in the provided text?
 <|eot_id|><|start_header_id|>assistant<|end_header_id|>
 The special magic number for worried-purse mentioned in the provided text is
 ```
 The model should complete with: `8930103`
--- a/docs/sparse_attention_guide.md
+++ b/docs/sparse_attention_guide.md
@@ -440,3 +440,42 @@ Required libraries:
 - `minference`: For MInference vertical_slash kernel
 Docker image `tzj/xattn:v0.5` has all dependencies pre-installed.
 ---
 ## Quest Sparse Policy (nano-vLLM)
 **Files**: `nanovllm/kvcache/sparse/quest.py`, `nanovllm/kvcache/sparse/policy.py`
 Quest policy is used in nano-vLLM for CPU offload mode. It selects Top-K blocks based on query-key similarity bounds using min/max key metadata.
 ### Scoring Mechanism
 ```python
 score_min = torch.einsum('hd,bhd->bh', q, key_min)  # [num_blocks, kv_heads]
 score_max = torch.einsum('hd,bhd->bh', q, key_max)  # [num_blocks, kv_heads]
 scores = torch.maximum(score_min, score_max).mean(dim=-1)  # [num_blocks] ← averaged!
 ```
 ### Critical Limitation - No Per-Head Scheduling
 The `.mean(dim=-1)` averages scores across all heads, making a **unified** block selection for all heads:
 ```
 Block A: head0 needs (+4), head1 doesn't (-4) → avg = 0 → NOT selected
 Block B: head0 doesn't (-4), head1 needs (+4) → avg = 0 → NOT selected
 Block C: both heads moderately need (+2, +2) → avg = +2 → selected
 ```
 ### Why Per-Head Scheduling is Infeasible
 1. **Memory Layout**: GPU cache stores all heads together `[block_size, kv_heads, head_dim]`
 2. **FlashAttention**: Requires complete heads - partial heads cause dimension mismatch
 3. **Block Granularity**: If any head needs a block, the entire block (all heads) must be loaded
 ### Policy Types
 | Policy | `supports_prefill` | `supports_decode` | Description |
 |--------|-------------------|-------------------|-------------|
 | `FullAttentionPolicy` | True | True | Loads all blocks (baseline) |
 | `QuestPolicy` | False | True | Decode-only Top-K selection |
--- a/docs/sparse_offload_integration.md
+++ b/docs/sparse_offload_integration.md
@@ -0,0 +1,386 @@
 # Sparse Policy Integration with Layerwise Offload
 This document describes the architecture and design of integrating sparse attention policies (MInference, Quest) with the layerwise CPU offload execution path.
 ## Design Goals
 1. **Extend sparse policies to offload path**: GPU-only path already supports sparse policies, but layerwise offload bypasses them
 2. **Maintain encapsulation**: All `copy_()` operations must be inside OffloadEngine, not exposed to model_runner
 3. **Distinguish policy types**: Some policies affect attention computation (MInference), others affect KV load strategy (Quest)
 4. **Extensible architecture**: Easy to add new sparse policies in the future
 ## Key Insight
 The existing sparse policy implementation works, but the layerwise offload path bypasses it:
 | Path | Attention Method | Sparse Support |
 |------|------------------|----------------|
 | GPU-only | `attention.py` → `sparse_prefill_attention()` | YES |
 | Layerwise offload | `model_runner.py` → `flash_attn_varlen_func()` | NO (direct call) |
 ## Two Types of Sparse Policies
 The fundamental difference between sparse policies:
 | Policy | Affects Attention Computation | Affects KV Load Strategy | `select_blocks()` Behavior |
 |--------|------------------------------|--------------------------|---------------------------|
 | **MInference** | YES (`sparse_prefill_attention`) | NO | `return available_blocks` (all) |
 | **Quest** | NO | YES | Returns Top-K subset |
 - **MInference**: Only changes how attention is computed, doesn't affect external load/offload flow
 - **Quest**: Selectively loads only some blocks, affects H2D transfer
 ## The `requires_block_selection` Interface Flag
 To distinguish these policy types, we add a flag to the base class:
 ```python
 # nanovllm/kvcache/sparse/policy.py
 class SparsePolicy(ABC):
    # Existing flags
    supports_prefill: bool = True
    supports_decode: bool = True
    # 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
 ```
 ### Policy Implementations
 ```python
 # MInference: prefill-only, no block selection
 class MInferencePolicy(SparsePolicy):
    supports_prefill = True
    supports_decode = False
    requires_block_selection = False  # Only affects attention computation
 # Quest: decode-only, requires block selection
 class QuestPolicy(SparsePolicy):
    supports_prefill = False
    supports_decode = True
    requires_block_selection = True  # Affects KV load strategy
 # Full attention: baseline
 class FullAttentionPolicy(SparsePolicy):
    supports_prefill = True
    supports_decode = True
    requires_block_selection = False  # Load all blocks
 ```
 ## OffloadEngine Encapsulation
 All KV cache operations are encapsulated in OffloadEngine. The model_runner never directly accesses internal storage.
 ### Prefill: Synchronous Offload with Hooks
 ```python
 # nanovllm/kvcache/offload_engine.py
 def offload_layer_kv_sync(
    self,
    layer_id: int,
    k: Tensor,
    v: Tensor,
    cpu_block_ids: List[int],
    total_tokens: int,
 ) -> 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
        # 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])
 ```
 ### Decode: Policy-Driven Block Loading
 ```python
 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.
    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,
        )
        # Select blocks using policy
        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
        )
    else:
        # Load all blocks (no selection)
        return self._load_blocks_to_buffer(
            buffer_idx, layer_id, cpu_block_ids, valid_tokens_per_block
        )
 ```
 ## Prefill Integration (MInference)
 MInference only affects attention computation, not the load/offload flow:
 ```python
 # nanovllm/engine/model_runner.py - run_layerwise_offload_prefill()
 def run_layerwise_offload_prefill(self, seqs):
    ...
    for layer_id in range(num_layers):
        # QKV projection + RoPE
        q, k = layer.self_attn.rotary_emb(positions, q, k)
        # Sparse or Full attention
        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:
            # Full attention using FlashAttention
            attn_output = flash_attn_varlen_func(q, k, v, ...)
        # MLP
        ...
        # Offload ALL KV (MInference doesn't affect this)
        offload_engine.offload_layer_kv_sync(layer_id, k, v, cpu_block_ids, total_tokens)
 ```
 ### Execution Flow Diagram
 ```
 ┌─────────────────────────────────────────────────────────────────┐
 │                    Layerwise Offload Prefill                     │
 │                      with MInference                             │
 └─────────────────────────────────────────────────────────────────┘
 For each layer:
 ┌──────────────┐    ┌──────────────┐    ┌────────────────────────┐
 │ QKV Proj     │───▶│ RoPE         │───▶│ sparse_prefill_attn()  │
 │              │    │              │    │ (MInference pattern)   │
 └──────────────┘    └──────────────┘    └───────────┬────────────┘
                                                    │
                    ┌──────────────┐    ┌───────────▼────────────┐
                    │ MLP          │◀───│ O Projection           │
                    │              │    │                        │
                    └──────┬───────┘    └────────────────────────┘
                           │
                    ┌──────▼───────┐
                    │ offload_     │    K, V still on GPU
                    │ layer_kv_    │───▶ Copy to CPU
                    │ sync()       │    (all blocks)
                    └──────────────┘
 ```
 ## Decode Integration (Quest - Infrastructure Ready)
 Quest affects block load strategy. The infrastructure is ready, full integration deferred.
 ```python
 # nanovllm/engine/model_runner.py - run_layerwise_offload_decode()
 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(
            i, i, cpu_block_table, valid_tokens_per_block
        )
    for layer_id in range(num_layers):
        current_buffer = layer_id % num_buffers
        # Wait for buffer load
        offload_engine.wait_buffer_load(current_buffer)
        # QKV projection
        q, k_new, v_new = ...
        # Get loaded KV from ring buffer
        k_prefill, v_prefill = offload_engine.get_buffer_kv(
            current_buffer, loaded_tokens[current_buffer]
        )
        # Attention
        ...
        # Mark buffer done
        offload_engine.record_buffer_compute_done(current_buffer)
        # Load next layer
        # Future: use load_layer_kv_to_buffer_with_policy(query=q) for Quest
        next_layer = layer_id + num_buffers
        if next_layer < num_layers:
            loaded_tokens[current_buffer] = offload_engine.load_layer_kv_to_buffer(
                current_buffer, next_layer, cpu_block_table, valid_tokens_per_block
            )
 ```
 ### Quest Integration (Future Work)
 When Quest is fully integrated:
 ```python
 # Load next layer with Quest block selection
 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
    )
 ```
 **Challenge**: First N layers are preloaded before query is available, so they must use full load.
 ## Configuration
 ### Enabling Sparse Policy
 ```python
 from nanovllm import LLM
 from nanovllm.config import SparsePolicyType
 # GPU-only with MInference
 llm = LLM(
    model_path,
    sparse_policy=SparsePolicyType.MINFERENCE,
    minference_adaptive_budget=0.3,  # 30% of seq_len
 )
 # Offload with MInference
 llm = LLM(
    model_path,
    enable_cpu_offload=True,
    num_gpu_blocks=2,
    sparse_policy=SparsePolicyType.MINFERENCE,
    minference_adaptive_budget=0.3,
 )
 ```
 ### MInference Parameters
 | Parameter | Default | Description |
 |-----------|---------|-------------|
 | `minference_adaptive_budget` | 0.3 | Budget as fraction of seq_len (0.3 = 30%) |
 | `minference_vertical_size` | 1000 | Fixed vertical size (when budget=None) |
 | `minference_slash_size` | 6096 | Fixed slash size (when budget=None) |
 | `minference_num_sink_tokens` | 30 | Always-kept initial tokens |
 | `minference_num_recent_diags` | 100 | Always-kept recent diagonals |
 ### Quest Parameters (for future decode integration)
 | Parameter | Default | Description |
 |-----------|---------|-------------|
 | `sparse_topk_blocks` | 8 | Top-K blocks to load |
 | `sparse_threshold_blocks` | 4 | Apply sparse only when blocks > threshold |
 ## Sparse Policy Hooks
 Sparse policies can implement hooks for metadata collection:
 ```python
 class SparsePolicy(ABC):
    def on_prefill_offload(
        self,
        block_id: int,
        layer_id: int,
        key: torch.Tensor,
        valid_tokens: int,
    ) -> None:
        """
        Hook called during prefill offload BEFORE KV is copied to CPU.
        Key tensor is still on GPU - can compute metadata efficiently.
        Used by Quest to compute min/max key statistics for block selection.
        """
        pass
    def on_decode_offload(
        self,
        block_id: int,
        keys: torch.Tensor,  # [num_layers, block_size, kv_heads, head_dim]
    ) -> None:
        """
        Hook called when decode buffer is offloaded to CPU.
        """
        pass
 ```
 ## File Changes Summary
 | 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()`, sparse hooks |
 | `nanovllm/engine/model_runner.py` | Integrate sparse policies in offload paths |
 ## Key Design Principles
 1. **Encapsulation**: All `copy_()` operations inside OffloadEngine
 2. **Interface Flag**: `requires_block_selection` declares policy type
 3. **Separation of Concerns**:
   - MInference: only `sparse_prefill_attention()` (compute-level)
   - Quest: `select_blocks()` + hooks (load-level)
 4. **Hooks Inside Engine**: Policy hooks called within OffloadEngine methods
 ## Test Results
 Verified on Qwen3-4B-Instruct-2507 with 32K input:
 ```
 # GPU-only + MInference
 test_needle.py --model Qwen3-4B --input-len 32768 --enable-minference
 - Prefill: 3383 tok/s
 - Output: "7492<|im_end|>"
 - Result: PASSED
 # Offload + MInference
 test_needle.py --model Qwen3-4B --input-len 32768 --enable-offload --enable-minference
 - Prefill: 5373 tok/s
 - Output: "7492<|im_end|>"
 - Result: PASSED
 ```
 Both configurations produce identical outputs, confirming correctness.
 ## Related Documents
 - [`sparse_attention_guide.md`](sparse_attention_guide.md): Algorithm details for sparse methods
 - [`architecture_guide.md`](architecture_guide.md): Overall system architecture
 - [`gpu_only_performance_issue.md`](gpu_only_performance_issue.md): Why offload is faster than GPU-only
--- a/docs/sparse_prefill_integration_plan.md
+++ b/docs/sparse_prefill_integration_plan.md
@@ -0,0 +1,367 @@
 # Sparse Prefill Attention Integration Plan
 ## Executive Summary
 本文档整合了 int-minference-1/2/3 三个分支的分析，提出统一的三种稀疏注意力策略（MInference、XAttention、FlexPrefill）集成方案。
 ---
 ## Part 1: 现状分析
 ### 1.1 x-attention 仓库策略对比
 | 策略 | Pattern 类型 | 估计方法 | Kernel Backend |
 |------|-------------|---------|----------------|
 | **MInference** | Vertical + Slash | Last-64-Q attention → 列/对角线求和 | `vertical_slash_sparse_attention` (minference lib) |
 | **XAttention** | Block Mask | Stride-based Q/K 下采样 → block 分数 | `block_sparse_attn_func` (MIT-HAN-LAB) |
 | **FlexPrefill** | Adaptive V+S | Last-block attention + JS 散度自适应 | `triton_block_wise_attention` (custom triton) |
 ### 1.2 关键发现：两种 Kernel 接口
 **接口 A: Index-Based (minference)**
 ```python
 # MInference 使用 vertical+slash indices
 vertical_indices = [heads, vertical_size]  # 重要 K 列位置
 slash_indices = [heads, slash_size]        # 对角线偏移
 output = vertical_slash_sparse_attention(q, k, v, vertical_indices, slash_indices)
 ```
 **接口 B: Block Mask-Based (block_sparse_attn)**
 ```python
 # XAttention/FlexPrefill 使用 boolean block mask
 block_mask = torch.bool[batch, heads, q_blocks, k_blocks]  # True = 计算
 output = block_sparse_attn_func(q, k, v, block_mask, ...)
 ```
 ### 1.3 当前 nanovllm MInference 实现
 **文件**: `nanovllm/kvcache/sparse/minference.py`
 **已实现功能**:
 - `estimate_pattern()`: 使用 last-64-Q 估计 vertical+slash pattern
 - `sparse_prefill_attention()`: 调用 minference kernel 执行稀疏注意力
 - 支持 GQA（通过 K/V repeat_interleave）
 - 支持 adaptive_budget 自适应预算
 **问题**:
 1. 与 XAttention/FlexPrefill 使用不同 kernel，无法统一接口
 2. `sparse_prefill_attention()` 将估计和执行耦合在一起
 3. 没有 BlockMask 中间表示，难以复用
 ---
 ## Part 2: 架构设计
 ### 2.1 设计原则
 1. **向后兼容**: 保持现有 `SparsePolicy` 接口不变
 2. **渐进式重构**: 添加新功能而非替换
 3. **统一中间表示**: 新策略使用 `BlockMask` 作为可选中间表示
 4. **可插拔 Kernel**: 支持多种 attention kernel backend
 ### 2.2 架构图
 ```
 ┌──────────────────────────────────────────────────────────────────────────────┐
 │                       Unified Sparse Prefill Framework                        │
 ├──────────────────────────────────────────────────────────────────────────────┤
 │                                                                               │
 │  ┌─────────────────┐  ┌─────────────────┐  ┌─────────────────┐               │
 │  │   MInference    │  │   XAttention    │  │   FlexPrefill   │  Strategies   │
 │  │   Policy        │  │   Policy        │  │   Policy        │               │
 │  └────────┬────────┘  └────────┬────────┘  └────────┬────────┘               │
 │           │                    │                    │                         │
 │           │ (indices)          │ (BlockMask)        │ (BlockMask)             │
 │           │                    │                    │                         │
 │           ▼                    └────────┬───────────┘                         │
 │  ┌─────────────────┐                    ▼                                     │
 │  │   minference    │  ┌─────────────────────────────────────────────────────┐│
 │  │   kernel        │  │              BlockMask Container                    ││
 │  └────────┬────────┘  │  [batch, num_heads, q_blocks, k_blocks] - boolean   ││
 │           │           └─────────────────────────────────────────────────────┘│
 │           │                              │                                    │
 │           │                              ▼                                    │
 │           │           ┌─────────────────────────────────────────────────────┐│
 │           │           │         block_sparse_attn_func                      ││
 │           │           │         (MIT-HAN-LAB kernel)                        ││
 │           │           └─────────────────────────────────────────────────────┘│
 │           │                              │                                    │
 │           └──────────────────────────────┼────────────────────────────────── │
 │                                          ▼                                    │
 │  ┌─────────────────────────────────────────────────────────────────────────┐ │
 │  │                        Attention Output                                  │ │
 │  │                   [seq_len, num_heads, head_dim]                         │ │
 │  └─────────────────────────────────────────────────────────────────────────┘ │
 │                                                                               │
 └──────────────────────────────────────────────────────────────────────────────┘
 ```
 ### 2.3 新增类设计
 ```python
 # nanovllm/kvcache/sparse/block_mask.py
@dataclass
 class BlockMask:
    """Block-level attention mask container."""
    mask: torch.Tensor      # [batch, heads, q_blocks, k_blocks]
    block_size: int
    seq_len: int
    num_q_blocks: int
    num_k_blocks: int
    def sparsity_ratio(self) -> float:
        """Fraction of blocks masked out."""
        return 1.0 - self.mask.float().mean().item()
    def to_flat_indices(self, head_idx: int) -> torch.Tensor:
        """Convert to flattened block indices for a given head."""
        pass
    @classmethod
    def from_vertical_slash(
        cls,
        vertical_idx: torch.Tensor,
        slash_idx: torch.Tensor,
        seq_len: int,
        block_size: int,
    ) -> "BlockMask":
        """Convert MInference-style indices to block mask."""
        pass
    def apply_causal(self) -> "BlockMask":
        """Apply causal constraint (lower triangular)."""
        pass
 ```
 ```python
 # nanovllm/kvcache/sparse/kernels/block_sparse.py
 def block_sparse_attention(
    q: torch.Tensor,           # [seq_len, num_heads, head_dim]
    k: torch.Tensor,           # [seq_len, num_kv_heads, head_dim]
    v: torch.Tensor,           # [seq_len, num_kv_heads, head_dim]
    block_mask: BlockMask,
 ) -> torch.Tensor:
    """
    Execute block sparse attention using MIT-HAN-LAB kernel.
    Handles:
    - GQA expansion (K/V heads < Q heads)
    - Tensor format conversion
    - Causal masking
    """
    from block_sparse_attn import block_sparse_attn_func
    # ... implementation
 ```
 ---
 ## Part 3: 实现计划
 ### Phase 1: 基础设施 (新增文件)
 **目标**: 添加 BlockMask 和 block_sparse_attn 封装
 **文件**:
 - `nanovllm/kvcache/sparse/block_mask.py` (NEW)
 - `nanovllm/kvcache/sparse/kernels/__init__.py` (NEW)
 - `nanovllm/kvcache/sparse/kernels/block_sparse.py` (NEW)
 **任务**:
 1. 实现 `BlockMask` 数据类
 2. 实现 `block_sparse_attention()` 封装函数
 3. 处理 GQA 和 tensor 格式转换
 4. 测试：使用全 True 的 block mask 验证输出正确
 ### Phase 2: XAttention 实现
 **目标**: 移植 x-attention 的 XAttention 策略
 **文件**:
 - `nanovllm/kvcache/sparse/xattention.py` (NEW)
 - `nanovllm/config.py` (添加 XATTENTION 枚举)
 - `nanovllm/kvcache/sparse/__init__.py` (更新工厂函数)
 **关键函数移植**:
 ```python
 # From x-attention/xattn/src/Xattention.py
 def xattn_estimate(q, k, block_size, stride, threshold, ...):
    # 1. Stride-based Q/K downsampling
    reshaped_k = cat([k[:, :, i::stride, :] for i in range(stride)], dim=-1)
    reshaped_q = cat([q[:, :, stride-1-i::stride, :] for i in range(stride)], dim=-1)
    # 2. Block-level attention scores
    attn_weights = matmul(reshaped_q, reshaped_k.T) / sqrt(d) / stride
    # 3. Threshold selection
    block_mask = find_blocks_chunked(attn_sum, threshold)
    return block_mask
 ```
 **配置参数**:
 ```python
 xattention_stride: int = 16           # Q/K 下采样步长
 xattention_threshold: float = 0.9     # 累积分数阈值
 xattention_block_size: int = 128      # Block 大小
 ```
 **测试**: `python tests/test_needle.py --input-len 32768 --enable-xattention`
 ### Phase 3: FlexPrefill 实现
 **目标**: 移植 x-attention 的 FlexPrefill 策略
 **文件**:
 - `nanovllm/kvcache/sparse/flexprefill.py` (NEW)
 - `nanovllm/config.py` (添加 FLEXPREFILL 枚举)
 **关键函数移植**:
 ```python
 # From x-attention/xattn/src/Flexprefill.py
 def get_active_blocks(q, k, gamma, tau, block_size, ...):
    # 1. Last-block attention analysis
    last_q = q[:, -block_size:, :, :]
    qk = einsum('bihd,bjhd->bhij', last_q, k)
    # 2. Vertical + slash pattern detection
    vertical = qk.mean(-2)  # Column importance
    slash = sum_all_diagonal_matrix(qk)  # Diagonal importance
    # 3. JS divergence for adaptive budget
    kl_div = js_divergence(avg_qk, vertical_pooled)
    is_sparse_head = kl_div > tau
    budget = gamma if is_sparse_head else 1.0
    # 4. Select blocks
    block_idx = transform_vertical_slash_idx(...)
    return block_mask
 ```
 **配置参数**:
 ```python
 flexprefill_gamma: float = 0.9        # 基础覆盖率
 flexprefill_tau: float = 0.1          # JS 散度阈值
 flexprefill_min_budget: int = 128     # 最小 token 预算
 flexprefill_block_size: int = 128     # Block 大小
 ```
 **测试**: `python tests/test_needle.py --input-len 32768 --enable-flexprefill`
 ### Phase 4: MInference 可选重构
 **目标**: (可选) 让 MInference 也可以使用 block_sparse_attn
 **修改文件**:
 - `nanovllm/kvcache/sparse/minference.py`
 **新增方法**:
 ```python
 class MInferencePolicy(SparsePolicy):
    def __init__(self, ..., use_block_sparse: bool = False):
        self.use_block_sparse = use_block_sparse
    def estimate_block_mask(self, q, k, layer_id) -> BlockMask:
        """Convert vertical+slash indices to BlockMask."""
        vertical_idx, slash_idx = self.estimate_pattern(q, k, layer_id)
        return BlockMask.from_vertical_slash(vertical_idx, slash_idx, ...)
    def sparse_prefill_attention(self, q, k, v, layer_id):
        if self.use_block_sparse:
            block_mask = self.estimate_block_mask(q, k, layer_id)
            return block_sparse_attention(q, k, v, block_mask)
        else:
            # 使用原有 minference kernel
            return self._minference_kernel_attention(q, k, v, layer_id)
 ```
 ### Phase 5: 集成和测试
 **任务**:
 1. 更新 `__init__.py` 工厂函数支持所有策略
 2. 更新 Config 添加所有配置参数
 3. 添加性能基准测试脚本
 4. 更新文档
 ---
 ## Part 4: 依赖管理
 ### 必需依赖
 ```
 # requirements.txt 新增
 block-sparse-attn    # MIT-HAN-LAB block sparse kernel
 triton>=2.0          # FlexPrefill Triton kernels
 ```
 ### 安装说明
 ```bash
 # block_sparse_attn from MIT-HAN-LAB
 pip install git+https://github.com/mit-han-lab/Block-Sparse-Attention.git
 # 或从本地安装（如果有）
 cd /home/zijie/Code/x-attention/Block-Sparse-Attention
 pip install -e .
 ```
 ---
 ## Part 5: 配置参数汇总
 ### SparsePolicyType 枚举
 ```python
 class SparsePolicyType(str, Enum):
    FULL = "full"              # 全注意力（无稀疏）
    QUEST = "quest"            # Decode-only Top-K
    MINFERENCE = "minference"  # Prefill vertical+slash
    XATTENTION = "xattention"  # Prefill stride-based block
    FLEXPREFILL = "flexprefill"  # Prefill adaptive JS-divergence
 ```
 ### 策略参数对照表
 | 策略 | 参数 | 默认值 | 说明 |
 |------|-----|--------|------|
 | MInference | `adaptive_budget` | 0.3 | 预算占 seq_len 比例 |
 | MInference | `vertical_size` | 1000 | 固定 vertical 大小 |
 | MInference | `slash_size` | 6096 | 固定 slash 大小 |
 | XAttention | `stride` | 16 | Q/K 下采样步长 |
 | XAttention | `threshold` | 0.9 | 累积分数阈值 |
 | XAttention | `block_size` | 128 | Block 大小 |
 | FlexPrefill | `gamma` | 0.9 | 基础覆盖率 |
 | FlexPrefill | `tau` | 0.1 | JS 散度阈值 |
 | FlexPrefill | `min_budget` | 128 | 最小 token 预算 |
 | FlexPrefill | `block_size` | 128 | Block 大小 |
 ---
 ## Part 6: 成功标准
 1. **正确性**: 所有三种策略通过 32K+ needle-in-haystack 测试
 2. **性能**: 稀疏 prefill 比全注意力快 (>1.5x speedup at 64K)
 3. **统一接口**: XAttention/FlexPrefill 使用 BlockMask + block_sparse_attn
 4. **向后兼容**: 现有 MInference 配置继续工作
 5. **可配置**: 所有策略参数可通过 LLM 配置设置
 ---
 ## Part 7: 风险评估
 | 风险 | 影响 | 可能性 | 缓解措施 |
 |------|-----|--------|---------|
 | block_sparse_attn 硬件兼容性 | 高 | 中 | 测试目标硬件，fallback 到 flash_attn |
 | MInference → block mask 精度损失 | 中 | 低 | 对比测试输出差异 |
 | Triton kernel 移植问题 | 中 | 中 | 使用非 Triton fallback |
 | 内存开销增加 | 低 | 低 | block_size=128 → 1KB/head for 128K |
 ---
 ## References
 - x-attention repo: `/home/zijie/Code/x-attention`
 - MIT-HAN-LAB Block-Sparse-Attention: `https://github.com/mit-han-lab/Block-Sparse-Attention`
 - MInference paper: https://arxiv.org/abs/2407.02490
 - Current nanovllm sparse implementation: `nanovllm/kvcache/sparse/`
--- a/docs/transformers_compatibility.md
+++ b/docs/transformers_compatibility.md
@@ -0,0 +1,279 @@
 # Transformers 低版本兼容性问题
 ## 概述
 本文档详细记录了 nano-vllm 在低版本 transformers（< 4.51.0）环境下的兼容性问题。这些问题源于 nano-vllm 使用了 transformers 4.51.0 才引入的 `Qwen3Config` 类。
 ## 问题背景
 ### 测试环境
 | 环境 | 版本 | 说明 |
 |------|------|------|
 | Docker 镜像 | `tzj/ruler:v0.3` | NVIDIA PyTorch 24.08 容器 |
 | transformers | 4.45.2 | 系统预装版本 |
 | Python | 3.10.12 | 系统版本 |
 | PyTorch | 2.5.0a0+872d972 | CUDA 12.6 |
 ### 冲突场景
 在 RULER benchmark 测试环境中，NeMo 框架依赖 transformers 4.45.2 和特定版本的 `huggingface_hub`。升级 transformers 到 4.51.0+ 会导致：
 ```
 ImportError: cannot import name 'ModelFilter' from 'huggingface_hub'
 ```
 因此需要 nano-vllm 适配低版本 transformers，以便在同一环境中运行。
 ## 详细问题分析
 ### 1. 核心问题：Qwen3Config 不存在
 **错误信息**：
 ```python
 ImportError: cannot import name 'Qwen3Config' from 'transformers'
 (/usr/local/lib/python3.10/dist-packages/transformers/__init__.py)
 ```
 **问题根源**：
 - `Qwen3Config` 是在 transformers **4.51.0** 版本中首次引入
 - transformers 4.45.2 只包含 `Qwen2` 系列模型
 **受影响版本**：
 | transformers 版本 | Qwen3 支持 | 可用 Qwen 模型 |
 |------------------|-----------|---------------|
 | < 4.51.0 | 不支持 | qwen2, qwen2_audio, qwen2_moe, qwen2_vl |
 | >= 4.51.0 | 支持 | qwen2 系列 + qwen3, qwen3_moe |
 ### 2. 影响范围
 #### 2.1 直接影响的文件
 | 文件路径 | 问题代码 | 影响 |
 |---------|---------|------|
 | `nanovllm/models/qwen3.py:4` | `from transformers import Qwen3Config` | 直接导入失败 |
 | `nanovllm/models/__init__.py:6` | `from nanovllm.models import qwen3` | 触发 qwen3 导入 |
 #### 2.2 级联影响
 由于 `nanovllm/models/__init__.py` 无条件导入了 `qwen3` 模块，会导致以下级联失败：
 ```python
 # 这些导入都会失败
 from nanovllm.models import llama          # FAILED
 from nanovllm.models import get_model_class # FAILED
 import nanovllm                            # FAILED
 ```
 **测试验证**：
 ```python
 # transformers 4.45.2 环境
 >>> from nanovllm.models.registry import register_model
 SUCCESS  # registry 本身可以导入
 >>> from nanovllm.config import Config
 SUCCESS  # config 不依赖 Qwen3Config
 >>> from nanovllm.models import llama
 FAILED: cannot import name 'Qwen3Config' from 'transformers'
 # 因为 models/__init__.py 先导入了 qwen3
 ```
 ### 3. Qwen3Config 使用位置
 在 `nanovllm/models/qwen3.py` 中的使用：
 ```python
 # Line 4
 from transformers import Qwen3Config
 # Line 128-129: 类型注解
 class Qwen3DecoderLayer(nn.Module):
    def __init__(self, config: Qwen3Config) -> None:
        ...
 # Line 170-171: 类型注解
 class Qwen3Model(nn.Module):
    def __init__(self, config: Qwen3Config) -> None:
        ...
 # Line 200-203: 类型注解
 class Qwen3ForCausalLM(nn.Module):
    def __init__(self, config: Qwen3Config) -> None:
        ...
 ```
 ### 4. Qwen3Config 属性使用
 代码中使用了以下 `Qwen3Config` 属性：
 | 属性 | 位置 | 用途 |
 |------|------|------|
 | `hidden_size` | Line 131, 147, 173 | 隐藏层维度 |
 | `num_attention_heads` | Line 132 | 注意力头数 |
 | `num_key_value_heads` | Line 133 | KV 头数 |
 | `max_position_embeddings` | Line 134 | 最大位置编码 |
 | `rms_norm_eps` | Line 135, 147, 148, 175 | RMSNorm epsilon |
 | `attention_bias` | Line 136 (getattr) | 是否使用注意力偏置 |
 | `head_dim` | Line 137 (getattr) | 注意力头维度 |
 | `rope_theta` | Line 138 (getattr) | RoPE base |
 | `rope_scaling` | Line 139 (getattr) | RoPE scaling 配置 |
 | `intermediate_size` | Line 144 | FFN 中间层维度 |
 | `hidden_act` | Line 145 | 激活函数类型 |
 | `vocab_size` | Line 173, 206 | 词表大小 |
 | `num_hidden_layers` | Line 174 | Transformer 层数 |
 | `tie_word_embeddings` | Line 207 | 是否共享词嵌入 |
 ## 解决方案建议
 ### 方案 1: 条件导入（推荐）
 修改 `nanovllm/models/__init__.py`：
 ```python
 """Model registry and model implementations."""
 from nanovllm.models.registry import register_model, get_model_class, MODEL_REGISTRY
 # Import models to trigger registration
 # Llama is always available
 from nanovllm.models import llama
 # Qwen3 requires transformers >= 4.51.0
 try:
    from nanovllm.models import qwen3
 except ImportError:
    import warnings
    warnings.warn(
        "Qwen3 models require transformers >= 4.51.0. "
        "Install with: pip install 'transformers>=4.51.0'"
    )
 __all__ = ["register_model", "get_model_class", "MODEL_REGISTRY"]
 ```
 修改 `nanovllm/models/qwen3.py`：
 ```python
 import torch
 from torch import nn
 import torch.distributed as dist
 # Conditional import for Qwen3Config
 try:
    from transformers import Qwen3Config
 except ImportError:
    # Create a placeholder for type hints when Qwen3Config is not available
    Qwen3Config = None
    raise ImportError(
        "Qwen3Config requires transformers >= 4.51.0. "
        "Current version does not support Qwen3 models."
    )
 # ... rest of the code
 ```
 ### 方案 2: 使用 AutoConfig（兼容性更好）
 修改 `nanovllm/models/qwen3.py` 以使用 `AutoConfig` 而非具体的 `Qwen3Config`：
 ```python
 from typing import TYPE_CHECKING, Any
 # Only import Qwen3Config for type checking
 if TYPE_CHECKING:
    from transformers import Qwen3Config
 # Runtime: use duck typing
 class Qwen3DecoderLayer(nn.Module):
    def __init__(self, config: Any) -> None:  # Accept any config-like object
        super().__init__()
        # Access attributes via getattr for safety
        self.self_attn = Qwen3Attention(
            hidden_size=config.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=config.num_key_value_heads,
            max_position=config.max_position_embeddings,
            rms_norm_eps=config.rms_norm_eps,
            qkv_bias=getattr(config, 'attention_bias', True),
            head_dim=getattr(config, 'head_dim', None),
            rope_theta=getattr(config, "rope_theta", 1000000),
            rope_scaling=getattr(config, "rope_scaling", None),
        )
        # ...
 ```
 ### 方案 3: 版本检查与优雅降级
 在 `nanovllm/__init__.py` 或启动时添加版本检查：
 ```python
 import transformers
 from packaging import version
 TRANSFORMERS_VERSION = version.parse(transformers.__version__)
 QWEN3_MIN_VERSION = version.parse("4.51.0")
 QWEN3_AVAILABLE = TRANSFORMERS_VERSION >= QWEN3_MIN_VERSION
 if not QWEN3_AVAILABLE:
    import warnings
    warnings.warn(
        f"transformers {transformers.__version__} does not support Qwen3 models. "
        f"Upgrade to >= 4.51.0 for Qwen3 support."
    )
 ```
 ## 适配优先级
 建议按以下优先级进行适配：
 1. **P0 - models/__init__.py**: 添加 try-except 使 Llama 模型可独立使用
 2. **P1 - qwen3.py**: 添加清晰的错误信息，说明版本要求
 3. **P2 - 类型注解**: 可选地改为 `Any` 或使用 `TYPE_CHECKING`
 4. **P3 - 文档**: 在 README 和 pyproject.toml 中说明版本依赖
 ## 测试验证
 适配后应验证以下场景：
 ### 测试 1: 低版本环境（transformers 4.45.2）
 ```bash
 # 预期结果：Llama 模型可用，Qwen3 提示版本不足
 docker run --rm \
    -v /path/to/nano-vllm:/workspace/nano-vllm \
    -e PYTHONPATH=/workspace/nano-vllm \
    tzj/ruler:v0.3 \
    python -c "
 from nanovllm.models import get_model_class, MODEL_REGISTRY
 print('Available models:', list(MODEL_REGISTRY.keys()))
 # Expected: ['LlamaForCausalLM']
 # Warning: Qwen3 models require transformers >= 4.51.0
 "
 ```
 ### 测试 2: 高版本环境（transformers >= 4.51.0）
 ```bash
 # 预期结果：Llama 和 Qwen3 模型均可用
 pip install 'transformers>=4.51.0'
 python -c "
 from nanovllm.models import get_model_class, MODEL_REGISTRY
 print('Available models:', list(MODEL_REGISTRY.keys()))
 # Expected: ['LlamaForCausalLM', 'Qwen3ForCausalLM', 'Qwen2ForCausalLM']
 "
 ```
 ## 相关参考
 - [Transformers Qwen3 文档](https://huggingface.co/docs/transformers/en/model_doc/qwen3)
 - [Qwen3 GitHub](https://github.com/QwenLM/Qwen3)
 - [Transformers 版本历史](https://github.com/huggingface/transformers/releases)
 ## 版本信息
 | 日期 | 版本 | 变更 |
 |------|------|------|
 | 2025-01-11 | 1.0 | 初始文档，记录 transformers 4.45.2 兼容性问题 |
--- a/findings.md
+++ b/findings.md
@@ -1,160 +1,288 @@
-# Findings: Multi-Model Support Analysis
+# Findings: nanovllm 多请求状态污染分析
-## Current Architecture Analysis
+## 重要说明
-### Model Loading Flow
+**nanovllm offload 模式不支持 batch**，只能单个 request 顺序执行。问题出在**请求切换**（前一个 request 完成后，开始下一个 request）时状态清理不完整。
 ```
 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
+## 1. 代码架构发现
-```json
+### 1.1 请求生命周期 (顺序执行)
-{
+
-  "architectures": ["LlamaForCausalLM"],
+**关键**: offload 模式下，每次只处理**一个 request**，不是 batch。
-  "model_type": "llama",
+
-  "attention_bias": false,
+```
-  "mlp_bias": false,
+LLMEngine.generate() [llm_engine.py:114-151]
-  "head_dim": 128,
+├── Observer.complete_reset()  # 重置性能统计
-  "hidden_size": 4096,
+├── for prompt in prompts:
-  "intermediate_size": 14336,
+│   └── add_request(prompt, sp)  # 添加到 scheduler 队列
-  "num_attention_heads": 32,
+├── while not is_finished():
-  "num_hidden_layers": 32,
+│   ├── scheduler.schedule()  # 获取下一个序列 (offload 模式: 1个)
-  "num_key_value_heads": 8,
+│   ├── model_runner.call("run", seqs, is_prefill)  # 执行单个请求
-  "hidden_act": "silu",
+│   └── scheduler.postprocess(seqs, token_ids)
-  "rms_norm_eps": 1e-05,
+│       └── if seq.is_finished:
-  "rope_theta": 500000.0,
+│           └── kvcache_manager.deallocate(seq)  # 释放资源 ← 问题点
-  "rope_scaling": {
+│           └── [开始处理下一个请求]  # ← 状态切换
-    "factor": 8.0,
+└── return outputs
    "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"`)：
+```
- 低频分量保持不变（对应短距离依赖）
+Request A (prefill) → Request A (decode × N) → Request A 完成
- 高频分量线性插值（对应长距离依赖）
+    ↓
- 参数: `factor`, `low_freq_factor`, `high_freq_factor`, `original_max_position_embeddings`
+deallocate(A)  ← 状态清理不完整!
    ↓
 Request B (prefill) → Request B 读取到 A 的残留状态 → 错误输出
 ```
-参考实现 (transformers):
+### 1.2 OffloadEngine 状态清单
 **位置**: `nanovllm/kvcache/offload_engine.py:40-145`
 | 成员变量 | 类型 | Shape | 生命周期 |
 |----------|------|-------|----------|
 | `layer_k_cache` | GPU Tensor | [num_buffers, max_seq_len, kv_heads, head_dim] | 整个引擎 |
 | `layer_v_cache` | GPU Tensor | [num_buffers, max_seq_len, kv_heads, head_dim] | 整个引擎 |
 | `decode_k_buffer` | GPU Tensor | [num_layers, block_size, kv_heads, head_dim] | 整个引擎 |
 | `decode_v_buffer` | GPU Tensor | [num_layers, block_size, kv_heads, head_dim] | 整个引擎 |
 | `k_cache_cpu` | CPU Tensor (pinned) | [num_layers, num_cpu_blocks, block_size, kv_heads, head_dim] | 整个引擎 |
 | `v_cache_cpu` | CPU Tensor (pinned) | [num_layers, num_cpu_blocks, block_size, kv_heads, head_dim] | 整个引擎 |
 | `compute_stream` | CUDA Stream | - | 整个引擎 |
 | `prefill_offload_streams` | List[CUDA Stream] | num_layers | 整个引擎 |
 | `prefill_offload_events` | List[CUDA Event] | num_layers | 整个引擎 |
 | `layer_load_streams` | List[CUDA Stream] | num_buffers | 整个引擎 |
 | `buffer_load_events` | List[CUDA Event] | num_buffers | 整个引擎 |
 | `buffer_compute_done_events` | List[CUDA Event] | num_buffers | 整个引擎 |
 **关键发现**:
 - **没有 reset() 方法**
 - **没有任何清理逻辑**
 - 所有 tensor 在初始化时 `torch.zeros()` 后永不清零
 ### 1.3 HybridKVCacheManager 状态清单
 **位置**: `nanovllm/kvcache/hybrid_manager.py`
 | 成员变量 | 作用 | 清理方式 |
 |----------|------|----------|
 | `logical_blocks` | 逻辑块列表 | `block.reset()` in deallocate |
 | `free_logical_ids` | 空闲逻辑块队列 | deallocate 归还 |
 | `free_cpu_blocks` | 空闲 CPU 块队列 | deallocate 归还 |
 | `cpu_block_to_logical` | CPU 块→逻辑块映射 | deallocate 删除 |
 | `prefilled_blocks` | 已 prefill 的块集合 | deallocate 中 discard |
 | `_decode_start_pos` | 序列→decode起始位置 | `clear_decode_tracking()` |
 | `_prefill_len` | 序列→prefill长度 | `clear_decode_tracking()` |
 **关键发现**:
 - `deallocate()` 没有调用 `clear_decode_tracking()`！
 - `_decode_start_pos` 和 `_prefill_len` 使用 `id(seq)` 作为 key
 - Python 对象 ID 可能在不同请求间重用
 ---
 ## 2. 请求切换机制分析
 ### 2.1 offload 模式的单 request 限制
 代码中明确限制：
 ```python
-def _compute_llama3_parameters(config, device, inv_freq):
+# model_runner.py:757, 880
-    factor = config.factor
+assert len(seqs) == 1, "Layer-wise offload only supports single sequence"
-    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
+### 2.2 请求切换时序
    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,
+│ Request A: [prefill] → [decode] → [decode] → ... → [完成]       │
-        inv_freq
+└─────────────────────────────────────────────────────────────────┘
                                                      ↓
                                          deallocate(seq_A)
                                          - blocks 释放 ✓
                                          - tracking 字典未清理 ✗
                                                      ↓
 ┌─────────────────────────────────────────────────────────────────┐
 │ Request B: [prefill] → [decode] → ...                           │
 │            ↑                                                     │
 │            如果 id(seq_B) == id(seq_A)，读到 A 的残留状态！      │
 └─────────────────────────────────────────────────────────────────┘
 ```
 ### 2.3 Python 对象 ID 重用
 Python 的内存管理会重用已释放对象的内存地址，导致：
 ```python
 seq_A = Sequence(...)  # id(seq_A) = 0x7f1234567890
 del seq_A              # 对象被释放，但字典中 key 保留
 seq_B = Sequence(...)  # id(seq_B) 可能 = 0x7f1234567890（相同地址）
 # _decode_start_pos[id(seq_B)] 返回 seq_A 的旧值！
 ```
 ---
 ## 3. 状态污染机制分析
 ### 3.1 decode buffer 污染路径
 **污染写入** (`run_layerwise_offload_decode:1010-1013`):
 ```python
 # 每次 decode step，将当前 token 的 KV 存入 decode buffer
 offload_engine.decode_k_buffer[layer_id, pos_in_block].copy_(ring_k[context_len])
 offload_engine.decode_v_buffer[layer_id, pos_in_block].copy_(ring_v[context_len])
 ```
 **污染读取** (`run_layerwise_offload_decode:969-976`):
 ```python
 # 如果有之前的 decode tokens，从 decode buffer 读取
 if num_prev_decode_tokens > 0:
    k_decode_prev, v_decode_prev = offload_engine.get_decode_kv(
        layer_id, decode_start_pos, pos_in_block
    )
-    smooth_factor = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
+    ring_k[total_prefill_tokens:total_prefill_tokens + num_prev_decode_tokens].copy_(k_decode_prev)
    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
 ```
---
+**问题场景**:
 1. 请求 A 的 decode 阶段在 `decode_k_buffer[layer, 0:N]` 写入 KV
 2. 请求 A 完成，buffer 数据保留
 3. 请求 B 开始，如果其 `decode_start_pos` 被错误计算为非零
 4. 请求 B 会读取请求 A 的旧数据
-## Weight Mapping Analysis
+### 3.2 decode_start_pos 计算逻辑
 **位置**: `hybrid_manager.py:485-505`
 ### Qwen3 packed_modules_mapping
 ```python
-packed_modules_mapping = {
+def get_decode_start_pos(self, seq: Sequence) -> int:
-    "q_proj": ("qkv_proj", "q"),
+    seq_id = id(seq)  # Python 对象 ID
-    "k_proj": ("qkv_proj", "k"),
+    if seq_id not in self._decode_start_pos:
-    "v_proj": ("qkv_proj", "v"),
+        # 第一次调用 - 计算起始位置
-    "gate_proj": ("gate_up_proj", 0),
+        prefill_len = len(seq) - 1  # 当前长度减去新 token
-    "up_proj": ("gate_up_proj", 1),
+        self._decode_start_pos[seq_id] = prefill_len % self._block_size
-}
+    return self._decode_start_pos[seq_id]
 ```
-### Llama Weight Names (from safetensors)
+**问题**:
-预期 Llama 权重命名与 Qwen3 类似：
+- 如果新请求的 `id(seq)` 恰好等于旧请求的 `id(seq)`（Python 内存重用）
- `model.layers.{i}.self_attn.q_proj.weight`
+- `_decode_start_pos` 中可能存在旧的值
- `model.layers.{i}.self_attn.k_proj.weight`
+- 会返回错误的 decode 起始位置
 - `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 相同，可以复用。
+### 3.3 clear_decode_tracking 未被调用
 **位置**: `hybrid_manager.py:538-549`
 ```python
 def clear_decode_tracking(self, seq: Sequence) -> None:
    seq_id = id(seq)
    self._decode_start_pos.pop(seq_id, None)
    self._prefill_len.pop(seq_id, None)
 ```
 **问题**:
 - 这个方法在 `deallocate()` 中**没有被调用**！
 - 查看 `deallocate()` (218-244 行)，没有 `clear_decode_tracking()` 调用
 - 这导致旧请求的 tracking 数据残留
 ---
-## Shared Components (Can Reuse)
+## 3. 失败模式分析
-| Component | File | Notes |
+### 3.1 观察到的失败模式
-|-----------|------|-------|
+
-| `RMSNorm` | `layers/layernorm.py` | 通用 |
+从测试结果:
-| `SiluAndMul` | `layers/activation.py` | 通用 |
+| Sample | Expected | Output | Status |
-| `Attention` | `layers/attention.py` | FlashAttention wrapper |
+|--------|----------|--------|--------|
-| `QKVParallelLinear` | `layers/linear.py` | 支持 bias=False |
+| 0 | 8930103 | `: 8930103.` | PASS (第一个请求) |
-| `RowParallelLinear` | `layers/linear.py` | 通用 |
+| 1 | 4194548 | `: 419 multiplication of 4548.` | **FAIL** |
-| `MergedColumnParallelLinear` | `layers/linear.py` | 通用 |
+| 2 | 8231838 | `:ное 8231838.` | PASS |
-| `VocabParallelEmbedding` | `layers/embed_head.py` | 通用 |
+
-| `ParallelLMHead` | `layers/embed_head.py` | 通用 |
+Sample 1 的输出 "419 multiplication of 4548" 显示数字被"拆分"了。
-| `load_model` | `utils/loader.py` | 通用 |
+
 **可能原因**:
 1. 在某个 decode step，attention 计算使用了错误的 KV
 2. 模型"看到"了旧请求的部分 context
 3. 导致生成逻辑出错
 ### 3.2 为什么第一个请求总是成功？
 1. 第一个请求时，所有 buffer 都是零初始化
 2. `decode_start_pos` 字典为空，正确计算
 3. 没有残留数据干扰
 ### 3.3 为什么后续请求可能成功？
 某些请求可能成功因为：
 1. `id(seq)` 没有与之前的请求冲突
 2. `pos_in_block` 不重叠，没读到旧数据
 3. 或者旧数据恰好对结果影响不大
 ---
-## Llama vs Qwen3 Implementation Diff
+## 4. 修复方向
-### Attention
+### 4.1 必须修复: deallocate 时清理状态
 | Feature | Qwen3Attention | LlamaAttention |
 |---------|----------------|----------------|
 | QKV bias | 可配置 (attention_bias) | 始终 False |
 | q_norm | 有 (when bias=False) | 无 |
 | k_norm | 有 (when bias=False) | 无 |
 | RoPE | Standard | Llama3 scaled |
-### MLP
+```python
-| Feature | Qwen3MLP | LlamaMLP |
+# hybrid_manager.py: deallocate()
-|---------|----------|----------|
+def deallocate(self, seq: Sequence) -> None:
-| gate/up bias | False | False |
+    # ... 现有逻辑 ...
 | down bias | False | False |
 | hidden_act | silu | silu |
-**结论**: Llama MLP 与 Qwen3 MLP 几乎相同，可以直接复用或简化。
+    # 添加: 清理 decode tracking
    self.clear_decode_tracking(seq)
    # 添加: 通知 offload engine 清理
    if self.offload_engine is not None:
        self.offload_engine.on_sequence_finished()
 ```
 ### 4.2 必须修复: OffloadEngine 添加清理方法
 ```python
 # offload_engine.py
 def on_sequence_finished(self):
    """请求完成时的清理"""
    # 清零 decode buffer
    self.decode_k_buffer.zero_()
    self.decode_v_buffer.zero_()
 ```
 ### 4.3 可选: 更激进的清理
 ```python
 def reset_all(self):
    """完全重置状态"""
    self.decode_k_buffer.zero_()
    self.decode_v_buffer.zero_()
    self.layer_k_cache.zero_()
    self.layer_v_cache.zero_()
    # 重置 CUDA events
    for event in self.buffer_compute_done_events:
        event.record()
 ```
 ---
-## Risk Assessment
+## 5. 待验证假设
-| Risk | Impact | Mitigation |
+| 假设 | 验证方法 | 优先级 |
-|------|--------|------------|
+|------|----------|--------|
-| RoPE 实现错误 | 高 - 导致错误输出 | 参考 transformers 实现，单元测试 |
+| decode_buffer 残留导致污染 | 在第二个请求开始时检查 buffer 是否为零 | 高 |
-| 权重映射错误 | 高 - 模型无法加载 | 检查 safetensors 键名 |
+| _decode_start_pos 字典残留 | 打印 deallocate 前后的字典内容 | 高 |
-| 注册表循环导入 | 中 - 启动失败 | 延迟导入 |
+| id(seq) 重用导致错误 | 打印每个请求的 seq id | 中 |
 | ring buffer 残留 | 检查每次 decode 前 ring buffer 内容 | 低 |
 ---
 ## 6. 参考代码位置
 | 功能 | 文件 | 行号 |
 |------|------|------|
 | OffloadEngine 初始化 | offload_engine.py | 40-145 |
 | deallocate | hybrid_manager.py | 218-244 |
 | clear_decode_tracking | hybrid_manager.py | 538-549 |
 | get_decode_start_pos | hybrid_manager.py | 485-505 |
 | run_layerwise_offload_decode | model_runner.py | 867-1057 |
 | decode buffer 写入 | model_runner.py | 1010-1013 |
 | decode buffer 读取 | model_runner.py | 969-976 |
--- a/nanovllm/config.py
+++ b/nanovllm/config.py
@@ -9,6 +9,7 @@ class SparsePolicyType(Enum):
    """Sparse attention policy types."""
    FULL = auto()   # No sparse attention (load all blocks)
    QUEST = auto()  # Query-aware Top-K block selection (decode only)
    MINFERENCE = auto()  # MInference vertical + slash sparse prefill (GPU-only)
@dataclass
@@ -31,6 +32,7 @@ class Config:
    offload_policy: str = "lru"  # "lru", "fifo", or full class path
    num_transfer_streams: int = 4  # Number of CUDA streams for async transfers
    num_gpu_blocks: int = -1  # User-specified GPU blocks count, -1 = auto (use max available)
    num_kv_buffers: int = 4  # Ring buffer size for layer-wise offload (decode H2D pipeline)
    # Computed fields for offload (set in __post_init__ or by ModelRunner)
    num_gpu_kvcache_blocks: int = -1
@@ -39,10 +41,18 @@ class Config:
    # Sparse attention configuration
    # Quest: decode-only sparse attention with Top-K block selection
    # FULL: no sparse attention (load all blocks)
    # MINFERENCE: MInference vertical + slash sparse prefill (GPU-only)
    sparse_policy: SparsePolicyType = SparsePolicyType.FULL
    sparse_topk_blocks: int = 8  # Top-K blocks for Quest
    sparse_threshold_blocks: int = 4  # Apply sparse only when blocks > threshold
    # MInference configuration (used when sparse_policy == MINFERENCE)
    minference_adaptive_budget: float = 0.3  # Budget as fraction of seq_len (None to use fixed sizes)
    minference_vertical_size: int = 1000  # Fixed vertical size (if adaptive_budget is None)
    minference_slash_size: int = 6096  # Fixed slash size (if adaptive_budget is None)
    minference_num_sink_tokens: int = 30  # Sink tokens to always keep
    minference_num_recent_diags: int = 100  # Recent diagonals to always keep
    def __post_init__(self):
        assert os.path.isdir(self.model)
        assert self.kvcache_block_size % 256 == 0
@@ -51,6 +61,15 @@ class Config:
        self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)
        assert self.max_num_batched_tokens >= self.max_model_len
        # CPU offload mode only supports single sequence (layer-wise processing)
        if self.enable_cpu_offload and self.max_num_seqs != 1:
            import logging
            logging.warning(
                f"CPU offload mode only supports single sequence. "
                f"Overriding max_num_seqs from {self.max_num_seqs} to 1."
            )
            self.max_num_seqs = 1
        # Override torch_dtype if user specified
        if self.dtype is not None:
            dtype_map = {
--- a/nanovllm/engine/llm_engine.py
+++ b/nanovllm/engine/llm_engine.py
@@ -34,14 +34,56 @@ class LLMEngine:
        # Set Sequence.block_size to match the KV cache block size
        Sequence.block_size = config.kvcache_block_size
        self.scheduler = Scheduler(config, self.model_runner.kvcache_manager)
-        atexit.register(self.exit)
+        self._closed = False
        atexit.register(self._atexit_handler)
-    def exit(self):
+    def _atexit_handler(self):
        """Handler for atexit - only runs if close() wasn't called."""
        if not self._closed:
            self.close()
    def close(self):
        """Explicitly close the engine and release all resources.
        This method is idempotent - calling it multiple times is safe.
        Supports: explicit close(), context manager, and __del__ fallback.
        """
        if self._closed:
            return
        self._closed = True
        # Unregister atexit to prevent double cleanup
        try:
            atexit.unregister(self._atexit_handler)
        except Exception:
            pass
        # Cleanup resources
        self.model_runner.call("exit")
        del self.model_runner
        for p in self.ps:
            p.join()
    def exit(self):
        """Alias for close() - kept for backward compatibility."""
        self.close()
    def __del__(self):
        """Destructor - attempt cleanup if not already done."""
        try:
            self.close()
        except Exception:
            pass
    def __enter__(self):
        """Context manager entry."""
        return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - ensures cleanup."""
        self.close()
        return False
    def add_request(self, prompt: str | list[int], sampling_params: SamplingParams):
        if isinstance(prompt, str):
            prompt = self.tokenizer.encode(prompt)
--- a/nanovllm/engine/model_runner.py
+++ b/nanovllm/engine/model_runner.py
--- 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
@@ -70,12 +71,20 @@ def create_kvcache_manager(config: "Config") -> KVCacheManager:
        threshold_blocks=getattr(config, 'sparse_threshold_blocks', 4),
    )
    # max_seq_len needs to be larger than max_model_len to accommodate decode tokens
    # When prefill uses ~max_model_len tokens, decode needs additional slots
    # Add max_new_tokens (default 512) buffer for decode phase
    max_new_tokens = getattr(config, 'max_new_tokens', 512)
    max_seq_len = config.max_model_len + max_new_tokens
    return HybridKVCacheManager(
        num_gpu_slots=num_gpu_blocks,
        num_cpu_blocks=num_cpu_blocks,
        block_size=config.kvcache_block_size,
        policy=eviction_policy,
        sparse_policy=sparse_policy,
        num_kv_buffers=getattr(config, 'num_kv_buffers', 4),
        max_seq_len=max_seq_len,
    )
--- 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"
--- a/nanovllm/kvcache/hybrid_manager.py
+++ b/nanovllm/kvcache/hybrid_manager.py
@@ -65,23 +65,22 @@ class LogicalBlock:
 class HybridKVCacheManager(KVCacheManager):
    """
-    Hybrid CPU-GPU KV cache manager with ring buffer design.
+    Hybrid CPU-GPU KV cache manager with layer-wise offload design.
    Architecture (CPU-primary mode):
    - CPU pool: Primary storage for all KV cache (num_cpu_blocks)
-    - GPU buffer: Ring buffer for computation only (num_gpu_slots)
+    - GPU ring buffer: For decode H2D pipeline (num_kv_buffers)
-    - Logical blocks: What sequences reference (num_cpu_blocks)
+    - Decode buffer: Per-layer accumulation of decode tokens (block_size)
    Design:
    - All KV cache is stored on CPU as primary storage
-    - GPU is used as a ring buffer for computation only (no persistent data)
+    - GPU ring buffer enables pipelined H2D transfers during decode
-    - During prefill: KV is written to GPU ring slot, then offloaded to CPU
+    - During prefill: KV is computed and offloaded layer-by-layer to CPU
-    - During decode: Previous KV is loaded from CPU to GPU for attention
+    - During decode: Previous KV is loaded from CPU via ring buffer pipeline
    - Ring buffer enables pipelined H2D transfers overlapped with computation
    Note:
    - Logical blocks map 1:1 with CPU blocks (total_blocks = num_cpu_blocks)
-    - GPU slots are transient compute buffers, not tracked in logical blocks
+    - GPU ring buffer is for decode pipeline, not persistent storage
    """
    def __init__(
@@ -91,25 +90,31 @@ class HybridKVCacheManager(KVCacheManager):
        block_size: int,
        policy: Optional[EvictionPolicy] = None,
        sparse_policy: "SparsePolicy" = None,
        num_kv_buffers: int = 4,
        max_seq_len: int = 131072,
    ):
        """
-        Initialize hybrid manager with CPU-primary ring buffer design.
+        Initialize hybrid manager with layer-wise offload design.
-        All KV cache is stored on CPU as primary storage. GPU slots are used
+        All KV cache is stored on CPU as primary storage. GPU ring buffer is used
-        as a ring buffer for computation only.
+        for decode H2D pipeline.
        Args:
-            num_gpu_slots: Number of GPU buffer slots (ring buffer for computation)
+            num_gpu_slots: Number of GPU buffer slots (kept for backward compat, not used)
            num_cpu_blocks: Number of CPU pool blocks (primary storage)
            block_size: Tokens per block
            policy: Eviction policy (default: LRU, used for prefix cache management)
            sparse_policy: Sparse attention policy (Quest for decode-only sparse)
            num_kv_buffers: Ring buffer size for decode H2D pipeline
            max_seq_len: Maximum sequence length for GPU buffer allocation
        """
        self._block_size = block_size
        self.num_gpu_slots = num_gpu_slots
        self.num_cpu_blocks = num_cpu_blocks
        self.num_kv_buffers = num_kv_buffers
        self.max_seq_len = max_seq_len
        # In CPU-primary mode, logical blocks map 1:1 with CPU blocks
-        # GPU slots are transient compute buffers, not tracked as logical blocks
+        # GPU ring buffer is for decode pipeline, not persistent storage
        self.total_blocks = num_cpu_blocks
        # Eviction policy
@@ -147,7 +152,7 @@ class HybridKVCacheManager(KVCacheManager):
        # Track blocks pending GPU load (for decode graph)
        self.pending_gpu_loads: Set[int] = set()  # logical_ids
-        # Track blocks that have been prefilled (KV written) for chunked prefill
+        # Track blocks that have been prefilled (KV offloaded to CPU)
        self.prefilled_blocks: Set[int] = set()  # logical_ids
        # Track decode starting position within block (for batched offload optimization)
@@ -182,13 +187,21 @@ class HybridKVCacheManager(KVCacheManager):
            num_kv_heads=num_kv_heads,
            head_dim=head_dim,
            dtype=dtype,
            num_kv_buffers=self.num_kv_buffers,
            max_seq_len=self.max_seq_len,
            sparse_policy=self.sparse_policy,
        )
    def get_layer_cache(self, layer_id: int) -> Tuple[Tensor, Tensor]:
-        """Get GPU K/V cache tensors for a layer."""
+        """
        Get GPU K/V cache tensors for a layer.
        Note: In layer-wise offload mode, this returns empty tensors as KV
        is managed directly by the offload engine's ring buffer.
        """
        assert self.offload_engine is not None
-        return self.offload_engine.get_layer_cache(layer_id)
+        # Return empty tensors - actual KV is in offload_engine's ring buffer
        return torch.empty(0), torch.empty(0)
    def can_allocate(self, seq: Sequence) -> bool:
        """Check if we can allocate blocks for a new sequence."""
@@ -231,6 +244,13 @@ class HybridKVCacheManager(KVCacheManager):
        seq.num_cached_tokens = 0
        seq.block_table.clear()
        # Clear decode tracking to prevent state pollution between requests
        self.clear_decode_tracking(seq)
        # Clear offload engine state (decode buffer, events)
        if self.offload_engine is not None:
            self.offload_engine.on_sequence_finished()
    def can_append(self, seq: Sequence) -> bool:
        """Check if we can append a token."""
        need_new_block = (len(seq) % self._block_size == 1)
@@ -279,8 +299,8 @@ class HybridKVCacheManager(KVCacheManager):
        """
        Prepare KV cache for attention computation.
-        In ring buffer mode, this is a no-op because chunked offload
+        In layer-wise offload mode, this is a no-op because KV transfers
-        paths handle H2D transfers directly in the attention layer.
+        are handled directly in model_runner's layer-by-layer methods.
        """
        pass
@@ -291,12 +311,12 @@ class HybridKVCacheManager(KVCacheManager):
        """
        Get GPU slot tables for sequences.
-        In ring buffer mode, all blocks are on CPU, so this raises an error
+        In layer-wise offload mode, all blocks are on CPU, so this raises an error
-        if called. Use run_chunked_offload_* methods instead.
+        if called. Use run_layerwise_offload_* methods instead.
        """
        raise RuntimeError(
-            "get_gpu_block_tables should not be called in ring buffer mode. "
+            "get_gpu_block_tables should not be called in layer-wise offload mode. "
-            "Use run_chunked_offload_prefill/decode instead."
+            "Use run_layerwise_offload_prefill/decode instead."
        )
    def post_attention_cleanup(
@@ -307,18 +327,18 @@ class HybridKVCacheManager(KVCacheManager):
        """
        Cleanup after attention.
-        In ring buffer mode, this is a no-op because offload is handled
+        In layer-wise offload mode, this is a no-op because offload is handled
-        directly in the chunked prefill/decode paths.
+        directly in model_runner's layer-by-layer methods.
        """
        pass
-    # ========== Ring Buffer CPU-primary Chunked Prefill Support ==========
+    # ========== Layer-wise Offload Support ==========
    def get_prefilled_cpu_blocks(self, seq: Sequence) -> List[int]:
        """
        Get list of CPU block IDs for blocks that have been prefilled.
-        Used for loading previous KV during chunked prefill.
+        Used for loading prefilled KV during decode.
        Returns:
            List of CPU block IDs in sequence order
@@ -329,17 +349,19 @@ class HybridKVCacheManager(KVCacheManager):
                block = self.logical_blocks[logical_id]
                if block.location == BlockLocation.CPU:
                    cpu_blocks.append(block.cpu_block_id)
-        # logger.debug(
+        # DEBUG: Log on first decode call
-        #     f"get_prefilled_cpu_blocks: prefilled_blocks={list(self.prefilled_blocks)}, "
+        logger.debug(
-        #     f"returned cpu_blocks={cpu_blocks}"
+            f"[DEBUG] get_prefilled_cpu_blocks: block_table={list(seq.block_table)}, "
-        # )
+            f"prefilled_blocks={list(self.prefilled_blocks)}, "
            f"returned cpu_blocks={cpu_blocks}"
        )
        return cpu_blocks
-    # ========== Ring Buffer CPU-primary support ==========
+    # ========== CPU Block Allocation ==========
    def allocate_cpu_only(self, seq: Sequence) -> None:
        """
-        Allocate CPU blocks for sequence (for ring buffer mode).
+        Allocate CPU blocks for sequence (for layer-wise offload mode).
        Unlike allocate(), here all blocks are allocated to CPU,
        GPU is only used as ring buffer for computation.
@@ -370,6 +392,10 @@ class HybridKVCacheManager(KVCacheManager):
            self.cpu_block_to_logical[cpu_block_id] = logical_id
            seq.block_table.append(logical_id)
        # DEBUG: Log allocated CPU blocks
        cpu_blocks = [self.logical_blocks[lid].cpu_block_id for lid in seq.block_table]
        logger.debug(f"[DEBUG] allocate_cpu_only: allocated cpu_blocks={cpu_blocks}")
            # NOTE: Prefix cache disabled in offload mode
            # If enabled, would compute hash and update:
            #   h = self.compute_hash(seq.block(i), prefix_hash)
@@ -417,6 +443,8 @@ class HybridKVCacheManager(KVCacheManager):
            if block.location == BlockLocation.CPU:
                cpu_block_ids.append(block.cpu_block_id)
                logical_ids.append(logical_id)
        # DEBUG: Log during prefill
        logger.debug(f"[DEBUG] get_all_cpu_blocks: returned cpu_block_ids={cpu_block_ids}")
        return cpu_block_ids, logical_ids
    def allocate_next_cpu_block(self, seq: Sequence) -> int:
@@ -468,20 +496,6 @@ class HybridKVCacheManager(KVCacheManager):
            return block.cpu_block_id
        return -1
    def get_write_slot_for_chunked_offload(self, seq: Sequence) -> int:
        """
        Get GPU slot for writing new KV during chunked offload decode.
        In ring buffer design, always use decode_slot (slot[0]) to write new KV.
        This avoids conflicts with loading operations which use slots[1:].
        Args:
            seq: Sequence
        Returns:
            GPU slot ID (always decode_slot = 0)
        """
        return self.offload_engine.decode_slot
    def get_decode_start_pos(self, seq: Sequence) -> int:
        """
@@ -503,6 +517,12 @@ class HybridKVCacheManager(KVCacheManager):
            # Decode starts at the next position
            prefill_len = len(seq) - 1  # Current len includes the new decode token
            self._decode_start_pos[seq_id] = prefill_len % self._block_size
            # DEBUG: Log first access
            logger.debug(
                f"[DEBUG] get_decode_start_pos FIRST ACCESS: seq_id={seq_id}, "
                f"len(seq)={len(seq)}, prefill_len={prefill_len}, "
                f"stored decode_start_pos={self._decode_start_pos[seq_id]}"
            )
        return self._decode_start_pos[seq_id]
    def reset_decode_start_pos(self, seq: Sequence) -> None:
@@ -535,6 +555,11 @@ class HybridKVCacheManager(KVCacheManager):
            # First decode step - store the prefill length
            # len(seq) - 1 because current len includes the first decode token
            self._prefill_len[seq_id] = len(seq) - 1
            # DEBUG: Log first access
            logger.debug(
                f"[DEBUG] get_prefill_len FIRST ACCESS: seq_id={seq_id}, "
                f"len(seq)={len(seq)}, stored prefill_len={self._prefill_len[seq_id]}"
            )
        return self._prefill_len[seq_id]
    def clear_decode_tracking(self, seq: Sequence) -> None:
@@ -547,6 +572,15 @@ class HybridKVCacheManager(KVCacheManager):
            seq: Sequence
        """
        seq_id = id(seq)
        # DEBUG: Log clearing and CPU blocks
        cpu_blocks = [self.logical_blocks[lid].cpu_block_id for lid in seq.block_table
                      if self.logical_blocks[lid].location == BlockLocation.CPU]
        logger.debug(
            f"[DEBUG] clear_decode_tracking: seq_id={seq_id}, "
            f"clearing decode_start_pos={self._decode_start_pos.get(seq_id, 'N/A')}, "
            f"prefill_len={self._prefill_len.get(seq_id, 'N/A')}, "
            f"cpu_blocks={cpu_blocks}"
        )
        self._decode_start_pos.pop(seq_id, None)
        self._prefill_len.pop(seq_id, None)
--- a/nanovllm/kvcache/offload_engine.py
+++ b/nanovllm/kvcache/offload_engine.py
--- a/nanovllm/kvcache/sparse/init.py
+++ b/nanovllm/kvcache/sparse/init.py
@@ -23,6 +23,7 @@ from nanovllm.config import SparsePolicyType
 from nanovllm.kvcache.sparse.policy import SparsePolicy, PolicyContext
 from nanovllm.kvcache.sparse.full_policy import FullAttentionPolicy
 from nanovllm.kvcache.sparse.quest import QuestPolicy, QuestConfig, BlockMetadataManager
 from nanovllm.kvcache.sparse.minference import MInferencePolicy
 def create_sparse_policy(policy_type: SparsePolicyType, **kwargs) -> SparsePolicy:
@@ -55,6 +56,15 @@ def create_sparse_policy(policy_type: SparsePolicyType, **kwargs) -> SparsePolic
        )
        return QuestPolicy(config)
    elif policy_type == SparsePolicyType.MINFERENCE:
        return MInferencePolicy(
            vertical_size=kwargs.get("vertical_size", 1000),
            slash_size=kwargs.get("slash_size", 6096),
            adaptive_budget=kwargs.get("adaptive_budget", 0.3),
            num_sink_tokens=kwargs.get("num_sink_tokens", 30),
            num_recent_diags=kwargs.get("num_recent_diags", 100),
        )
    else:
        raise ValueError(f"Unknown policy type: {policy_type}")
@@ -67,5 +77,6 @@ __all__ = [
    "QuestPolicy",
    "QuestConfig",
    "BlockMetadataManager",
    "MInferencePolicy",
    "create_sparse_policy",
 ]
--- a/nanovllm/kvcache/sparse/full_policy.py
+++ b/nanovllm/kvcache/sparse/full_policy.py
@@ -25,6 +25,7 @@ class FullAttentionPolicy(SparsePolicy):
    # Full attention supports both prefill and decode
    supports_prefill = True
    supports_decode = True
    requires_block_selection = False  # Load all blocks, no selective loading
    def select_blocks(
        self,
--- a/nanovllm/kvcache/sparse/minference.py
+++ b/nanovllm/kvcache/sparse/minference.py
@@ -0,0 +1,354 @@
 """
 MInference sparse attention policy.
 Implements vertical + slash sparse pattern estimation using the last 64 query tokens.
 Reference: MInference paper (https://arxiv.org/abs/2407.02490)
 """
 import math
 from typing import List, Tuple, Optional
 import torch
 import torch.nn.functional as F
 from nanovllm.kvcache.sparse.policy import SparsePolicy, PolicyContext
 class MInferencePolicy(SparsePolicy):
    """
    MInference sparse prefill policy using vertical + slash pattern.
    This policy estimates sparse attention patterns by analyzing attention
    scores from the last 64 query tokens, then selects:
    - Vertical: Key positions that are important across all queries
    - Slash: Diagonal bands (local context)
    The estimated pattern is then used to compute sparse attention.
    Note: This policy is designed for GPU-only prefill. For CPU offload,
    the pattern estimation and sparse attention will be handled differently.
    """
    supports_prefill = True
    supports_decode = False  # MInference is prefill-only sparse strategy
    requires_block_selection = False  # MInference only affects attention computation, not KV load
    def __init__(
        self,
        vertical_size: int = 1000,
        slash_size: int = 6096,
        adaptive_budget: Optional[float] = 0.3,
        num_sink_tokens: int = 30,
        num_recent_diags: int = 100,
    ):
        """
        Initialize MInference policy.
        Args:
            vertical_size: Number of vertical (column) positions to keep
            slash_size: Number of diagonal bands to keep
            adaptive_budget: If set, compute budget as fraction of seq_len
                            (overrides vertical_size and slash_size)
            num_sink_tokens: Number of initial sink tokens to always keep
            num_recent_diags: Number of recent diagonals to always keep
        """
        self.vertical_size = vertical_size
        self.slash_size = slash_size
        self.adaptive_budget = adaptive_budget
        self.num_sink_tokens = num_sink_tokens
        self.num_recent_diags = num_recent_diags
        # Cache for last-q causal mask
        self._last_q_mask_cache: dict = {}
    def _get_causal_mask(self, last_q: int, seq_len: int, device: torch.device) -> torch.Tensor:
        """Get causal mask for last-q attention."""
        cache_key = (last_q, seq_len, device)
        if cache_key not in self._last_q_mask_cache:
            # Create mask where last_q queries can attend to all previous positions
            # Shape: [last_q, seq_len]
            mask = torch.ones(last_q, seq_len, device=device, dtype=torch.bool)
            # Apply causal constraint for the last last_q positions
            # Query i (from last_q) can only attend to positions <= (seq_len - last_q + i)
            for i in range(last_q):
                mask[i, seq_len - last_q + i + 1:] = False
            self._last_q_mask_cache[cache_key] = mask
        return self._last_q_mask_cache[cache_key]
    def estimate_pattern(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        layer_id: int,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Estimate vertical + slash sparse pattern using last 64 query tokens.
        Memory-optimized for long sequences (64K+).
        Args:
            q: Query tensor [seq_len, num_heads, head_dim]
            k: Key tensor [seq_len, num_kv_heads, head_dim]
            layer_id: Current layer index (for potential layer-specific patterns)
        Returns:
            Tuple of (vertical_indices, slash_indices):
            - vertical_indices: [num_heads, vertical_size] - important K positions
            - slash_indices: [num_heads, slash_size] - diagonal offsets
        """
        seq_len = q.shape[0]
        num_heads = q.shape[1]
        head_dim = q.shape[2]
        num_kv_heads = k.shape[1]
        # Adaptive budget
        if self.adaptive_budget is not None:
            budget = int(seq_len * self.adaptive_budget)
            vertical_size = max(self.num_sink_tokens + 1, int(budget * 0.2))
            slash_size = max(self.num_recent_diags + 1, int(budget * 0.8))
        else:
            vertical_size = self.vertical_size
            slash_size = self.slash_size
        # Use last 64 Q tokens for estimation
        last_q = min(64, seq_len)
        q_last = q[-last_q:]  # [last_q, heads, dim] - this is a view, not a copy
        # Handle GQA: if num_kv_heads < num_heads, we need to expand K
        if num_kv_heads < num_heads:
            num_groups = num_heads // num_kv_heads
            k_work = k.repeat_interleave(num_groups, dim=1)
        else:
            k_work = k
        # Compute attention scores: [heads, last_q, seq_len]
        scale = 1.0 / math.sqrt(head_dim)
        qk = torch.einsum('qhd,khd->hqk', q_last, k_work) * scale
        # Free k_work if it was a copy
        if num_kv_heads < num_heads:
            del k_work
        # Apply causal mask for last positions (in-place)
        causal_mask = self._get_causal_mask(last_q, seq_len, q.device)
        qk.masked_fill_(~causal_mask.unsqueeze(0), float('-inf'))
        # Softmax (in-place where possible)
        qk = F.softmax(qk, dim=-1, dtype=torch.float32)
        # === Vertical pattern ===
        # Sum across query dimension -> importance of each K position
        vertical_scores = qk.sum(dim=1)  # [heads, seq_len]
        # Force keep first num_sink_tokens (attention sinks) - in-place
        vertical_scores[:, :self.num_sink_tokens] = float('inf')
        # Select top-k
        actual_vertical = min(vertical_size, seq_len)
        vertical_indices = vertical_scores.topk(actual_vertical, dim=-1).indices
        vertical_indices = vertical_indices.sort(dim=-1).values
        del vertical_scores
        # === Slash pattern ===
        # Create diagonal index matrix: [last_q, seq_len] with int32 to save memory
        q_indices = torch.arange(last_q, device=q.device, dtype=torch.int32).unsqueeze(1)
        k_indices = torch.arange(seq_len, device=q.device, dtype=torch.int32).unsqueeze(0)
        diag_indices = (seq_len - last_q + q_indices) - k_indices  # [last_q, seq_len]
        del q_indices
        # Create causal mask for slash computation
        q_pos = seq_len - last_q + torch.arange(last_q, device=q.device, dtype=torch.int32).unsqueeze(1)
        slash_causal_mask = k_indices <= q_pos
        del q_pos, k_indices
        # Clamp diagonal indices to valid range
        diag_indices = diag_indices.clamp(0, seq_len - 1)
        # Apply causal mask to qk (in-place) for slash computation
        qk[:, ~slash_causal_mask] = 0
        del slash_causal_mask
        # Accumulate scores per diagonal - process in batches to save memory
        slash_scores = torch.zeros(num_heads, seq_len, device=q.device, dtype=torch.float32)
        # Process heads in chunks to reduce peak memory for diag_indices_expanded
        chunk_size = min(8, num_heads)  # Process 8 heads at a time
        for h_start in range(0, num_heads, chunk_size):
            h_end = min(h_start + chunk_size, num_heads)
            n_heads_chunk = h_end - h_start
            # Expand diag_indices only for this chunk
            diag_chunk = diag_indices.unsqueeze(0).expand(n_heads_chunk, -1, -1).long()
            qk_chunk = qk[h_start:h_end]
            slash_scores[h_start:h_end].scatter_add_(
                1,
                diag_chunk.reshape(n_heads_chunk, -1),
                qk_chunk.reshape(n_heads_chunk, -1)
            )
            del diag_chunk, qk_chunk
        del diag_indices, qk
        # Force keep first num_recent_diags (in-place)
        slash_scores[:, :self.num_recent_diags] = float('inf')
        # Select top-k diagonal indices
        actual_slash = min(slash_size, seq_len)
        slash_indices = slash_scores.topk(actual_slash, dim=-1).indices
        slash_indices = slash_indices.sort(dim=-1).values
        del slash_scores
        return vertical_indices, slash_indices
    def select_blocks(
        self,
        available_blocks: List[int],
        ctx: PolicyContext,
    ) -> List[int]:
        """
        Select blocks for chunked CPU offload mode.
        For MInference in GPU-only mode, this method is not used.
        In CPU offload mode, it would select blocks based on the sparse pattern.
        For now, return all blocks (full attention fallback).
        """
        # MInference pattern is computed in attention.forward()
        # For CPU offload integration (Phase B), this would use the pattern
        return available_blocks
    def reset(self) -> None:
        """Reset policy state."""
        self._last_q_mask_cache.clear()
    def sparse_prefill_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        layer_id: int,
    ) -> torch.Tensor:
        """
        Compute MInference sparse attention for prefill.
        Uses vertical + slash pattern to compute sparse attention efficiently.
        Memory-optimized to handle long sequences (64K+) by freeing intermediate tensors.
        Args:
            q: Query tensor [seq_len, num_heads, head_dim]
            k: Key tensor [seq_len, num_kv_heads, head_dim]
            v: Value tensor [seq_len, num_kv_heads, head_dim]
            layer_id: Current transformer layer index
        Returns:
            Attention output [seq_len, num_heads, head_dim]
        """
        from minference.ops.pit_sparse_flash_attention_v2 import _triton_mixed_sparse_attention
        from minference.cuda import convert_vertical_slash_indexes
        seq_len = q.shape[0]
        num_heads = q.shape[1]
        head_dim = q.shape[2]
        num_kv_heads = k.shape[1]
        # Estimate sparse pattern (uses temporary memory for qk scores)
        vertical_indices, slash_indices = self.estimate_pattern(q, k, layer_id)
        # Free any cached memory from pattern estimation
        torch.cuda.empty_cache()
        # Triton sparse attention kernel parameters
        block_size_M = 64
        block_size_N = 64
        # Calculate padding
        pad = (block_size_M - seq_len) & (block_size_M - 1)
        need_head_pad = head_dim not in [16, 32, 64, 128, 256, 512]
        head_pad = (2 ** math.ceil(math.log2(head_dim)) - head_dim) if need_head_pad else 0
        # Handle GQA: expand K/V to match query heads
        # Do this BEFORE creating batched tensors to avoid double copies
        if num_kv_heads < num_heads:
            num_groups = num_heads // num_kv_heads
            # Use repeat_interleave for memory-efficient expansion
            k_work = k.repeat_interleave(num_groups, dim=1)
            v_work = v.repeat_interleave(num_groups, dim=1)
        else:
            k_work = k
            v_work = v
        # Transform Q to [batch, heads, seq, dim] format with padding in one step
        # This avoids creating intermediate copies
        if pad > 0 or head_pad > 0:
            q_batched = torch.nn.functional.pad(
                q.unsqueeze(0).transpose(1, 2),
                [0, head_pad, 0, pad, 0, 0, 0, 0]
            ).contiguous()
        else:
            q_batched = q.unsqueeze(0).transpose(1, 2).contiguous()
        # Transform K to batched format
        if pad > 0 or head_pad > 0:
            k_batched = torch.nn.functional.pad(
                k_work.unsqueeze(0).transpose(1, 2),
                [0, head_pad, 0, pad, 0, 0, 0, 0]
            ).contiguous()
        else:
            k_batched = k_work.unsqueeze(0).transpose(1, 2).contiguous()
        # Free k_work if it was a copy (GQA case)
        if num_kv_heads < num_heads:
            del k_work
        # Transform V to batched format
        if pad > 0 or head_pad > 0:
            v_batched = torch.nn.functional.pad(
                v_work.unsqueeze(0).transpose(1, 2),
                [0, head_pad, 0, pad, 0, 0, 0, 0]
            ).contiguous()
        else:
            v_batched = v_work.unsqueeze(0).transpose(1, 2).contiguous()
        # Free v_work if it was a copy (GQA case)
        if num_kv_heads < num_heads:
            del v_work
            torch.cuda.empty_cache()
        # Prepare indices for Triton kernel
        v_idx = vertical_indices.to(torch.int32).reshape((1, num_heads, -1))
        v_idx = v_idx.sort(dim=-1, descending=False)[0].contiguous()
        del vertical_indices
        s_idx = slash_indices.to(torch.int32).reshape((1, num_heads, -1))
        s_idx = s_idx.sort(dim=-1, descending=True)[0].contiguous()
        del slash_indices
        seqlens = torch.tensor([seq_len], dtype=torch.int32, device=q.device)
        sm_scale = head_dim ** -0.5
        # Convert vertical+slash indices to block sparse format
        block_count, block_offset, column_count, column_index = convert_vertical_slash_indexes(
            seqlens, v_idx, s_idx, seq_len, block_size_M, block_size_N,
        )
        del v_idx, s_idx
        # Call Triton mixed sparse attention kernel
        o = _triton_mixed_sparse_attention(
            q_batched, k_batched, v_batched, seqlens,
            block_count, block_offset, column_count, column_index,
            sm_scale, block_size_M, block_size_N,
        )
        # Free input tensors immediately after kernel call
        del q_batched, k_batched, v_batched
        del block_count, block_offset, column_count, column_index
        # Remove padding and convert back to [seq_len, num_heads, head_dim]
        o = o[..., :seq_len, :head_dim]
        o = o.transpose(1, 2).squeeze(0).contiguous()
        return o
    def __repr__(self) -> str:
        return (f"MInferencePolicy("
                f"adaptive_budget={self.adaptive_budget}, "
                f"vertical_size={self.vertical_size}, "
                f"slash_size={self.slash_size})")
--- a/nanovllm/kvcache/sparse/policy.py
+++ b/nanovllm/kvcache/sparse/policy.py
@@ -77,6 +77,12 @@ class SparsePolicy(ABC):
    supports_prefill: bool = True
    supports_decode: bool = True
    # Whether this policy requires selective block loading during decode
    # If True: OffloadEngine will call select_blocks() before loading KV from CPU
    # If False: OffloadEngine will load all blocks (select_blocks ignored for load)
    # Example: MInference=False (only affects attention), Quest=True (affects load)
    requires_block_selection: bool = False
    def initialize(
        self,
        num_layers: int,
@@ -183,5 +189,32 @@ class SparsePolicy(ABC):
        """
        pass
    def sparse_prefill_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        layer_id: int,
    ) -> torch.Tensor:
        """
        Compute sparse attention for prefill phase.
        This method is called when supports_prefill=True and the policy
        is used for GPU-only sparse prefill (no CPU offload).
        Args:
            q: Query tensor [seq_len, num_heads, head_dim]
            k: Key tensor [seq_len, num_kv_heads, head_dim]
            v: Value tensor [seq_len, num_kv_heads, head_dim]
            layer_id: Current transformer layer index
        Returns:
            Attention output [seq_len, num_heads, head_dim]
        """
        raise NotImplementedError(
            f"{self.__class__.__name__} does not implement sparse_prefill_attention. "
            "Set supports_prefill=False or implement this method."
        )
    def __repr__(self) -> str:
        return f"{self.__class__.__name__}()"
--- a/nanovllm/kvcache/sparse/quest.py
+++ b/nanovllm/kvcache/sparse/quest.py
@@ -158,6 +158,7 @@ class QuestPolicy(SparsePolicy):
    # Quest is decode-only
    supports_prefill = False
    supports_decode = True
    requires_block_selection = True  # Quest affects KV load strategy (selective block loading)
    def __init__(self, config: QuestConfig):
        """
--- a/nanovllm/layers/attention.py
+++ b/nanovllm/layers/attention.py
@@ -1,13 +1,8 @@
 import logging
 import torch
 import torch.cuda.nvtx
 from torch import nn
 from flash_attn.flash_attn_interface import flash_attn_varlen_func, flash_attn_with_kvcache
 from nanovllm.utils.context import get_context
 from nanovllm.kvcache.sparse.policy import PolicyContext
 logger = logging.getLogger(__name__)
 def store_kvcache(
@@ -60,12 +55,17 @@ def store_kvcache(
    valid_values_flat = valid_values.reshape(-1, D)
    # In-place scatter using index_copy_
    # 即使 valid_slots 为空张量，index_copy_ 也是安全的（不会修改数据）。
    k_cache_flat.index_copy_(0, valid_slots.long(), valid_keys_flat)
    v_cache_flat.index_copy_(0, valid_slots.long(), valid_values_flat)
 class Attention(nn.Module):
    """
    Attention layer for GPU-only mode.
    For CPU offload mode, attention is computed directly in model_runner's
    run_layerwise_offload_prefill/decode methods using FlashAttention.
    """
    def __init__(
        self,
@@ -87,635 +87,29 @@ class Attention(nn.Module):
        context = get_context()
        k_cache, v_cache = self.k_cache, self.v_cache
-        # Determine if we're in chunked offload mode
+        # Store KV to cache (for GPU-only mode)
-        is_chunked_offload = (
+        if k_cache.numel() and v_cache.numel():
-            context.is_chunked_prefill and
+            store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
            hasattr(context, 'kvcache_manager') and
            context.kvcache_manager is not None and
            hasattr(context.kvcache_manager, 'offload_engine')
        )
        #! Ensure synchronization before accessing k_cache/v_cache
        # torch.cuda.synchronize()
        #! =======================================================
        if is_chunked_offload and context.is_prefill:
            # Chunked prefill mode: write KV to per-layer prefill buffer (not GPU slot)
            # This enables fully async offloads since each layer has its own buffer.
            offload_engine = context.kvcache_manager.offload_engine
            compute_stream = offload_engine.compute_stream
            # Wait for default stream to ensure slot_mapping tensor transfer is complete
            compute_stream.wait_stream(torch.cuda.default_stream())
            with torch.cuda.stream(compute_stream):
                # Write KV to per-layer prefill buffer (contiguous write, no slot_mapping)
                # k, v shape: [num_tokens, kv_heads, head_dim]
                num_tokens = k.shape[0]
                offload_engine.prefill_k_buffer[self.layer_id, :num_tokens].copy_(k)
                offload_engine.prefill_v_buffer[self.layer_id, :num_tokens].copy_(v)
        elif is_chunked_offload:
            # Chunked decode mode: use compute_stream for store_kvcache
            # This ensures proper synchronization with per-layer offload
            compute_stream = context.kvcache_manager.offload_engine.compute_stream
            if k_cache.numel() and v_cache.numel():
                # CRITICAL: Wait for default stream to ensure slot_mapping tensor transfer is complete
                # slot_mapping is created with non_blocking=True on default stream, but we use it
                # on compute_stream. Without this sync, index_copy_ can get corrupted indices.
                compute_stream.wait_stream(torch.cuda.default_stream())
                with torch.cuda.stream(compute_stream):
                    store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
        else:
            # Normal mode: store on default stream
            if k_cache.numel() and v_cache.numel():
                store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
        if context.is_prefill:
-            if context.is_chunked_prefill:
+            if context.block_tables is not None:    # prefix cache
                # Chunked prefill: merge attention from previous KV
                o = self._chunked_prefill_attention(q, k, v, context)
            elif context.block_tables is not None:    # prefix cache
                k, v = k_cache, v_cache
                o = flash_attn_varlen_func(q, k, v,
                                           max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q,
                                           max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k,
                                           softmax_scale=self.scale, causal=True, block_table=context.block_tables)
            elif context.sparse_prefill_policy is not None:
                # Sparse prefill (GPU-only) - delegate to policy
                o = context.sparse_prefill_policy.sparse_prefill_attention(
                    q, k, v, self.layer_id
                )
            else:
                o = flash_attn_varlen_func(q, k, v,
                                           max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q,
                                           max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k,
                                           softmax_scale=self.scale, causal=True, block_table=context.block_tables)
        else:    # decode
-            if context.is_chunked_prefill:
+            o = flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache,
-                # Chunked decode: need to load all KV from CPU+GPU
+                                        cache_seqlens=context.context_lens, block_table=context.block_tables,
-                # Store current decode token to per-layer decode buffer
+                                        softmax_scale=self.scale, causal=True)
                # This is needed because GPU cache has no layer dimension,
                # so all layers would overwrite each other in decode_slot.
                kvcache_manager = context.kvcache_manager
                offload_engine = kvcache_manager.offload_engine
                pos_in_block = context.decode_pos_in_block
                # k, v shape: [1, kv_heads, head_dim]
                offload_engine.decode_k_buffer[self.layer_id, pos_in_block].copy_(k.squeeze(0))
                offload_engine.decode_v_buffer[self.layer_id, pos_in_block].copy_(v.squeeze(0))
                o = self._chunked_decode_attention(q, k, v, context)
            else:
                o = flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache,
                                            cache_seqlens=context.context_lens, block_table=context.block_tables,
                                            softmax_scale=self.scale, causal=True)
        return o
    def _chunked_prefill_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        context,
    ) -> torch.Tensor:
        """
        Compute attention with per-layer prefill buffer for async offload.
        Optimized design:
        - Current chunk's KV is written to per-layer prefill buffer (not GPU slot)
        - Previous chunks' KV are loaded from CPU using GPU slots
        - Each layer offloads from its own buffer - no waiting required!
        For each layer:
        1. Current chunk's KV is in prefill_buffer[layer_id] (just written by model)
        2. Load previous chunks from CPU using available slots (pipeline)
        3. Compute attention against previous KV (no causal mask)
        4. Compute attention against current KV from prefill buffer (causal)
        5. Merge all results using online softmax
        6. Async offload prefill buffer to CPU (no waiting!)
        """
        from nanovllm.kvcache.chunked_attention import flash_attn_with_lse, merge_attention_outputs
        current_chunk_idx = context.current_chunk_idx
        torch.cuda.nvtx.range_push(f"ChunkedPrefill: L{self.layer_id} Chunk{current_chunk_idx}")
        # q shape: [total_tokens, num_heads, head_dim]
        q_batched = q.unsqueeze(0)  # [1, total_tokens, heads, dim]
        num_tokens = k.shape[0]
        o_acc = None
        lse_acc = None
        kvcache_manager = context.kvcache_manager
        seq = context.chunked_seq if hasattr(context, 'chunked_seq') else None
        offload_engine = kvcache_manager.offload_engine if kvcache_manager is not None else None
        if kvcache_manager is not None and seq is not None and self.layer_id >= 0:
            # Get prefilled CPU blocks (blocks from previous chunks)
            cpu_block_table = kvcache_manager.get_prefilled_cpu_blocks(seq)
            # Apply sparse policy if enabled (Quest returns all blocks for prefill since query=None)
            sparse_policy = kvcache_manager.sparse_policy
            if cpu_block_table and sparse_policy is not None:
                num_chunks = getattr(context, 'num_chunks', current_chunk_idx + 1)
                policy_ctx = PolicyContext(
                    query_chunk_idx=current_chunk_idx,
                    num_query_chunks=num_chunks,
                    layer_id=self.layer_id,
                    query=None,  # Prefill typically doesn't use query for selection
                    is_prefill=True,
                    block_size=kvcache_manager.block_size,
                    total_kv_len=len(cpu_block_table) * kvcache_manager.block_size,
                )
                cpu_block_table = sparse_policy.select_blocks(
                    cpu_block_table, policy_ctx
                )
            if cpu_block_table:
                # Get available load slots (all slots can be used since we use prefill buffer)
                load_slots = list(range(offload_engine.num_ring_slots))
                pipeline_depth = len(load_slots)
                if pipeline_depth == 0:
                    # Only 1 slot total, cannot pipeline - use sync loading
                    o_acc, lse_acc = self._sync_load_previous_chunks(
                        q_batched, cpu_block_table, offload_engine
                    )
                else:
                    # Use ring buffer pipeline
                    o_acc, lse_acc = self._ring_buffer_pipeline_load(
                        q_batched, cpu_block_table, load_slots, offload_engine,
                        current_chunk_idx
                    )
        # Get compute stream for all attention operations
        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)
        if compute_stream is not None:
            with torch.cuda.stream(compute_stream):
                torch.cuda.nvtx.range_push(f"FlashAttn: L{self.layer_id} CurrentChunk (causal)")
                # Get KV from per-layer prefill buffer
                k_batched, v_batched = offload_engine.get_prefill_buffer_slice(self.layer_id, num_tokens)
                current_o, current_lse = flash_attn_with_lse(
                    q_batched,
                    k_batched,
                    v_batched,
                    softmax_scale=self.scale,
                    causal=True,
                )
                torch.cuda.nvtx.range_pop()
        else:
            torch.cuda.nvtx.range_push(f"FlashAttn: L{self.layer_id} CurrentChunk (causal)")
            k_batched = k.unsqueeze(0)
            v_batched = v.unsqueeze(0)
            current_o, current_lse = flash_attn_with_lse(
                q_batched,
                k_batched,
                v_batched,
                softmax_scale=self.scale,
                causal=True,
            )
            torch.cuda.nvtx.range_pop()
        # Merge with accumulated (all on compute_stream for consistency)
        if o_acc is None:
            final_o = current_o
        else:
            if compute_stream is not None:
                with torch.cuda.stream(compute_stream):
                    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:
                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()
        torch.cuda.nvtx.range_pop()  # ChunkedPrefill
        # Per-layer ASYNC offload: offload prefill buffer to CPU
        # No waiting required! Each layer has its own buffer and stream.
        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]
                # Async offload - no waiting, fully parallel across layers
                offload_engine.offload_prefill_buffer_async(
                    self.layer_id, cpu_block_id, num_tokens
                )
        # Sync default stream with compute_stream before returning
        # This ensures the result is ready for the rest of the model (layernorm, MLP)
        if compute_stream is not None:
            torch.cuda.default_stream().wait_stream(compute_stream)
        # Remove batch dimension: [1, total_tokens, heads, dim] -> [total_tokens, heads, dim]
        return final_o.squeeze(0)
    def _sync_load_previous_chunks(
        self,
        q_batched: torch.Tensor,
        cpu_block_table: list,
        offload_engine,
    ):
        """Synchronous loading fallback when pipeline_depth=0."""
        from nanovllm.kvcache.chunked_attention import flash_attn_with_lse, merge_attention_outputs
        o_acc, lse_acc = None, None
        compute_stream = offload_engine.compute_stream
        for block_idx, cpu_block_id in enumerate(cpu_block_table):
            # Load to slot 0 (single slot)
            offload_engine.load_to_slot_layer(0, self.layer_id, cpu_block_id)
            offload_engine.wait_slot_layer(0)
            # IMPORTANT: Must use compute_stream to match wait_slot_layer
            with torch.cuda.stream(compute_stream):
                prev_k, prev_v = offload_engine.get_kv_for_slot(0)
                prev_o, prev_lse = flash_attn_with_lse(
                    q_batched, prev_k, prev_v,
                    softmax_scale=self.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)
        return o_acc, lse_acc
    def _ring_buffer_pipeline_load(
        self,
        q_batched: torch.Tensor,
        cpu_block_table: list,
        load_slots: list,
        offload_engine,
        current_chunk_idx: int = -1,
    ):
        """
        Ring buffer async pipeline loading with double buffering.
        Uses compute_done events to ensure safe buffer reuse:
        - Before loading to slot X, wait for previous compute on slot X to finish
        - Before computing on slot X, wait for load to slot X to finish
        Timeline with 2 slots (A, B):
        ┌──────────────┐
        │ Load B0→A    │
        └──────────────┘
                       ┌──────────────┐ ┌──────────────┐
                       │ Load B1→B    │ │ Load B2→A    │ ...
                       └──────────────┘ └──────────────┘
                                      ↘               ↘
                        ┌──────────────┐ ┌──────────────┐
                        │ Compute(A)   │ │ Compute(B)   │ ...
                        └──────────────┘ └──────────────┘
        The load_to_slot_layer internally waits for compute_done[slot] before
        starting the transfer, ensuring no data race.
        """
        from nanovllm.kvcache.chunked_attention import flash_attn_with_lse, merge_attention_outputs
        num_blocks = len(cpu_block_table)
        if num_blocks == 0:
            return None, None
        pipeline_depth = len(load_slots)
        if pipeline_depth == 0:
            return None, None
        o_acc, lse_acc = None, None
        if pipeline_depth == 1:
            # Only 1 slot available, cannot pipeline - use synchronous mode
            # IMPORTANT: Must use compute_stream to match synchronization in
            # load_to_slot_layer (waits for compute_done) and wait_slot_layer
            slot = load_slots[0]
            compute_stream = offload_engine.compute_stream
            for block_idx in range(num_blocks):
                cpu_block_id = cpu_block_table[block_idx]
                offload_engine.load_to_slot_layer(slot, self.layer_id, cpu_block_id)
                offload_engine.wait_slot_layer(slot)
                with torch.cuda.stream(compute_stream):
                    # Debug: call hooks on compute_stream (synchronized with transfer)
                    if offload_engine.debug_mode:
                        offload_engine._call_debug_hooks(slot, self.layer_id, cpu_block_id)
                    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=self.scale,
                        causal=False,
                    )
                    # Record compute done so next load can safely reuse this slot
                    offload_engine.record_slot_compute_done(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)
            return o_acc, lse_acc
        # N-way pipeline: use ALL available slots for maximum overlap
        # Pipeline depth = num_slots - 1 (num_slots blocks in flight)
        num_slots = len(load_slots)
        # Phase 1: Pre-load up to num_slots blocks to fill the pipeline
        # This starts all transfers in parallel, utilizing full PCIe bandwidth
        num_preload = min(num_slots, num_blocks)
        for i in range(num_preload):
            offload_engine.load_to_slot_layer(load_slots[i], self.layer_id, cpu_block_table[i])
        # Phase 2: Main loop - compute and immediately reuse slot for next transfer
        # Use dedicated compute_stream (not default stream) to enable overlap with transfers
        compute_stream = offload_engine.compute_stream
        for block_idx in range(num_blocks):
            torch.cuda.nvtx.range_push(f"PipelineBlock: L{self.layer_id} B{block_idx}")
            # 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 (on compute_stream)
            offload_engine.wait_slot_layer(current_slot)
            # Compute attention on current slot's data
            # IMPORTANT: Use dedicated compute_stream to avoid implicit sync with default stream
            with torch.cuda.stream(compute_stream):
                # Debug: call hooks on compute_stream (synchronized with transfer)
                if offload_engine.debug_mode:
                    offload_engine._call_debug_hooks(current_slot, self.layer_id, cpu_block_id)
                torch.cuda.nvtx.range_push(f"FlashAttn: L{self.layer_id} PrevBlock{block_idx}")
                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=self.scale,
                    causal=False,
                )
                torch.cuda.nvtx.range_pop()
                # Record compute done - this allows the next transfer to safely overwrite this slot
                offload_engine.record_slot_compute_done(current_slot)
            # Immediately start loading the NEXT block into this slot (if more blocks remain)
            # Key insight: reuse current_slot immediately after compute is done!
            next_block_idx = block_idx + num_slots
            if next_block_idx < num_blocks:
                offload_engine.load_to_slot_layer(current_slot, self.layer_id, cpu_block_table[next_block_idx])
            # Merge with accumulated (also on compute_stream for consistency)
            with torch.cuda.stream(compute_stream):
                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)
            torch.cuda.nvtx.range_pop()  # PipelineBlock
        return o_acc, lse_acc
    def _chunked_decode_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        context,
    ) -> torch.Tensor:
        """
        Compute decode attention using cross-layer pipeline.
        Optimization: Uses double-buffered layer cache to overlap H2D transfer
        with computation across layers:
        - Layer N computes while Layer N+1's data is being loaded
        - Each layer only waits for its own data, not all layers' data
        This reduces effective latency from O(num_layers * transfer_time) to
        O(transfer_time + num_layers * compute_time) when transfer < compute.
        """
        from nanovllm.kvcache.chunked_attention import flash_attn_with_lse, merge_attention_outputs
        # q shape: [batch_size, num_heads, head_dim] (single decode token per sequence)
        q_batched = q.unsqueeze(1)  # [batch, 1, heads, dim]
        kvcache_manager = context.kvcache_manager
        seq = context.chunked_seq
        # Get only PREFILLED CPU blocks (exclude the current decode block)
        cpu_block_table = kvcache_manager.get_prefilled_cpu_blocks(seq)
        if self.layer_id == 0:
            logger.debug(f"Decode attention: cpu_block_table={cpu_block_table}, seq.block_table={list(seq.block_table)}")
        if not cpu_block_table:
            raise RuntimeError("Chunked decode attention failed: no prefilled CPU blocks available")
        # Calculate valid tokens in the last CPU block
        # CRITICAL: Use original prefill length, not current seq length!
        # CPU blocks are fixed after prefill, their content doesn't change during decode.
        block_size = kvcache_manager.block_size
        num_prefill_blocks = len(cpu_block_table)
        total_prefill_tokens = kvcache_manager.get_prefill_len(seq)  # Original prefill length
        last_block_valid_tokens = total_prefill_tokens % block_size
        if last_block_valid_tokens == 0 and total_prefill_tokens > 0:
            last_block_valid_tokens = block_size  # Last block was exactly full
        # Apply sparse policy if enabled (Quest does Top-K selection for decode)
        sparse_policy = kvcache_manager.sparse_policy
        if sparse_policy is not None:
            policy_ctx = PolicyContext(
                query_chunk_idx=0,
                num_query_chunks=1,
                layer_id=self.layer_id,
                query=q_batched,
                is_prefill=False,
                block_size=kvcache_manager.block_size,
                total_kv_len=len(cpu_block_table) * kvcache_manager.block_size,
            )
            cpu_block_table = sparse_policy.select_blocks(
                cpu_block_table, policy_ctx
            )
        offload_engine = kvcache_manager.offload_engine
        # Use cross-layer pipeline if active (initialized in model_runner)
        if offload_engine.is_pipeline_active():
            o_acc, lse_acc = self._decode_with_layer_pipeline(
                q_batched, cpu_block_table, offload_engine,
                block_size, last_block_valid_tokens
            )
        else:
            # Fallback to original ring buffer pipeline
            load_slots = offload_engine.decode_load_slots
            o_acc, lse_acc = self._decode_ring_buffer_pipeline(
                q_batched, cpu_block_table, load_slots, offload_engine,
                block_size, last_block_valid_tokens
            )
        # Now attend to accumulated decode tokens from per-layer decode buffer
        pos_in_block = context.decode_pos_in_block
        start_pos = context.decode_start_pos_in_block
        num_accumulated = pos_in_block - start_pos + 1
        # Sync compute_stream with default stream before reading decode_buffer
        compute_stream = offload_engine.compute_stream
        compute_stream.wait_stream(torch.cuda.default_stream())
        with torch.cuda.stream(compute_stream):
            if num_accumulated > 0:
                # Read from per-layer decode buffer
                decode_k = offload_engine.decode_k_buffer[self.layer_id, start_pos:pos_in_block+1]
                decode_v = offload_engine.decode_v_buffer[self.layer_id, start_pos:pos_in_block+1]
                decode_k = decode_k.unsqueeze(0)
                decode_v = decode_v.unsqueeze(0)
                decode_o, decode_lse = flash_attn_with_lse(
                    q_batched, decode_k, decode_v,
                    softmax_scale=self.scale,
                    causal=False,
                )
                if o_acc is None:
                    o_acc = decode_o
                else:
                    o_acc, _ = merge_attention_outputs(o_acc, lse_acc, decode_o, decode_lse)
        if o_acc is None:
            raise RuntimeError("Chunked decode attention failed: no KV available")
        # Sync back to default stream before returning
        torch.cuda.default_stream().wait_stream(compute_stream)
        return o_acc
    def _decode_ring_buffer_pipeline(
        self,
        q_batched: torch.Tensor,
        cpu_block_table: list,
        load_slots: list,
        offload_engine,
        block_size: int,
        last_block_valid_tokens: int,
    ):
        """
        Ring buffer pipeline for decode prefill loading (same mechanism as prefill).
        Loads one block at a time, computes attention, and merges results.
        Uses the same load_to_slot_layer / wait_slot_layer / get_kv_for_slot
        methods as prefill for proven correctness.
        """
        from nanovllm.kvcache.chunked_attention import flash_attn_with_lse, merge_attention_outputs
        num_blocks = len(cpu_block_table)
        if num_blocks == 0:
            return None, None
        if not load_slots:
            return None, None
        o_acc, lse_acc = None, None
        num_slots = len(load_slots)
        compute_stream = offload_engine.compute_stream
        # Phase 1: Pre-load up to num_slots blocks
        num_preload = min(num_slots, num_blocks)
        for i in range(num_preload):
            offload_engine.load_to_slot_layer(load_slots[i], self.layer_id, cpu_block_table[i])
        # Phase 2: Process blocks with pipeline
        for block_idx in range(num_blocks):
            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)
            with torch.cuda.stream(compute_stream):
                # Get KV from slot
                prev_k, prev_v = offload_engine.get_kv_for_slot(current_slot)
                # Handle partial last block
                is_last_block = (block_idx == num_blocks - 1)
                if is_last_block and last_block_valid_tokens < block_size:
                    prev_k = prev_k[:, :last_block_valid_tokens, :, :]
                    prev_v = prev_v[:, :last_block_valid_tokens, :, :]
                # Compute attention
                prev_o, prev_lse = flash_attn_with_lse(
                    q_batched, prev_k, prev_v,
                    softmax_scale=self.scale,
                    causal=False,
                )
                # Record compute done for slot reuse
                offload_engine.record_slot_compute_done(current_slot)
            # Start loading next block (pipeline)
            next_block_idx = block_idx + num_slots
            if next_block_idx < num_blocks:
                offload_engine.load_to_slot_layer(current_slot, self.layer_id, cpu_block_table[next_block_idx])
            # Merge with accumulated
            with torch.cuda.stream(compute_stream):
                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)
        return o_acc, lse_acc
    def _decode_with_layer_pipeline(
        self,
        q_batched: torch.Tensor,
        cpu_block_table: list,
        offload_engine,
        block_size: int,
        last_block_valid_tokens: int,
    ):
        """
        Decode using cross-layer pipeline for optimized H2D transfer.
        This method uses pre-loaded layer buffers instead of loading
        blocks one by one. The pipeline loads the next layer's data
        while the current layer computes, achieving transfer/compute overlap.
        The key insight is that each layer needs the SAME blocks but from
        different layers of CPU cache. By double-buffering and pipelining
        across layers, we reduce total latency.
        """
        from nanovllm.kvcache.chunked_attention import flash_attn_with_lse, merge_attention_outputs
        num_blocks = len(cpu_block_table)
        if num_blocks == 0:
            return None, None
        compute_stream = offload_engine.compute_stream
        # Get KV from pre-loaded layer buffer (triggers next layer loading)
        prev_k, prev_v = offload_engine.get_decode_layer_kv(self.layer_id, num_blocks)
        # prev_k, prev_v shape: [num_blocks, block_size, kv_heads, head_dim]
        # Reshape to [1, num_blocks * block_size, kv_heads, head_dim]
        total_tokens = num_blocks * block_size
        # Handle partial last block
        if last_block_valid_tokens < block_size:
            # Only use valid tokens from last block
            actual_tokens = (num_blocks - 1) * block_size + last_block_valid_tokens
            # Flatten and truncate
            prev_k_flat = prev_k.reshape(-1, prev_k.shape[-2], prev_k.shape[-1])[:actual_tokens]
            prev_v_flat = prev_v.reshape(-1, prev_v.shape[-2], prev_v.shape[-1])[:actual_tokens]
        else:
            prev_k_flat = prev_k.reshape(-1, prev_k.shape[-2], prev_k.shape[-1])
            prev_v_flat = prev_v.reshape(-1, prev_v.shape[-2], prev_v.shape[-1])
        # Add batch dimension: [1, total_tokens, kv_heads, head_dim]
        prev_k_batched = prev_k_flat.unsqueeze(0)
        prev_v_batched = prev_v_flat.unsqueeze(0)
        # Compute attention on all prefilled blocks at once
        with torch.cuda.stream(compute_stream):
            o_acc, lse_acc = flash_attn_with_lse(
                q_batched, prev_k_batched, prev_v_batched,
                softmax_scale=self.scale,
                causal=False,
            )
        return o_acc, lse_acc
--- a/nanovllm/models/init.py
+++ b/nanovllm/models/init.py
@@ -3,7 +3,13 @@
 from nanovllm.models.registry import register_model, get_model_class, MODEL_REGISTRY
 # Import models to trigger registration
-from nanovllm.models import qwen3
+# Qwen3 requires transformers>=4.51.0 for Qwen3Config
 try:
    from nanovllm.models import qwen3
 except ImportError as e:
    import warnings
    warnings.warn(f"Qwen3 model not available (requires transformers>=4.51.0): {e}")
 from nanovllm.models import llama
 __all__ = ["register_model", "get_model_class", "MODEL_REGISTRY"]
--- a/nanovllm/utils/context.py
+++ b/nanovllm/utils/context.py
@@ -1,5 +1,5 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
-from typing import Optional, List, Tuple, Any
+from typing import Any
 import torch
@@ -14,26 +14,9 @@ class Context:
    context_lens: torch.Tensor | None = None
    block_tables: torch.Tensor | None = None
-    # Chunked prefill support
+    # Sparse prefill attention support (GPU-only path)
-    is_chunked_prefill: bool = False
+    # When set, uses policy.sparse_prefill_attention() instead of FlashAttention
-    # Previous KV chunks info: List of (start_pos, end_pos) for blocks on CPU
+    sparse_prefill_policy: Any = None  # SparsePolicy instance with supports_prefill=True
    prev_kv_ranges: List[Tuple[int, int]] = field(default_factory=list)
    # Current chunk's position offset (for causal mask)
    chunk_offset: int = 0
    # Reference to kvcache manager for loading previous KV (HybridKVCacheManager)
    kvcache_manager: Any = None
    # Current layer's previous K/V chunks (loaded from CPU)
    # Set by model_runner before each layer's forward
    prev_kv_chunks: List[Tuple[torch.Tensor, torch.Tensor]] = field(default_factory=list)
    # Current sequence being processed (for chunked prefill to load KV)
    chunked_seq: Any = None
    # Position within block for decode (used for reading from Decode region)
    decode_pos_in_block: int = 0
    # Starting position within block where decode tokens began (for accumulated token tracking)
    # Used when batching decode offloads - we need to attend to all accumulated tokens
    decode_start_pos_in_block: int = 0
    # Current chunk index for ring buffer pipeline (prefill only)
    current_chunk_idx: int = 0
 _CONTEXT = Context()
@@ -52,14 +35,7 @@ def set_context(
    slot_mapping=None,
    context_lens=None,
    block_tables=None,
-    is_chunked_prefill=False,
+    sparse_prefill_policy=None,
    prev_kv_ranges=None,
    chunk_offset=0,
    kvcache_manager=None,
    chunked_seq=None,
    decode_pos_in_block=0,
    decode_start_pos_in_block=0,
    current_chunk_idx=0,
 ):
    global _CONTEXT
    _CONTEXT = Context(
@@ -71,14 +47,7 @@ def set_context(
        slot_mapping=slot_mapping,
        context_lens=context_lens,
        block_tables=block_tables,
-        is_chunked_prefill=is_chunked_prefill,
+        sparse_prefill_policy=sparse_prefill_policy,
        prev_kv_ranges=prev_kv_ranges or [],
        chunk_offset=chunk_offset,
        kvcache_manager=kvcache_manager,
        chunked_seq=chunked_seq,
        decode_pos_in_block=decode_pos_in_block,
        decode_start_pos_in_block=decode_start_pos_in_block,
        current_chunk_idx=current_chunk_idx,
    )
--- a/notes.md
+++ b/notes.md
@@ -0,0 +1,324 @@
 # Notes: Sparsity Integration into Layerwise Offload
 ## Current Architecture Analysis
 ### GPU-Only Path vs Offload Path
 | Aspect | GPU-Only | Layerwise Offload |
 |--------|----------|-------------------|
 | KV Storage | GPU blocks (paged) | CPU pinned + GPU ring buffer |
 | Prefill | All layers → then attention | Per-layer: attention → offload |
 | Decode | FlashAttn with block table | Ring buffer H2D → FlashAttn |
 | Sparse Support | MInference via `attention.py` | Not integrated |
 ### MInference Flow (GPU-Only)
 ```
 attention.py:101-105:
  if context.sparse_prefill_policy is not None:
      o = context.sparse_prefill_policy.sparse_prefill_attention(q, k, v, layer_id)
 minference.py:sparse_prefill_attention():
  1. estimate_pattern(q, k, layer_id) -> vertical_indices, slash_indices
  2. _triton_mixed_sparse_attention(q, k, v, indices)
  3. return output
 ```
 ### Quest Flow (GPU Block Mode)
 ```
 hybrid_manager.py (if using CPU offload with Quest):
  select_blocks(available_blocks, ctx) -> selected block IDs
  -> load selected blocks to GPU
  -> standard FlashAttn with loaded blocks
 ```
 ### Layerwise Offload Prefill Flow
 ```
 model_runner.py:run_layerwise_offload_prefill():
  for layer_id in range(num_layers):
      # QKV projection
      q, k, v = qkv_proj(hidden_ln)
      # RoPE
      q, k = rotary_emb(positions, q, k)
      # FULL attention (no sparsity!)
      attn_output = flash_attn_varlen_func(q, k, v, ...)
      # MLP
      hidden_states = mlp(attn_out + residual)
      # Sync offload ALL k, v to CPU
      for block_id in cpu_block_ids:
          k_cache_cpu[layer_id, block_id].copy_(k[start:end])
          v_cache_cpu[layer_id, block_id].copy_(v[start:end])
 ```
 ### Layerwise Offload Decode Flow
 ```
 model_runner.py:run_layerwise_offload_decode():
  # Preload first N layers to ring buffer
  for i in range(num_buffers):
      offload_engine.load_layer_kv_to_buffer(i, i, cpu_block_table, valid_tokens)
  for layer_id in range(num_layers):
      current_buffer = layer_id % num_buffers
      # Wait for buffer load
      offload_engine.wait_buffer_load(current_buffer)
      # Get prefilled KV from ring buffer (ALL blocks loaded)
      k_prefill, v_prefill = offload_engine.get_buffer_kv(current_buffer, total_prefill_tokens)
      # QKV for new token
      q, k_new, v_new = qkv_proj(hidden_ln)
      # Concat and full attention
      k_full = torch.cat([k_prefill, k_decode_prev, k_new])
      attn_output = flash_attn_varlen_func(q, k_full, v_full, ...)
      # Start loading next layer
      offload_engine.load_layer_kv_to_buffer(current_buffer, layer_id + num_buffers, ...)
 ```
 ## Integration Points
 ### 1. Prefill Sparse Integration Point
 **Location:** `model_runner.py:535-543`
 **Current:**
 ```python
 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,
 )
 ```
 **After Integration:**
 ```python
 if self.sparse_policy and self.sparse_policy.supports_offload_prefill:
    attn_output, k_sparse, v_sparse = self.sparse_policy.offload_prefill_attention(
        q, k, v, layer_id
    )
    k_to_offload = k_sparse if k_sparse is not None else k
    v_to_offload = v_sparse if v_sparse is not None else v
 else:
    attn_output = flash_attn_varlen_func(q, k, v, ...)
    k_to_offload, v_to_offload = k, v
 ```
 ### 2. Decode Sparse Integration Point
 **Location:** `model_runner.py:636-637` and `model_runner.py:704-706`
 **Current (preload):**
 ```python
 for i in range(num_preload):
    offload_engine.load_layer_kv_to_buffer(
        i, i, cpu_block_table, valid_tokens_per_block
    )
 ```
 **After Integration:**
 ```python
 for i in range(num_preload):
    layer_to_load = i
    if self.sparse_policy and self.sparse_policy.supports_offload_decode:
        # Prepare q for this layer (need to compute ahead)
        # OR: use previous layer's pattern as estimate
        selected_blocks = self.sparse_policy.select_offload_blocks(
            None,  # q not available yet at preload
            layer_to_load,
            cpu_block_table,
            valid_tokens_per_block
        )
    else:
        selected_blocks = cpu_block_table
    offload_engine.load_sparse_layer_kv_to_buffer(
        i, layer_to_load, selected_blocks, valid_tokens_per_block
    )
 ```
 **Challenge:** Q is not available during preload phase!
 **Solutions:**
 1. Skip sparse preload, only sparse for non-preloaded layers
 2. Use previous decode step's pattern as estimate
 3. Add preload hook to sparse policy
 ### 3. Offload Engine Extension
 **New Method in OffloadEngine:**
 ```python
 def load_sparse_layer_kv_to_buffer(
    self,
    buffer_idx: int,
    layer_id: int,
    selected_cpu_block_ids: List[int],
    original_valid_tokens: List[int],
 ) -> int:
    """
    Load only selected blocks from CPU to buffer.
    Returns:
        Total tokens loaded (may be less than full sequence)
    """
    stream = self.layer_load_streams[buffer_idx]
    with torch.cuda.stream(stream):
        stream.wait_event(self.buffer_compute_done_events[buffer_idx])
        # Build mapping: original block -> selected position
        offset = 0
        for i, cpu_block_id in enumerate(selected_cpu_block_ids):
            # Find original index to get valid tokens
            valid_tokens = original_valid_tokens[i]  # Need mapping
            self.layer_k_cache[buffer_idx, offset:offset+valid_tokens].copy_(
                self.k_cache_cpu[layer_id, cpu_block_id, :valid_tokens],
                non_blocking=True
            )
            # ... v_cache same
            offset += valid_tokens
        self.buffer_load_events[buffer_idx].record(stream)
    return offset  # Caller needs to know actual loaded tokens
 ```
 ## Metadata Flow for Quest
 ### During Prefill Offload
 **Current:** No metadata collection in offload path
 **Required:** Call `on_prefill_offload()` for each block
 ```python
 # In run_layerwise_offload_prefill()
 for i, cpu_block_id in enumerate(cpu_block_ids):
    start = i * block_size
    end = min(start + block_size, total_tokens)
    actual_size = end - start
    # BEFORE offload: update Quest metadata
    if self.sparse_policy and hasattr(self.sparse_policy, 'on_prefill_offload'):
        self.sparse_policy.on_prefill_offload(
            cpu_block_id, layer_id, k[start:end], actual_size
        )
    # Offload
    offload_engine.k_cache_cpu[layer_id, cpu_block_id, :actual_size].copy_(k[start:end])
    offload_engine.v_cache_cpu[layer_id, cpu_block_id, :actual_size].copy_(v[start:end])
 ```
 ### Quest Metadata Shape
 ```python
 # BlockMetadataManager
 key_min: [num_blocks, num_layers, num_kv_heads, head_dim]  # Min key per block per layer
 key_max: [num_blocks, num_layers, num_kv_heads, head_dim]  # Max key per block per layer
 ```
 **Memory:** 2 * num_blocks * num_layers * kv_heads * head_dim * 2 bytes
 - Example: 1000 blocks * 28 layers * 4 heads * 128 dim * 2 * 2 = ~57 MB
 ## Performance Considerations
 ### MInference Prefill Overhead
 | Operation | Time (64K seq) |
 |-----------|----------------|
 | Pattern estimation (last-64) | ~5ms |
 | Triton sparse attention | ~80ms |
 | Full FlashAttention | ~100ms |
 | **Net Speedup** | ~15-20% |
 ### Quest Decode Overhead
 | Operation | Time |
 |-----------|------|
 | Block scoring (GPU metadata) | ~0.1ms |
 | Top-K selection | ~0.05ms |
 | Sparse H2D load (8 blocks) | ~2ms |
 | Full H2D load (100 blocks) | ~20ms |
 | **Net Speedup** | ~10x H2D |
 ### Memory Trade-offs
 | Mode | GPU Memory | CPU Memory | H2D Bandwidth |
 |------|------------|------------|---------------|
 | Full offload | Ring buffer | Full KV | High |
 | Sparse offload | Ring buffer | Full KV | Low (subset) |
 | Aggressive sparse | Ring buffer | Sparse KV | Very low |
 ## Edge Cases
 ### 1. Short Sequences (< sparse threshold)
 ```python
 if total_tokens < sparse_threshold:
    # Fall back to full attention
    use_sparse = False
 ```
 ### 2. First Decode Step (no previous Q)
 Quest can't score blocks without Q. Options:
 - Use average embedding as proxy
 - Load all blocks for first step
 - Use prefill pattern as estimate
 ### 3. Variable Sequence Lengths in Batch
 Layerwise offload currently only supports batch_size=1:
 ```python
 assert len(seqs) == 1, "Layer-wise offload only supports single sequence"
 ```
 Sparse integration should maintain this constraint.
 ### 4. Ring Buffer vs Sparse Load Mismatch
 Ring buffer assumes fixed `total_prefill_tokens`:
 ```python
 k_prefill, v_prefill = offload_engine.get_buffer_kv(buffer_idx, total_prefill_tokens)
 ```
 Sparse load has variable token count. Need:
 ```python
 # Track actual loaded tokens per buffer
 loaded_tokens[buffer_idx] = sparse_load_count
 k_prefill, v_prefill = offload_engine.get_buffer_kv(buffer_idx, loaded_tokens[buffer_idx])
 ```
 ## Testing Strategy
 ### Unit Tests
 1. `test_sparse_policy_interface.py` - Verify new interface methods
 2. `test_minference_offload.py` - MInference in offload mode
 3. `test_quest_offload.py` - Quest block selection in offload mode
 ### Integration Tests
 1. `test_offload_sparse_e2e.py` - Full prefill+decode with sparsity
 2. `test_accuracy_comparison.py` - Compare outputs: full vs sparse
 ### Benchmarks
 1. `bench_offload_sparse.py` - Compare:
   - Full offload (baseline)
   - MInference prefill + Quest decode
   - Aggressive sparse offload
--- a/progress.md
+++ b/progress.md
@@ -1,76 +1,155 @@
-# Progress Log: Multi-Model Support
+# Progress Log: nanovllm 多请求状态污染问题
-## Session: 2026-01-10
+## Session: 2026-01-12
-### Initial Analysis Complete
+### 资源分配
-**Time**: Session start
+| 资源 | 分配 |
 |------|------|
 | **GPU** | **1** (严格限制，不可更改) |
-**Actions:**
+### 任务目标
-1. Read `nanovllm/engine/model_runner.py` - 确认硬编码位置 (line 35)
+研究 nanovllm CPU offload 模式下多请求之间状态影响导致准确率下降的问题。
 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
+### 10:00 - 启动分析
-| Phase | Status | Notes |
+**完成**:
-|-------|--------|-------|
+- [x] 读取 `docs/offload_accuracy_issue.md` 了解问题背景
-| 1. Model Registry | **COMPLETED** | `registry.py`, `__init__.py` |
+- [x] 激活 Serena MCP 项目
-| 2. Llama3 RoPE | **COMPLETED** | `rotary_embedding.py` |
+- [x] 获取关键组件符号概览
-| 3. Llama Model | **COMPLETED** | `llama.py` |
+
-| 4. ModelRunner | **COMPLETED** | Dynamic loading |
+**关键文件已分析**:
-| 5. Qwen3 Register | **COMPLETED** | `@register_model` decorator |
+- `nanovllm/kvcache/offload_engine.py` - OffloadEngine 类
-| 6. Testing | **COMPLETED** | Both Llama & Qwen3 pass |
+- `nanovllm/kvcache/hybrid_manager.py` - HybridKVCacheManager 类
 - `nanovllm/engine/model_runner.py` - ModelRunner 类
 - `nanovllm/engine/llm_engine.py` - LLMEngine 类
 - `nanovllm/engine/scheduler.py` - Scheduler 类
 ---
-## Test Results
+### 10:15 - 深入代码分析
-### 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
+| `OffloadEngine.__init__` | offload_engine.py:40-145 | 初始化所有 buffer，无 reset 方法 |
-Expected: 7492
+| `deallocate` | hybrid_manager.py:218-244 | 只清理逻辑块，不清理 OffloadEngine |
-Output: 7492
+| `clear_decode_tracking` | hybrid_manager.py:538-549 | 清理 tracking 字典，但未被调用 |
-Status: PASSED
+| `run_layerwise_offload_decode` | model_runner.py:867-1057 | 包含 decode buffer 读写逻辑 |
-Prefill: 3295 tok/s
+| `generate` | llm_engine.py:114-151 | 请求循环逻辑 |
-```
+| `postprocess` | scheduler.py:93-99 | 调用 deallocate |
 **关键发现 #1**: OffloadEngine 没有 reset() 方法
 **关键发现 #2**: deallocate() 没有调用 clear_decode_tracking()
 **关键发现 #3**: decode_buffer 在请求间不清理，可能导致状态污染
 ---
-## Files Modified This Session
+### 10:30 - 根因定位
-| File | Action | Description |
+**确认的问题**:
-|------|--------|-------------|
+
-| `nanovllm/models/registry.py` | created | Model registry with `@register_model` decorator |
+1. **decode buffer 残留**
-| `nanovllm/models/__init__.py` | created | Export registry functions, import models |
+   - 位置: `offload_engine.decode_k_buffer`, `decode_v_buffer`
-| `nanovllm/models/llama.py` | created | Llama model implementation |
+   - 写入: `model_runner.py:1010-1013`
-| `nanovllm/models/qwen3.py` | modified | Added `@register_model` decorator |
+   - 读取: `model_runner.py:969-976`
-| `nanovllm/layers/rotary_embedding.py` | modified | Added Llama3 RoPE scaling |
+   - 问题: 旧请求的 KV 数据可能被新请求读取
-| `nanovllm/engine/model_runner.py` | modified | Dynamic model loading via registry |
+
-| `.claude/rules/gpu-testing.md` | created | GPU testing rules |
+2. **tracking 字典未清理**
-| `task_plan.md` | created | Implementation plan |
+   - 位置: `hybrid_manager._decode_start_pos`, `_prefill_len`
-| `findings.md` | created | Technical findings |
+   - 问题: 使用 `id(seq)` 作为 key，可能重用
-| `progress.md` | created | Progress tracking |
+
 3. **缺失的清理调用**
   - `clear_decode_tracking()` 在 `deallocate()` 中未被调用
 ---
 ### 10:45 - 创建规划文件
 **创建的文件**:
 - [x] `task_plan.md` - 完整的任务规划和阶段
 - [x] `findings.md` - 详细的代码分析发现
 - [x] `progress.md` - 本文件
 ---
 ### 11:00 - Sequential Thinking 深入分析
 **使用 sequential thinking 验证分析结果**:
 - 确认 deallocate() 确实没有调用 clear_decode_tracking()
 - 分析 _decode_start_pos 和 _prefill_len 字典的生命周期
 - 确定 id(seq) 重用是问题的触发条件
 ---
 ### 11:15 - 完成规划文件
 **更新的文件**:
 - [x] `task_plan.md` - 添加完整的 debug 方案和实施计划
 - [x] `findings.md` - 详细的代码分析和修复方向
 - [x] `progress.md` - 更新到当前进度
 ---
 ## 下一步 (待用户确认)
 **执行顺序**:
 1. **实施修复** - 修改 `deallocate()` 添加 `clear_decode_tracking(seq)`
 2. **快速验证** - 20 样本连续执行（一次调用，不重启框架）→ 目标 20/20
 3. **完整验证** - 100 样本 → 目标 100/100 (最终验收)
 4. **防御性修复** (可选) - 添加 `OffloadEngine.on_sequence_finished()`
 **核心修改** (一行代码):
 ```python
 # hybrid_manager.py:deallocate() 末尾添加
 self.clear_decode_tracking(seq)
 ```
 **验收标准**:
 | 测试 | 样本数 | 通过要求 |
 |------|--------|----------|
 | 快速验证 | 20 | 20/20 (100%) |
 | 完整验证 | 100 | 100/100 (100%) |
 ---
 ## 错误记录
 | 时间 | 错误 | 解决方案 |
 |------|------|----------|
 | 10:05 | Serena MCP 未激活 | 调用 activate_project |
 ---
 ## 文件修改记录
 | 文件 | 操作 | 状态 |
 |------|------|------|
 | task_plan.md | 创建+更新 | 完成 |
 | findings.md | 创建 | 完成 |
 | progress.md | 创建+更新 | 完成 |
 ---
 ## 分析结论
 **重要澄清**: nanovllm offload 模式**不支持 batch**，只能单个 request 顺序执行。问题出在**请求切换**时状态清理不完整。
 **根本原因已确认**: `deallocate()` 没有调用 `clear_decode_tracking()`，导致 `_decode_start_pos` 和 `_prefill_len` 字典残留，当 Python 对象 ID 重用时，新请求会错误地使用旧请求的配置。
 **修复方案已设计**: 在 `deallocate()` 末尾添加 `self.clear_decode_tracking(seq)` 调用。
 ---
 ## 关键理解
 问题不是 "batch 处理"，而是：
 ```
 Request A 完成 → deallocate(A) [状态未完全清理] → Request B 开始 → B 读到 A 的残留状态
 ```
--- a/task_plan.md
+++ b/task_plan.md
@@ -1,144 +1,359 @@
-# Task Plan: Multi-Model Support for nanovllm
+# Task Plan: nanovllm CPU Offload 多请求状态污染问题
-## Goal
+## 问题概述
 扩展 nanovllm 框架以支持多种模型（当前只支持 Qwen3），特别是添加 Llama-3.1-8B-Instruct 支持，并建立可扩展的模型添加范式。
-## Current State Analysis
+**重要说明**: nanovllm offload 模式目前**不支持 batch**，只能单个 request 顺序执行。问题出在**请求切换**时的状态清理。
-### 硬编码问题位置
+| 模式 | 测试方式 | 准确率 |
- `nanovllm/engine/model_runner.py:35`: 直接实例化 `Qwen3ForCausalLM(hf_config)`
+|------|----------|--------|
- `nanovllm/engine/model_runner.py:9`: 硬编码导入 `from nanovllm.models.qwen3 import Qwen3ForCausalLM`
+| CPU Offload | 独立进程 (每请求一个进程) | **100%** |
 | CPU Offload | 同进程顺序多请求 | 66% |
 | Non-Offload | 同进程顺序多请求 | 100% |
-### Qwen3 vs Llama 3.1 架构差异
+**结论**: 单请求推理正确，问题在于**请求切换**时状态清理不完整。
 | 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
 ---
-## Phases
+## Phase 1: 代码分析 (complete)
-### Phase 1: Create Model Registry Pattern [pending]
+### 1.1 识别状态管理组件
 **Files to modify:**
 - `nanovllm/models/__init__.py` (new)
 - `nanovllm/models/registry.py` (new)
-**Tasks:**
+**已分析的关键组件**:
 1. 创建模型注册表机制
 2. 定义模型注册装饰器 `@register_model`
 3. 实现 `get_model_class(hf_config)` 函数，根据 `architectures` 字段自动选择模型
-**Design:**
+| 组件 | 文件 | 状态数据 |
-```python
+|------|------|----------|
-MODEL_REGISTRY: dict[str, type] = {}
+| `OffloadEngine` | `nanovllm/kvcache/offload_engine.py` | ring buffer, decode buffer, CUDA events |
 | `HybridKVCacheManager` | `nanovllm/kvcache/hybrid_manager.py` | logical blocks, prefilled_blocks, _decode_start_pos, _prefill_len |
 | `LLMEngine` | `nanovllm/engine/llm_engine.py` | generate() 循环，请求生命周期 |
 | `Scheduler` | `nanovllm/engine/scheduler.py` | postprocess() 调用 deallocate() |
-def register_model(*architectures):
+### 1.2 请求生命周期分析
    """Decorator to register a model class for given architecture names."""
    def decorator(cls):
        for arch in architectures:
            MODEL_REGISTRY[arch] = cls
        return cls
    return decorator
-def get_model_class(hf_config) -> type:
+```
-    """Get model class based on HF config architectures."""
+generate()
-    for arch in hf_config.architectures:
+  → 多个请求添加到 scheduler
-        if arch in MODEL_REGISTRY:
+  → while not finished:
-            return MODEL_REGISTRY[arch]
+      → schedule() 获取下一批 seqs
-    raise ValueError(f"Unsupported architecture: {hf_config.architectures}")
+      → model_runner.run() 执行推理
      → postprocess() 处理完成的请求
          → 如果完成: kvcache_manager.deallocate(seq)
 ```
-### Phase 2: Add Llama3 RoPE Scaling Support [pending]
+---
 **Files to modify:**
 - `nanovllm/layers/rotary_embedding.py`
-**Tasks:**
+## Phase 2: 根本原因分析 (complete)
-1. 实现 `Llama3RotaryEmbedding` 类，支持 llama3 rope_type
+
-2. 修改 `get_rope()` 函数，根据 rope_scaling 类型选择实现
+### 2.1 核心问题: OffloadEngine 缺少 reset() 方法
-3. 保持向后兼容（rope_scaling=None 使用原实现）
+
 **关键发现**: `OffloadEngine` 没有任何重置/清理方法！
 当请求完成时，`HybridKVCacheManager.deallocate()` 被调用，但它只清理：
 - 逻辑块状态 (`block.reset()`)
 - 物理块引用 (`free_cpu_blocks`, `cpu_block_to_logical`)
 - prefilled_blocks 集合
 - _decode_start_pos / _prefill_len 字典
 **未被清理的状态** (存在于 OffloadEngine):
 | 状态 | Shape | 问题 |
 |------|-------|------|
 | `layer_k_cache` | [num_buffers, max_seq_len, kv_heads, head_dim] | 包含旧请求的 KV |
 | `layer_v_cache` | [num_buffers, max_seq_len, kv_heads, head_dim] | 包含旧请求的 KV |
 | `decode_k_buffer` | [num_layers, block_size, kv_heads, head_dim] | 包含旧请求的 decode KV |
 | `decode_v_buffer` | [num_layers, block_size, kv_heads, head_dim] | 包含旧请求的 decode KV |
 ### 2.2 具体污染场景
 在 `run_layerwise_offload_decode()` (model_runner.py:867-1057):
 **Llama3 RoPE Scaling Formula:**
 ```python
-# From transformers:
+# 第 969-976 行: 读取之前的 decode KV
-# low_freq_factor, high_freq_factor, original_max_position_embeddings
+if num_prev_decode_tokens > 0:
-# Adjust frequencies based on wavelength thresholds
+    k_decode_prev, v_decode_prev = offload_engine.get_decode_kv(
        layer_id, decode_start_pos, pos_in_block
    )
    ring_k[...].copy_(k_decode_prev)  # 可能读取旧请求的数据!
 ```
-### Phase 3: Implement Llama Model [pending]
+**场景**:
-**Files to create:**
+1. 请求 A (32K tokens) 完成，decode_buffer 保留其 KV 数据
- `nanovllm/models/llama.py`
+2. 请求 B 开始，其 `decode_start_pos` 可能非零（如果继承了旧状态）
 3. 请求 B 在第一个 decode step 时错误地读取了请求 A 的 decode buffer 数据
-**Tasks:**
+### 2.3 潜在问题点
 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]
+1. **decode_start_pos 计算错误**:
-**Files to modify:**
+   - `get_decode_start_pos()` 使用 `id(seq)` 作为 key
- `nanovllm/engine/model_runner.py`
+   - Python 对象 ID 可能在请求之间重用
   - 如果新 seq 对象的 ID 与旧 seq 相同，可能错误继承旧的 start_pos
-**Tasks:**
+2. **decode buffer 残留数据**:
-1. 移除硬编码 `from nanovllm.models.qwen3 import Qwen3ForCausalLM`
+   - 如果 `pos_in_block` 在新请求中与旧请求重叠
-2. 导入 `from nanovllm.models import get_model_class`
+   - `get_decode_kv()` 会返回旧请求的数据
 3. 替换 `self.model = Qwen3ForCausalLM(hf_config)` 为:
   ```python
   model_class = get_model_class(hf_config)
   self.model = model_class(hf_config)
   ```
-### Phase 5: Register Qwen3 Model [pending]
+3. **ring buffer 残留数据**:
-**Files to modify:**
+   - 虽然每次 decode 会从 CPU 加载，但 decode buffer 的数据会被复制过来
- `nanovllm/models/qwen3.py`
+   - 如果 decode buffer 有残留，会污染 ring buffer
 **Tasks:**
 1. 导入 `from nanovllm.models.registry import register_model`
 2. 添加 `@register_model("Qwen3ForCausalLM", "Qwen2ForCausalLM")` 装饰器
 ### Phase 6: Test with Llama-3.1-8B-Instruct [pending]
 **Files:**
 - `tests/test_needle.py` (existing, use for validation)
 **Tasks:**
 1. 运行 needle 测试: `python tests/test_needle.py --model ~/models/Llama-3.1-8B-Instruct`
 2. 验证模型加载正确
 3. 验证推理输出正确
 ---
-## Errors Encountered
+## Phase 3: Debug 方案设计 (complete)
-| Error | Attempt | Resolution |
+
-|-------|---------|------------|
+### 3.1 确认的根本原因
-| (none yet) | | |
+
 通过代码分析，确认了两个根本原因：
 **根本原因 1 (主要)**: `deallocate()` 不调用 `clear_decode_tracking()`
 - 位置: `hybrid_manager.py:218-244`
 - 影响: `_decode_start_pos` 和 `_prefill_len` 字典残留
 - 后果: 如果 `id(seq)` 重用，返回错误的 decode 配置
 **根本原因 2 (次要)**: decode_buffer 不清理
 - 位置: `offload_engine.py`
 - 影响: `decode_k_buffer/v_buffer` 保留旧 KV
 - 后果: 可能被根本原因 1 触发读取
 ### 3.2 Debug 方案 A: 验证字典残留 (推荐先做)
 **目标**: 验证 `_decode_start_pos` 字典是否有残留
 **诊断代码** (添加到 `hybrid_manager.py`):
 ```python
 # 在 get_decode_start_pos() 开头添加
 def get_decode_start_pos(self, seq: Sequence) -> int:
    seq_id = id(seq)
    # DEBUG: 检查是否命中旧值
    if seq_id in self._decode_start_pos:
        logger.warning(f"[DEBUG] get_decode_start_pos: CACHE HIT! seq_id={seq_id}, "
                       f"cached_value={self._decode_start_pos[seq_id]}, "
                       f"expected={(len(seq) - 1) % self._block_size}")
    # ... 原有逻辑
 ```
 **诊断代码** (添加到 `deallocate()` 末尾):
 ```python
 def deallocate(self, seq: Sequence) -> None:
    # ... 现有逻辑 ...
    # DEBUG: 打印未清理的状态
    seq_id = id(seq)
    if seq_id in self._decode_start_pos:
        logger.warning(f"[DEBUG] deallocate: _decode_start_pos NOT CLEARED! "
                       f"seq_id={seq_id}, value={self._decode_start_pos[seq_id]}")
 ```
 ### 3.3 Debug 方案 B: 最小复现测试
 **文件**: `tests/test_multi_request_offload_debug.py`
 ```python
 """最小复现批量模式失败"""
 import os
 import sys
 sys.path.insert(0, os.getcwd())
 from nanovllm import LLM
 from nanovllm.sampling import SamplingParams
 # 使用 RULER NIAH 的两个样本
 PROMPTS = [
    # Sample 0 (通常成功)
    "...",  # 从 niah_single_1_32k.jsonl 加载
    # Sample 1 (通常失败)
    "...",
 ]
 EXPECTED = ["8930103", "4194548"]
 def main():
    llm = LLM(
        "~/models/Llama-3.1-8B-Instruct",
        max_model_len=33792,
        max_num_batched_tokens=33792,
        enable_cpu_offload=True,
        num_gpu_blocks=4,
        kvcache_block_size=1024,
        enforce_eager=True,
    )
    params = SamplingParams(temperature=0.1, max_tokens=50)
    # 连续处理两个请求
    for i, (prompt, expected) in enumerate(zip(PROMPTS, EXPECTED)):
        print(f"\n{'='*60}")
        print(f"Sample {i}: Expected = {expected}")
        # 打印关键状态
        kvm = llm.model_runner.kvcache_manager
        print(f"  _decode_start_pos 字典大小: {len(kvm._decode_start_pos)}")
        print(f"  _prefill_len 字典大小: {len(kvm._prefill_len)}")
        outputs = llm.generate([prompt], params, use_tqdm=False)
        output_text = outputs[0]["text"]
        passed = expected in output_text
        print(f"  Output: {output_text[:100]}...")
        print(f"  Status: {'PASS' if passed else 'FAIL'}")
 if __name__ == "__main__":
    main()
 ```
 ### 3.4 Debug 方案 C: 快速修复验证
 **目标**: 验证修复 `deallocate()` 是否解决问题
 **修改** (`hybrid_manager.py:218-244`):
 ```python
 def deallocate(self, seq: Sequence) -> None:
    """Release all blocks for a sequence."""
    for logical_id in reversed(seq.block_table):
        # ... 现有逻辑 ...
    seq.num_cached_tokens = 0
    seq.block_table.clear()
    # === 新增: 清理 decode tracking ===
    self.clear_decode_tracking(seq)
 ```
 **验证命令**:
 ```bash
 CUDA_VISIBLE_DEVICES=0 PYTHONPATH=.:$PYTHONPATH python tests/test_ruler_niah.py \
    --model ~/models/Llama-3.1-8B-Instruct \
    --enable-offload \
    --sample-indices 0,1,2,3,4 \
    --verbose
 ```
 ### 3.5 Debug 方案 D: 添加 OffloadEngine 清理 (防御性)
 **目标**: 进一步隔离请求状态
 **添加方法** (`offload_engine.py`):
 ```python
 def on_sequence_finished(self):
    """清理请求完成后的状态"""
    # 清零 decode buffer (防止残留数据被读取)
    self.decode_k_buffer.zero_()
    self.decode_v_buffer.zero_()
    logger.debug("OffloadEngine: decode buffer cleared")
 ```
 **调用点** (`hybrid_manager.py:deallocate` 末尾):
 ```python
 # 清理 OffloadEngine 状态
 if self.offload_engine is not None:
    self.offload_engine.on_sequence_finished()
 ```
 ---
-## Success Criteria
+## Phase 4: 实施计划 (pending)
- [x] 分析完成：理解当前架构和需要的改动
+
- [ ] Phase 1: 模型注册表实现
+### 推荐执行顺序
- [ ] Phase 2: Llama3 RoPE scaling 支持
+
- [ ] Phase 3: Llama 模型实现
+1. **Step 4.1**: 实施修复
- [ ] Phase 4: ModelRunner 动态加载
+   - 修改 `hybrid_manager.py:deallocate()` 添加 `clear_decode_tracking(seq)`
- [ ] Phase 5: Qwen3 模型注册
+
- [ ] Phase 6: Llama needle 测试通过
+2. **Step 4.2**: 快速验证 (20 样本连续执行)
   - **一次调用** `test_ruler_niah.py`，连续执行 20 个样本
   - **不重启框架**，验证请求切换是否正确
   - 目标: 20/20 全部通过
 3. **Step 4.3**: 完整验证 (100 样本)
   - 运行 100 个样本的 RULER NIAH 测试
   - 目标: 100/100 全部通过 (准确率从 66% → 100%)
 4. **Step 4.4**: 防御性修复 (可选)
   - 添加 `OffloadEngine.on_sequence_finished()` 方法
   - 清零 decode buffer 作为额外保险
 ### 具体修改
 **文件 1**: `nanovllm/kvcache/hybrid_manager.py`
 位置: `deallocate()` 方法末尾 (第 244 行后)
 ```python
 def deallocate(self, seq: Sequence) -> None:
    """Release all blocks for a sequence."""
    for logical_id in reversed(seq.block_table):
        # ... 现有逻辑 (218-242 行) ...
    seq.num_cached_tokens = 0
    seq.block_table.clear()
    # ============ 新增: 清理 decode tracking ============
    self.clear_decode_tracking(seq)
 ```
 **文件 2** (可选): `nanovllm/kvcache/offload_engine.py`
 位置: 在类末尾添加新方法
 ```python
 def on_sequence_finished(self):
    """清理请求完成后的状态 (防御性清理)"""
    self.decode_k_buffer.zero_()
    self.decode_v_buffer.zero_()
 ```
 ---
-## Notes
+## 关键文件清单
- 保持现有 Qwen3 功能不变
+
- 遵循现有代码风格
+| 文件 | 相关行号 | 说明 |
- 复用现有 layers 组件（Linear, RMSNorm, Embedding 等）
+|------|----------|------|
- 只添加必要的代码，不过度工程化
+| `nanovllm/kvcache/hybrid_manager.py` | 218-244 | `deallocate()` - **需要修改** |
 | `nanovllm/kvcache/hybrid_manager.py` | 538-549 | `clear_decode_tracking()` - 已存在 |
 | `nanovllm/kvcache/hybrid_manager.py` | 485-505 | `get_decode_start_pos()` - 问题读取点 |
 | `nanovllm/kvcache/hybrid_manager.py` | 519-537 | `get_prefill_len()` - 问题读取点 |
 | `nanovllm/kvcache/offload_engine.py` | 40-145 | `__init__` - 状态初始化 |
 | `nanovllm/kvcache/offload_engine.py` | (新增) | `on_sequence_finished()` - 可选防御 |
 | `nanovllm/engine/model_runner.py` | 867-1057 | `run_layerwise_offload_decode()` |
 | `nanovllm/engine/model_runner.py` | 969-976 | decode buffer 读取 (污染点) |
 ---
 ## 验证命令
 **指定 GPU: 1** (严格限制，不可更改)
 ```bash
 # 快速验证 (20 样本连续执行，不重启框架)
 # 目标: 20/20 通过
 CUDA_VISIBLE_DEVICES=1 PYTHONPATH=.:$PYTHONPATH python tests/test_ruler_niah.py \
    --model ~/models/Llama-3.1-8B-Instruct \
    --enable-offload \
    --sample-indices 0-19 \
    --verbose
 # 完整验证 (100 样本)
 # 目标: 100/100 通过 (最终验收)
 CUDA_VISIBLE_DEVICES=1 PYTHONPATH=.:$PYTHONPATH python tests/test_ruler_niah.py \
    --model ~/models/Llama-3.1-8B-Instruct \
    --enable-offload \
    --quiet
 ```
 **验收标准**:
 | 测试 | 样本数 | 通过要求 | 说明 |
 |------|--------|----------|------|
 | 快速验证 | 20 | 20/20 (100%) | 一次调用，连续执行，验证请求切换 |
 | 完整验证 | 100 | 100/100 (100%) | 最终验收 |
 ---
 ## 当前状态
 - [x] Phase 1: 代码分析
 - [x] Phase 2: 根本原因分析
 - [x] Phase 3: Debug 方案设计
 - [x] Phase 4: 实施计划 ✅ 100/100 PASSED
 ### 验证结果
 | 测试 | 结果 | 日期 |
 |------|------|------|
 | 20 样本快速验证 | ✅ 20/20 (100%) | 2026-01-13 |
 | 100 样本完整验证 | ✅ 100/100 (100%) | 2026-01-13 |
--- a/tests/run_parallel_niah.sh
+++ b/tests/run_parallel_niah.sh
@@ -0,0 +1,112 @@
 #!/bin/bash
 # Run NIAH tests in parallel on 6 GPUs
 # This tests the dynamic port allocation fix
 set -e
 MODEL="${1:-/home/zijie/models/Llama-3.1-8B-Instruct}"
 PROJECT_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
 echo "=========================================="
 echo "Parallel NIAH Test on 6 GPUs"
 echo "=========================================="
 echo "Model: $MODEL"
 echo "Project: $PROJECT_ROOT"
 echo ""
 # Sample distribution (100 samples total):
 # GPU 0: 0-16   (17 samples)
 # GPU 1: 17-33  (17 samples)
 # GPU 2: 34-50  (17 samples)
 # GPU 3: 51-67  (17 samples)
 # GPU 4: 68-83  (16 samples)
 # GPU 5: 84-99  (16 samples)
 declare -a RANGES=("0-16" "17-33" "34-50" "51-67" "68-83" "84-99")
 declare -a PIDS=()
 # Create log directory
 LOG_DIR="$PROJECT_ROOT/logs"
 mkdir -p "$LOG_DIR"
 # Start all 6 processes
 for gpu in {0..5}; do
    range="${RANGES[$gpu]}"
    log_file="$LOG_DIR/gpu${gpu}_${range}.log"
    echo "Starting GPU $gpu: samples $range -> $log_file"
    CUDA_VISIBLE_DEVICES=$gpu PYTHONPATH="$PROJECT_ROOT:$PYTHONPATH" \
        python "$PROJECT_ROOT/tests/test_ruler_niah.py" \
        --model "$MODEL" \
        --sample-indices "$range" \
        --enable-offload \
        --num-gpu-blocks 4 \
        --quiet \
        > "$log_file" 2>&1 &
    PIDS+=($!)
    # Small delay to stagger starts
    sleep 2
 done
 echo ""
 echo "All 6 processes started. Waiting for completion..."
 echo "PIDs: ${PIDS[*]}"
 echo ""
 # Wait for all processes and collect results
 declare -a RESULTS=()
 ALL_PASSED=true
 for i in {0..5}; do
    pid="${PIDS[$i]}"
    range="${RANGES[$i]}"
    log_file="$LOG_DIR/gpu${i}_${range}.log"
    if wait $pid; then
        RESULTS+=("GPU $i ($range): PASSED")
        echo "GPU $i completed successfully"
    else
        RESULTS+=("GPU $i ($range): FAILED (exit code $?)")
        ALL_PASSED=false
        echo "GPU $i FAILED!"
    fi
 done
 echo ""
 echo "=========================================="
 echo "RESULTS SUMMARY"
 echo "=========================================="
 for result in "${RESULTS[@]}"; do
    echo "$result"
 done
 echo ""
 # Show accuracy from each log
 echo "Accuracy per GPU:"
 for i in {0..5}; do
    range="${RANGES[$i]}"
    log_file="$LOG_DIR/gpu${i}_${range}.log"
    if [ -f "$log_file" ]; then
        accuracy=$(grep -E "Accuracy:|accuracy" "$log_file" | tail -1 || echo "N/A")
        port=$(grep "Auto-assigned distributed port" "$log_file" | head -1 || echo "N/A")
        echo "  GPU $i ($range): $accuracy | $port"
    fi
 done
 echo ""
 if $ALL_PASSED; then
    echo "=========================================="
    echo "ALL 6 TESTS PASSED!"
    echo "Dynamic port allocation works correctly."
    echo "=========================================="
    exit 0
 else
    echo "=========================================="
    echo "SOME TESTS FAILED!"
    echo "Check logs in $LOG_DIR"
    echo "=========================================="
    exit 1
 fi
--- a/tests/test_minference_gpu.py
+++ b/tests/test_minference_gpu.py
@@ -0,0 +1,163 @@
 """
 Needle-in-haystack test with MInference sparse attention.
 Tests: MInference sparse prefill on GPU-only path (no CPU offload).
 This validates that MInference's vertical + slash sparse pattern can
 correctly retrieve information from long context.
 """
 import os
 os.environ["NANOVLLM_LOG_LEVEL"] = "INFO"
 import argparse
 from nanovllm import LLM, SamplingParams
 from nanovllm.config import SparsePolicyType
 from utils import generate_needle_prompt, check_needle_answer
 def run_minference_test(
    model_path: str,
    max_model_len: int = 16384,
    input_len: int = 8192,
    needle_position: float = 0.5,
    needle_value: str = "7492",
    adaptive_budget: float = 0.3,
    max_new_tokens: int = 32,
    verbose: bool = True,
 ) -> bool:
    """
    Run needle test with MInference sparse prefill attention.
    Args:
        model_path: Path to model
        max_model_len: Maximum model context length
        input_len: Target input sequence length
        needle_position: Where to place needle (0.0-1.0)
        needle_value: The secret value to find
        adaptive_budget: MInference budget as fraction of seq_len
        max_new_tokens: Maximum tokens to generate
        verbose: Print detailed output
    Returns:
        True if test passed, False otherwise
    """
    if verbose:
        print(f"\n{'='*60}")
        print(f"MInference Sparse Prefill Test (GPU-only)")
        print(f"{'='*60}")
        print(f"Model: {model_path}")
        print(f"Max model len: {max_model_len}")
        print(f"Input length: {input_len}")
        print(f"Needle position: {needle_position:.0%}")
        print(f"Needle value: {needle_value}")
        print(f"Adaptive budget: {adaptive_budget}")
        print(f"{'='*60}\n")
    # Initialize LLM with MInference sparse attention
    llm = LLM(
        model_path,
        enforce_eager=True,
        max_model_len=max_model_len,
        max_num_batched_tokens=max_model_len,
        enable_cpu_offload=False,  # GPU-only
        sparse_policy=SparsePolicyType.MINFERENCE,
        minference_adaptive_budget=adaptive_budget,
    )
    # Generate needle prompt
    prompt, expected = generate_needle_prompt(
        tokenizer=llm.tokenizer,
        target_length=input_len,
        needle_position=needle_position,
        needle_value=needle_value,
    )
    # Generate output
    sampling_params = SamplingParams(
        temperature=0.6,
        max_tokens=max_new_tokens,
    )
    outputs = llm.generate([prompt], sampling_params, use_tqdm=True)
    # Check result
    output_text = outputs[0]["text"]
    output_token_ids = outputs[0]["token_ids"]
    passed = check_needle_answer(output_text, expected)
    if verbose:
        print(f"\n{'='*60}")
        print(f"Result")
        print(f"{'='*60}")
        print(f"Expected: {expected}")
        print(f"Output tokens ({len(output_token_ids)}): {output_token_ids[:20]}")
        print(f"Output: {output_text[:200]}...")
        print(f"Status: {'PASSED' if passed else 'FAILED'}")
        print(f"{'='*60}\n")
    return passed
 if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Needle-in-haystack test with MInference sparse prefill"
    )
    parser.add_argument(
        "--model", "-m",
        type=str,
        default=os.path.expanduser("~/models/Qwen3-4B-Instruct-2507/"),
        help="Path to model"
    )
    parser.add_argument(
        "--max-model-len",
        type=int,
        default=16 * 1024,
        help="Maximum model context length"
    )
    parser.add_argument(
        "--input-len",
        type=int,
        default=8 * 1024,
        help="Target input sequence length"
    )
    parser.add_argument(
        "--needle-position",
        type=float,
        default=0.5,
        help="Needle position (0.0=start, 0.5=middle, 1.0=end)"
    )
    parser.add_argument(
        "--needle-value",
        type=str,
        default="7492",
        help="The secret value to hide"
    )
    parser.add_argument(
        "--adaptive-budget",
        type=float,
        default=0.3,
        help="MInference adaptive budget (fraction of seq_len)"
    )
    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=32,
        help="Maximum tokens to generate"
    )
    args = parser.parse_args()
    passed = run_minference_test(
        model_path=args.model,
        max_model_len=args.max_model_len,
        input_len=args.input_len,
        needle_position=args.needle_position,
        needle_value=args.needle_value,
        adaptive_budget=args.adaptive_budget,
        max_new_tokens=args.max_new_tokens,
        verbose=True,
    )
    if passed:
        print("test_minference_gpu: PASSED")
    else:
        print("test_minference_gpu: FAILED")
        exit(1)
--- a/tests/test_needle.py
+++ b/tests/test_needle.py
@@ -31,8 +31,14 @@ def run_needle_test(
    max_new_tokens: int = 32,
    enable_cpu_offload: bool = False,
    enable_quest: bool = False,
    enable_minference: bool = False,
    sparse_topk: int = 8,
    sparse_threshold: int = 4,
    minference_budget: float = 0.3,
    minference_vertical: int = 1000,
    minference_slash: int = 6096,
    gpu_utilization: float = 0.9,
    enforce_eager: bool = True,
    verbose: bool = True,
 ) -> bool:
    """
@@ -49,14 +55,25 @@ def run_needle_test(
        max_new_tokens: Maximum tokens to generate
        enable_cpu_offload: Enable CPU offload mode
        enable_quest: Enable Quest sparse attention (decode-only Top-K)
        enable_minference: Enable MInference sparse prefill (GPU-only)
        sparse_topk: Top-K blocks for Quest
        sparse_threshold: Apply sparse only when blocks > threshold
        minference_budget: MInference adaptive budget (fraction of seq_len, None=fixed mode)
        minference_vertical: Fixed vertical_size (only used when budget=None)
        minference_slash: Fixed slash_size (only used when budget=None)
        gpu_utilization: GPU memory utilization fraction
        verbose: Print detailed output
    Returns:
        True if test passed, False otherwise
    """
-    sparse_policy = SparsePolicyType.QUEST if enable_quest else SparsePolicyType.FULL
+    # Determine sparse policy
    if enable_minference:
        sparse_policy = SparsePolicyType.MINFERENCE
    elif enable_quest:
        sparse_policy = SparsePolicyType.QUEST
    else:
        sparse_policy = SparsePolicyType.FULL
    if verbose:
        print(f"\n{'='*60}")
@@ -69,24 +86,40 @@ def run_needle_test(
        print(f"Needle position: {needle_position:.0%}")
        print(f"Needle value: {needle_value}")
        print(f"CPU offload: {enable_cpu_offload}")
-        if enable_cpu_offload:
+        print(f"Sparse policy: {sparse_policy.name}")
-            print(f"Sparse policy: {sparse_policy.name} (topk={sparse_topk}, threshold={sparse_threshold})")
+        if enable_cpu_offload and enable_quest:
            print(f"  Quest: topk={sparse_topk}, threshold={sparse_threshold}")
        if enable_minference:
            if minference_budget is not None:
                print(f"  MInference: adaptive (budget={minference_budget})")
            else:
                print(f"  MInference: fixed (vertical={minference_vertical}, slash={minference_slash})")
        print(f"{'='*60}\n")
    # 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,
        "kvcache_block_size": block_size,
        "gpu_memory_utilization": gpu_utilization,
    }
    if enable_cpu_offload:
        llm_kwargs["num_gpu_blocks"] = num_gpu_blocks
        llm_kwargs["sparse_policy"] = sparse_policy
        llm_kwargs["sparse_topk_blocks"] = sparse_topk
        llm_kwargs["sparse_threshold_blocks"] = sparse_threshold
    # Set sparse policy (can be used with or without offload)
    if enable_minference or enable_quest:
        llm_kwargs["sparse_policy"] = sparse_policy
    # MInference params (works with both GPU-only and offload mode)
    if enable_minference:
        llm_kwargs["minference_adaptive_budget"] = minference_budget
        llm_kwargs["minference_vertical_size"] = minference_vertical
        llm_kwargs["minference_slash_size"] = minference_slash
    llm = LLM(model_path, **llm_kwargs)
    # 2. Generate needle prompt
@@ -186,6 +219,11 @@ if __name__ == "__main__":
        action="store_true",
        help="Enable Quest sparse attention (decode-only Top-K selection)"
    )
    parser.add_argument(
        "--enable-minference",
        action="store_true",
        help="Enable MInference sparse prefill (GPU-only, vertical+slash pattern)"
    )
    parser.add_argument(
        "--sparse-topk",
        type=int,
@@ -198,8 +236,49 @@ if __name__ == "__main__":
        default=4,
        help="Apply sparse only when blocks > threshold"
    )
    parser.add_argument(
        "--minference-budget",
        type=float,
        default=0.3,
        help="MInference adaptive budget (fraction of seq_len, 0.3=30%% compute, 0=fixed mode)"
    )
    parser.add_argument(
        "--minference-vertical",
        type=int,
        default=1000,
        help="Fixed vertical_size (only used when budget=0)"
    )
    parser.add_argument(
        "--minference-slash",
        type=int,
        default=6096,
        help="Fixed slash_size (only used when budget=0)"
    )
    parser.add_argument(
        "--gpu-utilization",
        type=float,
        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,
@@ -211,8 +290,14 @@ if __name__ == "__main__":
        max_new_tokens=args.max_new_tokens,
        enable_cpu_offload=args.enable_offload,
        enable_quest=args.enable_quest,
        enable_minference=args.enable_minference,
        sparse_topk=args.sparse_topk,
        sparse_threshold=args.sparse_threshold,
        minference_budget=minference_budget,
        minference_vertical=args.minference_vertical,
        minference_slash=args.minference_slash,
        gpu_utilization=args.gpu_utilization,
        enforce_eager=enforce_eager,
        verbose=True,
    )
--- a/tests/test_port_conflict.py
+++ b/tests/test_port_conflict.py
@@ -0,0 +1,198 @@
 """Test for torch distributed port conflict fix.
 This test verifies that:
 1. Multiple independent processes can run simultaneously (dynamic port allocation)
 2. Sequential LLM creation in same process works (proper cleanup)
 Usage:
    # Test parallel processes (requires 2 GPUs)
    python tests/test_port_conflict.py --model ~/models/Qwen3-4B --gpus 4,5 --test parallel
    # Test sequential creation in same process
    CUDA_VISIBLE_DEVICES=4 python tests/test_port_conflict.py --model ~/models/Qwen3-4B --test sequential
 """
 import argparse
 import os
 import subprocess
 import sys
 import time
 def test_sequential_creation(model_path: str, enable_offload: bool = True):
    """Test creating multiple LLM instances sequentially in same process."""
    # Add project root to path
    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    sys.path.insert(0, project_root)
    from nanovllm import LLM, SamplingParams
    print("=" * 60)
    print("Test: Sequential LLM Creation (same process)")
    print("=" * 60)
    for i in range(3):
        print(f"\n--- Creating LLM instance {i+1}/3 ---")
        llm_kwargs = {"enable_cpu_offload": enable_offload}
        if enable_offload:
            llm_kwargs["num_gpu_blocks"] = 2
        llm = LLM(model_path, **llm_kwargs)
        # Simple generation
        outputs = llm.generate(
            ["Hello, how are you?"],
            SamplingParams(max_tokens=20)
        )
        print(f"Output: {outputs[0]['text'][:50]}...")
        # Explicit cleanup
        llm.close()
        print(f"Instance {i+1} closed successfully")
    print("\n" + "=" * 60)
    print("PASSED: test_sequential_creation")
    print("=" * 60)
 def test_context_manager(model_path: str, enable_offload: bool = True):
    """Test LLM with context manager."""
    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    sys.path.insert(0, project_root)
    from nanovllm import LLM, SamplingParams
    print("=" * 60)
    print("Test: Context Manager")
    print("=" * 60)
    for i in range(2):
        print(f"\n--- Context manager instance {i+1}/2 ---")
        llm_kwargs = {"enable_cpu_offload": enable_offload}
        if enable_offload:
            llm_kwargs["num_gpu_blocks"] = 2
        with LLM(model_path, **llm_kwargs) as llm:
            outputs = llm.generate(
                ["What is 2+2?"],
                SamplingParams(max_tokens=20)
            )
            print(f"Output: {outputs[0]['text'][:50]}...")
        print(f"Instance {i+1} auto-closed via context manager")
    print("\n" + "=" * 60)
    print("PASSED: test_context_manager")
    print("=" * 60)
 def test_parallel_processes(model_path: str, gpus: str, enable_offload: bool = True):
    """Test running multiple nanovllm processes in parallel."""
    gpu_list = [int(g.strip()) for g in gpus.split(",")]
    if len(gpu_list) < 2:
        print("ERROR: Need at least 2 GPUs for parallel test")
        return False
    print("=" * 60)
    print(f"Test: Parallel Processes (GPUs: {gpu_list})")
    print("=" * 60)
    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    # Script to run in each subprocess
    script = f'''
 import sys
 sys.path.insert(0, "{project_root}")
 import os
 from nanovllm import LLM, SamplingParams
 gpu = os.environ.get("CUDA_VISIBLE_DEVICES", "?")
 print(f"[GPU {{gpu}}] Starting LLM...")
 llm_kwargs = {{"enable_cpu_offload": {enable_offload}}}
 if {enable_offload}:
    llm_kwargs["num_gpu_blocks"] = 2
 llm = LLM("{model_path}", **llm_kwargs)
 print(f"[GPU {{gpu}}] LLM initialized, generating...")
 outputs = llm.generate(["Hello world"], SamplingParams(max_tokens=10))
 print(f"[GPU {{gpu}}] Output: {{outputs[0]['text'][:30]}}...")
 llm.close()
 print(f"[GPU {{gpu}}] Done")
 '''
    # Start processes on different GPUs
    procs = []
    for i, gpu in enumerate(gpu_list[:2]):  # Use first 2 GPUs
        print(f"\nStarting process on GPU {gpu}...")
        env = os.environ.copy()
        env["CUDA_VISIBLE_DEVICES"] = str(gpu)
        p = subprocess.Popen(
            [sys.executable, "-c", script],
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True
        )
        procs.append((gpu, p))
        time.sleep(2)  # Stagger starts to see concurrent running
    # Wait and collect results
    all_passed = True
    for gpu, p in procs:
        stdout, _ = p.communicate(timeout=300)
        print(f"\n--- GPU {gpu} output ---")
        print(stdout)
        if p.returncode != 0:
            print(f"ERROR: GPU {gpu} process failed with code {p.returncode}")
            all_passed = False
        else:
            print(f"GPU {gpu} process completed successfully")
    print("\n" + "=" * 60)
    if all_passed:
        print("PASSED: test_parallel_processes")
    else:
        print("FAILED: test_parallel_processes")
    print("=" * 60)
    return all_passed
 def main():
    parser = argparse.ArgumentParser(description="Test port conflict fix")
    parser.add_argument("--model", "-m", required=True, help="Path to model")
    parser.add_argument("--gpus", default="0,1", help="GPUs to use for parallel test (comma-separated)")
    parser.add_argument("--test", choices=["sequential", "context", "parallel", "all"],
                        default="all", help="Which test to run")
    parser.add_argument("--no-offload", action="store_true", help="Disable CPU offload")
    args = parser.parse_args()
    enable_offload = not args.no_offload
    model_path = os.path.expanduser(args.model)
    print(f"Model: {model_path}")
    print(f"CPU Offload: {enable_offload}")
    print(f"GPUs for parallel test: {args.gpus}")
    print()
    if args.test in ["sequential", "all"]:
        test_sequential_creation(model_path, enable_offload)
        print()
    if args.test in ["context", "all"]:
        test_context_manager(model_path, enable_offload)
        print()
    if args.test in ["parallel", "all"]:
        test_parallel_processes(model_path, args.gpus, enable_offload)
 if __name__ == "__main__":
    main()
--- a/tests/test_ruler.py
+++ b/tests/test_ruler.py
@@ -0,0 +1,392 @@
 """
 RULER benchmark comprehensive test for LLM.
 Tests multiple RULER tasks:
 - NIAH (Needle-In-A-Haystack): single, multikey, multiquery, multivalue
 - QA (Question Answering): qa_1, qa_2
 - CWE (Common Word Extraction)
 - FWE (Frequent Word Extraction)
 - VT (Variable Tracking)
 Usage:
    # Test all datasets with 2 samples each (debug mode)
    python tests/test_ruler.py --enable-offload --num-samples 2
    # Test specific datasets
    python tests/test_ruler.py --enable-offload --datasets niah_single_1,qa_1
    # Test all samples in all datasets
    python tests/test_ruler.py --enable-offload
 """
 import os
 os.environ["NANOVLLM_LOG_LEVEL"] = "INFO"
 import argparse
 import json
 import re
 import gc
 import time
 import torch
 from pathlib import Path
 from typing import List, Dict, Tuple, Optional
 from nanovllm import LLM, SamplingParams
 # ============================================================
 # Constants
 # ============================================================
 DEFAULT_DATA_DIR = Path(__file__).parent / "data/ruler_32k"
 DEFAULT_MODEL = os.path.expanduser("~/models/Llama-3.1-8B-Instruct")
 # Note: max_model_len must be > max_input_len to leave room for output tokens
 # 32k benchmark has inputs up to 32760 tokens, so we need 32768 + 128 = 32896
 DEFAULT_MAX_MODEL_LEN = 32896
 DEFAULT_MAX_NEW_TOKENS = 128  # Larger for multi-value tasks
 # Task categories for evaluation
 NIAH_TASKS = ["niah_single_1", "niah_single_2", "niah_single_3",
              "niah_multikey_1", "niah_multikey_2", "niah_multikey_3",
              "niah_multiquery", "niah_multivalue"]
 QA_TASKS = ["qa_1", "qa_2"]
 RECALL_TASKS = ["cwe", "fwe", "vt"]
 ALL_TASKS = NIAH_TASKS + QA_TASKS + RECALL_TASKS
 # ============================================================
 # Data Loading
 # ============================================================
 def load_samples(filepath: Path, indices: Optional[List[int]] = None) -> List[dict]:
    """Load samples from a JSONL file."""
    if not filepath.exists():
        raise FileNotFoundError(f"Data file not found: {filepath}")
    samples = []
    with open(filepath) as f:
        for i, line in enumerate(f):
            if indices is None or i in indices:
                sample = json.loads(line)
                sample["_local_idx"] = i
                samples.append(sample)
    return samples
 def count_samples(filepath: Path) -> int:
    """Count total samples in JSONL file."""
    with open(filepath) as f:
        return sum(1 for _ in f)
 # ============================================================
 # Evaluation Functions (Following RULER Official Metrics)
 # Ref: https://github.com/NVIDIA/RULER/blob/main/scripts/eval/synthetic/constants.py
 # ============================================================
 def string_match_all(output_text: str, expected_list: List[str]) -> float:
    """
    RULER official metric for NIAH, VT, CWE, FWE tasks.
    Formula: sum([1.0 if r.lower() in pred.lower() else 0.0 for r in ref]) / len(ref)
    Returns recall score (0.0 to 1.0): fraction of expected values found in output.
    """
    output_clean = output_text.replace('<|im_end|>', '').replace('\r', ' ').replace('\n', ' ')
    output_lower = output_clean.lower()
    if not expected_list:
        return 1.0
    found = sum(1.0 if exp.strip().lower() in output_lower else 0.0 for exp in expected_list)
    return found / len(expected_list)
 def string_match_part(output_text: str, expected_list: List[str]) -> float:
    """
    RULER official metric for QA tasks.
    Formula: max([1.0 if r.lower() in pred.lower() else 0.0 for r in ref])
    Returns 1.0 if ANY expected value is found, 0.0 otherwise.
    """
    output_clean = output_text.replace('<|im_end|>', '').replace('\r', ' ').replace('\n', ' ')
    output_lower = output_clean.lower()
    if not expected_list:
        return 1.0
    return max(1.0 if exp.strip().lower() in output_lower else 0.0 for exp in expected_list)
 def evaluate_output(output_text: str, expected_outputs: List[str], task_name: str) -> Tuple[bool, float]:
    """
    Evaluate model output using RULER official metrics.
    - QA tasks: string_match_part (any match = full score)
    - All other tasks: string_match_all (recall-based score)
    Returns (passed, score) where passed = score >= 0.5
    """
    if task_name in QA_TASKS:
        score = string_match_part(output_text, expected_outputs)
    else:
        # NIAH, VT, CWE, FWE all use string_match_all
        score = string_match_all(output_text, expected_outputs)
    passed = score >= 0.5  # Consider pass if score >= 50%
    return passed, score
 # ============================================================
 # Test Runner
 # ============================================================
 def run_task_test(
    llm: LLM,
    task_name: str,
    data_dir: Path,
    sample_indices: Optional[List[int]] = None,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    verbose: bool = True,
 ) -> Dict:
    """
    Run test for a single RULER task.
    Returns dict with: task, correct, total, score, results
    """
    data_file = data_dir / task_name / "validation.jsonl"
    samples = load_samples(data_file, sample_indices)
    if verbose:
        print(f"\n  Testing {task_name}: {len(samples)} samples")
    sampling_params = SamplingParams(
        temperature=0.1,
        max_tokens=max_new_tokens,
    )
    correct = 0
    total_score = 0.0
    results = []
    for sample in samples:
        idx = sample.get("index", sample["_local_idx"])
        prompt = sample["input"]
        expected = sample["outputs"]
        # Generate
        outputs = llm.generate([prompt], sampling_params, use_tqdm=False)
        output_text = outputs[0]["text"]
        # Evaluate
        passed, score = evaluate_output(output_text, expected, task_name)
        if passed:
            correct += 1
        total_score += score
        results.append({
            "index": idx,
            "expected": expected,
            "output": output_text[:200],
            "passed": passed,
            "score": score,
        })
        if verbose:
            status = "PASS" if passed else "FAIL"
            exp_preview = str(expected[0])[:30] if expected else "N/A"
            out_preview = output_text[:50].replace('\n', ' ')
            print(f"    [{idx}] {status} (score={score:.2f}) exp={exp_preview}... out={out_preview}...")
    avg_score = total_score / len(samples) if samples else 0.0
    return {
        "task": task_name,
        "correct": correct,
        "total": len(samples),
        "accuracy": correct / len(samples) if samples else 0.0,
        "avg_score": avg_score,
        "results": results,
    }
 def run_ruler_benchmark(
    model_path: str,
    data_dir: Path,
    datasets: Optional[List[str]] = None,
    num_samples: Optional[int] = None,
    max_model_len: int = DEFAULT_MAX_MODEL_LEN,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    enable_cpu_offload: bool = False,
    num_gpu_blocks: int = 4,
    block_size: int = 1024,
    gpu_utilization: float = 0.9,
    enforce_eager: bool = True,
    verbose: bool = True,
 ) -> Dict:
    """
    Run RULER benchmark on multiple tasks.
    Args:
        model_path: Path to the model
        data_dir: Directory containing task subdirectories
        datasets: List of task names to test (None = all)
        num_samples: Number of samples per task (None = all)
        ...other LLM config params...
    Returns:
        Dict with overall results and per-task results
    """
    # Determine tasks to run
    if datasets is None:
        tasks = [t for t in ALL_TASKS if (data_dir / t / "validation.jsonl").exists()]
    else:
        tasks = datasets
    # Sample indices
    sample_indices = list(range(num_samples)) if num_samples else None
    print(f"\n{'='*60}")
    print(f"RULER Benchmark")
    print(f"{'='*60}")
    print(f"Model: {model_path}")
    print(f"Data dir: {data_dir}")
    print(f"Tasks: {len(tasks)}")
    print(f"Samples per task: {num_samples if num_samples else 'all'}")
    print(f"CPU offload: {enable_cpu_offload}")
    print(f"{'='*60}")
    # Initialize LLM
    print("\nInitializing LLM...")
    llm_kwargs = {
        "max_model_len": max_model_len,
        "max_num_batched_tokens": max_model_len,
        "enforce_eager": enforce_eager,
        "gpu_memory_utilization": gpu_utilization,
        "kvcache_block_size": block_size,
        "enable_cpu_offload": enable_cpu_offload,
    }
    if enable_cpu_offload:
        llm_kwargs["num_gpu_blocks"] = num_gpu_blocks
    llm = LLM(model_path, **llm_kwargs)
    # Run tests
    start_time = time.time()
    task_results = []
    for task_name in tasks:
        result = run_task_test(
            llm=llm,
            task_name=task_name,
            data_dir=data_dir,
            sample_indices=sample_indices,
            max_new_tokens=max_new_tokens,
            verbose=verbose,
        )
        task_results.append(result)
        if verbose:
            print(f"  -> {task_name}: {result['correct']}/{result['total']} "
                  f"({result['accuracy']*100:.1f}%) avg_score={result['avg_score']:.3f}")
    total_time = time.time() - start_time
    # Cleanup
    del llm
    gc.collect()
    torch.cuda.empty_cache()
    # Aggregate results
    total_correct = sum(r["correct"] for r in task_results)
    total_samples = sum(r["total"] for r in task_results)
    overall_accuracy = total_correct / total_samples if total_samples > 0 else 0.0
    avg_score = sum(r["avg_score"] for r in task_results) / len(task_results) if task_results else 0.0
    # Print summary
    print(f"\n{'='*60}")
    print(f"RULER Benchmark Results")
    print(f"{'='*60}")
    print(f"\n{'Task':<20} {'Correct':<10} {'Accuracy':<12} {'Avg Score':<12}")
    print(f"{'-'*54}")
    for r in task_results:
        print(f"{r['task']:<20} {r['correct']}/{r['total']:<7} {r['accuracy']*100:>6.1f}%      {r['avg_score']:.3f}")
    print(f"{'-'*54}")
    print(f"{'TOTAL':<20} {total_correct}/{total_samples:<7} {overall_accuracy*100:>6.1f}%      {avg_score:.3f}")
    print(f"\nTime: {total_time:.1f}s")
    print(f"{'='*60}\n")
    return {
        "total_correct": total_correct,
        "total_samples": total_samples,
        "overall_accuracy": overall_accuracy,
        "avg_score": avg_score,
        "time": total_time,
        "task_results": task_results,
    }
 # ============================================================
 # CLI Entry Point
 # ============================================================
 if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="RULER benchmark comprehensive test",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--model", "-m", type=str, default=DEFAULT_MODEL,
                        help=f"Path to model (default: {DEFAULT_MODEL})")
    parser.add_argument("--data-dir", type=str, default=str(DEFAULT_DATA_DIR),
                        help=f"Path to data directory (default: {DEFAULT_DATA_DIR})")
    parser.add_argument("--datasets", type=str, default="",
                        help="Comma-separated list of datasets to test (default: all)")
    parser.add_argument("--num-samples", type=int, default=0,
                        help="Number of samples per dataset (default: 0 = all)")
    parser.add_argument("--max-model-len", type=int, default=DEFAULT_MAX_MODEL_LEN,
                        help=f"Maximum model context length (default: {DEFAULT_MAX_MODEL_LEN})")
    parser.add_argument("--max-new-tokens", type=int, default=DEFAULT_MAX_NEW_TOKENS,
                        help=f"Maximum tokens to generate (default: {DEFAULT_MAX_NEW_TOKENS})")
    parser.add_argument("--enable-offload", action="store_true",
                        help="Enable CPU offload mode")
    parser.add_argument("--num-gpu-blocks", type=int, default=4,
                        help="Number of GPU blocks for CPU offload (default: 4)")
    parser.add_argument("--block-size", type=int, default=1024,
                        help="KV cache block size (default: 1024)")
    parser.add_argument("--gpu-utilization", type=float, default=0.9,
                        help="GPU memory utilization (default: 0.9)")
    parser.add_argument("--use-cuda-graph", action="store_true",
                        help="Enable CUDA graph")
    parser.add_argument("--quiet", "-q", action="store_true",
                        help="Quiet mode")
    args = parser.parse_args()
    # Parse datasets
    datasets = args.datasets.split(",") if args.datasets else None
    num_samples = args.num_samples if args.num_samples > 0 else None
    results = run_ruler_benchmark(
        model_path=os.path.expanduser(args.model),
        data_dir=Path(args.data_dir),
        datasets=datasets,
        num_samples=num_samples,
        max_model_len=args.max_model_len,
        max_new_tokens=args.max_new_tokens,
        enable_cpu_offload=args.enable_offload,
        num_gpu_blocks=args.num_gpu_blocks,
        block_size=args.block_size,
        gpu_utilization=args.gpu_utilization,
        enforce_eager=not args.use_cuda_graph,
        verbose=not args.quiet,
    )
    # Exit code
    if results["overall_accuracy"] >= 0.5:
        print("test_ruler: PASSED")
    else:
        print(f"test_ruler: FAILED (accuracy={results['overall_accuracy']*100:.1f}%)")
        exit(1)
--- a/tests/test_ruler_niah.py
+++ b/tests/test_ruler_niah.py
@@ -0,0 +1,527 @@
 """
 RULER NIAH benchmark test for LLM.
 Tests: Long context retrieval capability using pre-generated RULER benchmark data.
 The NIAH (Needle-In-A-Haystack) task tests the model's ability to retrieve a
 specific magic number from a large context (~32K tokens).
 Usage:
    # Test all samples with CPU offload
    python tests/test_ruler_niah.py --enable-offload
    # Test specific samples
    python tests/test_ruler_niah.py --sample-indices 0,1,2 --enable-offload
    # Test with custom model
    python tests/test_ruler_niah.py --model /path/to/model --enable-offload
    # Group mode: test in batches with separate LLM initialization per group
    python tests/test_ruler_niah.py --enable-offload --group-size 5
 """
 import os
 os.environ["NANOVLLM_LOG_LEVEL"] = "INFO"
 import argparse
 import json
 from pathlib import Path
 from typing import List, Tuple, Optional
 from nanovllm import LLM, SamplingParams
 from utils import check_needle_answer
 # ============================================================
 # Constants
 # ============================================================
 DEFAULT_DATA_FILE = Path(__file__).parent / "data/ruler_niah/niah_single_1_32k.jsonl"
 DEFAULT_MODEL = os.path.expanduser("~/models/Llama-3.1-8B-Instruct")
 DEFAULT_MAX_MODEL_LEN = 32768
 DEFAULT_MAX_NEW_TOKENS = 50
 # ============================================================
 # Data Loading
 # ============================================================
 def load_ruler_samples(filepath: Path, indices: Optional[List[int]] = None) -> List[dict]:
    """
    Load RULER NIAH samples from a JSONL file.
    Args:
        filepath: Path to the JSONL file
        indices: Optional list of sample indices to load. If None, load all.
    Returns:
        List of sample dicts with keys: index, input, outputs, length
    """
    if not filepath.exists():
        raise FileNotFoundError(
            f"Data file not found: {filepath}\n"
            f"Please copy RULER NIAH data to this location. See docs/ruler_niah_standalone_test.md"
        )
    samples = []
    with open(filepath) as f:
        for i, line in enumerate(f):
            if indices is None or i in indices:
                sample = json.loads(line)
                samples.append(sample)
    if not samples:
        raise ValueError(f"No samples loaded from {filepath}")
    return samples
 def count_samples(filepath: Path) -> int:
    """Count total samples in JSONL file."""
    with open(filepath) as f:
        return sum(1 for _ in f)
 # ============================================================
 # Test Function
 # ============================================================
 def run_ruler_niah_test(
    model_path: str,
    data_file: Path,
    sample_indices: Optional[List[int]] = None,
    max_model_len: int = DEFAULT_MAX_MODEL_LEN,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    enable_cpu_offload: bool = False,
    num_gpu_blocks: int = 4,
    block_size: int = 1024,
    gpu_utilization: float = 0.9,
    enforce_eager: bool = True,
    verbose: bool = True,
 ) -> Tuple[int, int]:
    """
    Run RULER NIAH test on loaded samples.
    Args:
        model_path: Path to the model
        data_file: Path to JSONL data file
        sample_indices: List of sample indices to test (None = all)
        max_model_len: Maximum model context length
        max_new_tokens: Maximum tokens to generate
        enable_cpu_offload: Enable CPU offload mode
        num_gpu_blocks: Number of GPU blocks for offload
        block_size: KV cache block size
        gpu_utilization: GPU memory utilization fraction
        enforce_eager: Disable CUDA graphs
        verbose: Print detailed output
    Returns:
        (correct, total): Number of correct and total samples
    """
    # Load samples
    samples = load_ruler_samples(data_file, sample_indices)
    total = len(samples)
    if verbose:
        print(f"\n{'='*60}")
        print(f"RULER NIAH Test")
        print(f"{'='*60}")
        print(f"Model: {model_path}")
        print(f"Data file: {data_file}")
        print(f"Samples: {total}")
        print(f"Max model len: {max_model_len}")
        print(f"Max new tokens: {max_new_tokens}")
        print(f"CPU offload: {enable_cpu_offload}")
        if enable_cpu_offload:
            print(f"  num_gpu_blocks: {num_gpu_blocks}")
            print(f"  block_size: {block_size}")
        print(f"Enforce eager: {enforce_eager}")
        print(f"{'='*60}\n")
    # Check max_model_len vs data length
    max_data_len = max(s.get("length", 0) for s in samples)
    if max_model_len < max_data_len:
        print(f"WARNING: max_model_len ({max_model_len}) < max data length ({max_data_len})")
        print(f"         This may cause truncation or errors.\n")
    # Initialize LLM
    if verbose:
        print("Initializing LLM...")
    llm_kwargs = {
        "max_model_len": max_model_len,
        "max_num_batched_tokens": max_model_len,
        "enforce_eager": enforce_eager,
        "gpu_memory_utilization": gpu_utilization,
        "kvcache_block_size": block_size,
        "enable_cpu_offload": enable_cpu_offload,
    }
    if enable_cpu_offload:
        llm_kwargs["num_gpu_blocks"] = num_gpu_blocks
    llm = LLM(model_path, **llm_kwargs)
    # Sampling params
    # Note: nano-vllm doesn't support greedy (temperature=0), use low temperature instead
    sampling_params = SamplingParams(
        temperature=0.1,  # Low temperature for near-deterministic output
        max_tokens=max_new_tokens,
    )
    # Test each sample
    correct = 0
    results = []
    for i, sample in enumerate(samples):
        sample_idx = sample.get("index", i)
        prompt = sample["input"]
        expected = sample["outputs"][0]
        data_len = sample.get("length", "unknown")
        if verbose:
            print(f"\nSample {sample_idx}: Expected={expected}, Length={data_len}")
        # Generate
        outputs = llm.generate([prompt], sampling_params, use_tqdm=False)
        output_text = outputs[0]["text"]
        output_tokens = outputs[0]["token_ids"]
        # Check result
        passed = check_needle_answer(output_text, expected)
        if passed:
            correct += 1
        results.append({
            "index": sample_idx,
            "expected": expected,
            "output": output_text,
            "passed": passed,
        })
        if verbose:
            status = "PASS" if passed else "FAIL"
            output_preview = output_text[:100].replace('\n', ' ')
            print(f"  Output ({len(output_tokens)} tokens): {output_preview}...")
            print(f"  Status: {status}")
    # Summary
    if verbose:
        print(f"\n{'='*60}")
        print(f"Results: {correct}/{total} PASSED ({100*correct/total:.1f}%)")
        print(f"{'='*60}\n")
        if correct < total:
            print("Failed samples:")
            for r in results:
                if not r["passed"]:
                    print(f"  Sample {r['index']}: expected={r['expected']}, got={r['output'][:50]}...")
    return correct, total
 # ============================================================
 # Grouped Test Function
 # ============================================================
 def run_grouped_test(
    model_path: str,
    data_file: Path,
    group_size: int = 5,
    total_samples: Optional[int] = None,
    max_model_len: int = DEFAULT_MAX_MODEL_LEN,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    enable_cpu_offload: bool = False,
    num_gpu_blocks: int = 4,
    block_size: int = 1024,
    gpu_utilization: float = 0.9,
    enforce_eager: bool = True,
 ) -> Tuple[int, int, List[dict]]:
    """
    Run RULER NIAH test in groups, with separate LLM initialization per group.
    This mode is useful for:
    - Avoiding state accumulation issues
    - Testing LLM initialization stability
    - Running large-scale tests with memory cleanup between groups
    Args:
        model_path: Path to the model
        data_file: Path to JSONL data file
        group_size: Number of samples per group
        total_samples: Total samples to test (None = all in file)
        Other args: Same as run_ruler_niah_test
    Returns:
        (total_correct, total_tested, group_results): Results summary
    """
    import time
    import gc
    import torch
    # Count total samples in file
    file_sample_count = count_samples(data_file)
    if total_samples is None:
        total_samples = file_sample_count
    else:
        total_samples = min(total_samples, file_sample_count)
    num_groups = (total_samples + group_size - 1) // group_size
    print(f"\n{'='*60}")
    print(f"RULER NIAH Grouped Test")
    print(f"{'='*60}")
    print(f"Model: {model_path}")
    print(f"Data file: {data_file}")
    print(f"Total samples: {total_samples}")
    print(f"Group size: {group_size}")
    print(f"Number of groups: {num_groups}")
    print(f"CPU offload: {enable_cpu_offload}")
    print(f"{'='*60}\n")
    total_correct = 0
    total_tested = 0
    group_results = []
    all_failed = []
    test_start_time = time.time()
    for group_idx in range(num_groups):
        start_idx = group_idx * group_size
        end_idx = min(start_idx + group_size, total_samples)
        sample_indices = list(range(start_idx, end_idx))
        print(f"\n{'='*60}")
        print(f"Group {group_idx + 1}/{num_groups}: Samples {start_idx}-{end_idx - 1}")
        print(f"{'='*60}")
        group_start_time = time.time()
        # Run test for this group
        correct, tested = run_ruler_niah_test(
            model_path=model_path,
            data_file=data_file,
            sample_indices=sample_indices,
            max_model_len=max_model_len,
            max_new_tokens=max_new_tokens,
            enable_cpu_offload=enable_cpu_offload,
            num_gpu_blocks=num_gpu_blocks,
            block_size=block_size,
            gpu_utilization=gpu_utilization,
            enforce_eager=enforce_eager,
            verbose=True,
        )
        group_time = time.time() - group_start_time
        total_correct += correct
        total_tested += tested
        group_result = {
            "group": group_idx + 1,
            "samples": f"{start_idx}-{end_idx - 1}",
            "correct": correct,
            "total": tested,
            "accuracy": 100 * correct / tested if tested > 0 else 0,
            "time": group_time,
        }
        group_results.append(group_result)
        print(f"\nGroup {group_idx + 1} Summary: {correct}/{tested} PASSED ({group_result['accuracy']:.1f}%) in {group_time:.1f}s")
        # Force cleanup between groups
        gc.collect()
        torch.cuda.empty_cache()
        # Small delay to ensure port is released
        if group_idx < num_groups - 1:
            time.sleep(3)
    total_time = time.time() - test_start_time
    # Final summary
    print(f"\n{'='*60}")
    print(f"FINAL SUMMARY")
    print(f"{'='*60}")
    print(f"\nGroup Results:")
    print(f"{'Group':<8} {'Samples':<12} {'Result':<12} {'Accuracy':<10} {'Time':<10}")
    print(f"{'-'*52}")
    for r in group_results:
        print(f"{r['group']:<8} {r['samples']:<12} {r['correct']}/{r['total']:<9} {r['accuracy']:.1f}%{'':<5} {r['time']:.1f}s")
    print(f"{'-'*52}")
    overall_accuracy = 100 * total_correct / total_tested if total_tested > 0 else 0
    print(f"{'TOTAL':<8} {'0-' + str(total_tested-1):<12} {total_correct}/{total_tested:<9} {overall_accuracy:.1f}%{'':<5} {total_time:.1f}s")
    print(f"{'='*60}\n")
    return total_correct, total_tested, group_results
 # ============================================================
 # CLI Entry Point
 # ============================================================
 def parse_indices(s: str) -> List[int]:
    """Parse comma-separated indices like '0,1,2' or range like '0-4'."""
    if not s:
        return None
    indices = []
    for part in s.split(','):
        if '-' in part:
            start, end = part.split('-')
            indices.extend(range(int(start), int(end) + 1))
        else:
            indices.append(int(part))
    return indices
 if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="RULER NIAH benchmark test for long context LLM",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
 Examples:
  # Test all samples with CPU offload (recommended for 24GB GPUs)
  python tests/test_ruler_niah.py --enable-offload
  # Test specific samples
  python tests/test_ruler_niah.py --sample-indices 0,1,2 --enable-offload
  # Test with CUDA graph enabled
  python tests/test_ruler_niah.py --enable-offload --use-cuda-graph
        """
    )
    parser.add_argument(
        "--model", "-m",
        type=str,
        default=DEFAULT_MODEL,
        help=f"Path to model (default: {DEFAULT_MODEL})"
    )
    parser.add_argument(
        "--data-file",
        type=str,
        default=str(DEFAULT_DATA_FILE),
        help=f"Path to JSONL data file (default: {DEFAULT_DATA_FILE})"
    )
    parser.add_argument(
        "--sample-indices",
        type=str,
        default="",
        help="Sample indices to test (e.g., '0,1,2' or '0-4'). Default: all"
    )
    parser.add_argument(
        "--max-model-len",
        type=int,
        default=DEFAULT_MAX_MODEL_LEN,
        help=f"Maximum model context length (default: {DEFAULT_MAX_MODEL_LEN})"
    )
    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=DEFAULT_MAX_NEW_TOKENS,
        help=f"Maximum tokens to generate (default: {DEFAULT_MAX_NEW_TOKENS})"
    )
    parser.add_argument(
        "--enable-offload",
        action="store_true",
        help="Enable CPU offload mode (required for 24GB GPUs with 32K context)"
    )
    parser.add_argument(
        "--num-gpu-blocks",
        type=int,
        default=4,
        help="Number of GPU blocks for CPU offload (default: 4)"
    )
    parser.add_argument(
        "--block-size",
        type=int,
        default=1024,
        help="KV cache block size (default: 1024)"
    )
    parser.add_argument(
        "--gpu-utilization",
        type=float,
        default=0.9,
        help="GPU memory utilization fraction (default: 0.9)"
    )
    parser.add_argument(
        "--enforce-eager",
        action="store_true",
        default=True,
        help="Force eager execution, disable CUDA graphs (default: True)"
    )
    parser.add_argument(
        "--use-cuda-graph",
        action="store_true",
        help="Enable CUDA graph (overrides --enforce-eager)"
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        default=True,
        help="Print detailed output (default: True)"
    )
    parser.add_argument(
        "--quiet", "-q",
        action="store_true",
        help="Quiet mode, only print final result"
    )
    parser.add_argument(
        "--group-size",
        type=int,
        default=0,
        help="Enable grouped testing mode with specified group size. Each group initializes LLM separately. (default: 0 = disabled)"
    )
    parser.add_argument(
        "--total-samples",
        type=int,
        default=0,
        help="Total number of samples to test in group mode (default: 0 = all samples in file)"
    )
    args = parser.parse_args()
    # Process arguments
    sample_indices = parse_indices(args.sample_indices)
    enforce_eager = not args.use_cuda_graph
    verbose = not args.quiet
    # Check if group mode is enabled
    if args.group_size > 0:
        # Grouped testing mode
        total_samples = args.total_samples if args.total_samples > 0 else None
        correct, total, _ = run_grouped_test(
            model_path=os.path.expanduser(args.model),
            data_file=Path(args.data_file),
            group_size=args.group_size,
            total_samples=total_samples,
            max_model_len=args.max_model_len,
            max_new_tokens=args.max_new_tokens,
            enable_cpu_offload=args.enable_offload,
            num_gpu_blocks=args.num_gpu_blocks,
            block_size=args.block_size,
            gpu_utilization=args.gpu_utilization,
            enforce_eager=enforce_eager,
        )
    else:
        # Standard testing mode
        correct, total = run_ruler_niah_test(
            model_path=os.path.expanduser(args.model),
            data_file=Path(args.data_file),
            sample_indices=sample_indices,
            max_model_len=args.max_model_len,
            max_new_tokens=args.max_new_tokens,
            enable_cpu_offload=args.enable_offload,
            num_gpu_blocks=args.num_gpu_blocks,
            block_size=args.block_size,
            gpu_utilization=args.gpu_utilization,
            enforce_eager=enforce_eager,
            verbose=verbose,
        )
    # Final status
    if correct == total:
        print("test_ruler_niah: PASSED")
    else:
        print(f"test_ruler_niah: FAILED ({correct}/{total})")
        exit(1)
--- a/tests/test_ruler_niah.sh
+++ b/tests/test_ruler_niah.sh
@@ -0,0 +1,242 @@
 #!/bin/bash
 #
 # RULER NIAH Parallel Test Script
 #
 # Runs RULER NIAH benchmark across multiple GPUs in parallel.
 # Each sample is tested independently (separate Python process per sample).
 #
 # Usage:
 #   ./tests/test_ruler_niah.sh [OPTIONS]
 #
 # Options:
 #   --gpus "0,1,2,3"     GPUs to use (default: "0,1,2,3")
 #   --total N            Total samples to test (default: 100)
 #   --model PATH         Model path (default: ~/models/Llama-3.1-8B-Instruct)
 #   --output FILE        Output log file (default: /tmp/ruler_niah_results.log)
 #
 # Note: Removed 'set -e' because ((var++)) returns 1 when var=0, which triggers exit
 # Default configuration
 GPUS="0,1,2,3"
 TOTAL_SAMPLES=100
 MODEL_PATH="$HOME/models/Llama-3.1-8B-Instruct"
 OUTPUT_LOG="/tmp/ruler_niah_results.log"
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
 # Parse arguments
 while [[ $# -gt 0 ]]; do
    case $1 in
        --gpus)
            GPUS="$2"
            shift 2
            ;;
        --total)
            TOTAL_SAMPLES="$2"
            shift 2
            ;;
        --model)
            MODEL_PATH="$2"
            shift 2
            ;;
        --output)
            OUTPUT_LOG="$2"
            shift 2
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
 done
 # Convert GPU string to array
 IFS=',' read -ra GPU_ARRAY <<< "$GPUS"
 NUM_GPUS=${#GPU_ARRAY[@]}
 echo "============================================================"
 echo "RULER NIAH Parallel Test"
 echo "============================================================"
 echo "GPUs: ${GPUS} (${NUM_GPUS} GPUs)"
 echo "Total samples: ${TOTAL_SAMPLES}"
 echo "Model: ${MODEL_PATH}"
 echo "Output log: ${OUTPUT_LOG}"
 echo "Project root: ${PROJECT_ROOT}"
 echo "============================================================"
 echo ""
 # Create output directory
 mkdir -p "$(dirname "$OUTPUT_LOG")"
 # Initialize result tracking
 RESULT_DIR="/tmp/ruler_niah_results_$$"
 mkdir -p "$RESULT_DIR"
 # Function to run a single sample on a specific GPU
 run_sample() {
    local gpu=$1
    local sample_idx=$2
    local result_file="$RESULT_DIR/sample_${sample_idx}.result"
    # Run test with unique port based on GPU
    local port=$((2333 + gpu))
    NANOVLLM_DIST_PORT=$port \
    CUDA_VISIBLE_DEVICES=$gpu \
    PYTHONPATH="$PROJECT_ROOT:$PYTHONPATH" \
    python "$SCRIPT_DIR/test_ruler_niah.py" \
        --model "$MODEL_PATH" \
        --enable-offload \
        --sample-indices "$sample_idx" \
        --quiet \
        2>&1
    local exit_code=$?
    if [ $exit_code -eq 0 ]; then
        echo "PASS" > "$result_file"
    else
        echo "FAIL" > "$result_file"
    fi
    return $exit_code
 }
 # Function to run samples on a specific GPU
 run_gpu_worker() {
    local gpu=$1
    local gpu_idx=$2
    local log_file="$RESULT_DIR/gpu_${gpu}.log"
    echo "[GPU $gpu] Starting worker (gpu_idx=$gpu_idx)" | tee -a "$log_file"
    # Calculate which samples this GPU handles
    local sample_idx=$gpu_idx
    local pass_count=0
    local fail_count=0
    while [ $sample_idx -lt $TOTAL_SAMPLES ]; do
        echo "[GPU $gpu] Testing sample $sample_idx..." | tee -a "$log_file"
        local start_time=$(date +%s)
        if run_sample $gpu $sample_idx >> "$log_file" 2>&1; then
            echo "[GPU $gpu] Sample $sample_idx: PASS" | tee -a "$log_file"
            ((pass_count++))
        else
            echo "[GPU $gpu] Sample $sample_idx: FAIL" | tee -a "$log_file"
            ((fail_count++))
        fi
        local end_time=$(date +%s)
        local duration=$((end_time - start_time))
        echo "[GPU $gpu] Sample $sample_idx completed in ${duration}s" | tee -a "$log_file"
        # Move to next sample for this GPU (stride by number of GPUs)
        sample_idx=$((sample_idx + NUM_GPUS))
        # Small delay to avoid port conflicts
        sleep 2
    done
    echo "[GPU $gpu] Worker finished: $pass_count passed, $fail_count failed" | tee -a "$log_file"
    echo "$pass_count $fail_count" > "$RESULT_DIR/gpu_${gpu}.summary"
 }
 # Start time
 START_TIME=$(date +%s)
 echo "Starting parallel test at $(date '+%Y-%m-%d %H:%M:%S')"
 echo ""
 # Launch workers for each GPU in background
 PIDS=()
 for i in "${!GPU_ARRAY[@]}"; do
    gpu=${GPU_ARRAY[$i]}
    echo "Launching worker on GPU $gpu..."
    run_gpu_worker $gpu $i &
    PIDS+=($!)
 done
 echo ""
 echo "All workers launched. Waiting for completion..."
 echo "Monitor progress with: tail -f $RESULT_DIR/gpu_*.log"
 echo ""
 # Wait for all workers to complete
 for pid in "${PIDS[@]}"; do
    wait $pid
 done
 # End time
 END_TIME=$(date +%s)
 DURATION=$((END_TIME - START_TIME))
 echo ""
 echo "============================================================"
 echo "FINAL RESULTS"
 echo "============================================================"
 # Aggregate results
 TOTAL_PASS=0
 TOTAL_FAIL=0
 for gpu in "${GPU_ARRAY[@]}"; do
    if [ -f "$RESULT_DIR/gpu_${gpu}.summary" ]; then
        read pass fail < "$RESULT_DIR/gpu_${gpu}.summary"
        TOTAL_PASS=$((TOTAL_PASS + pass))
        TOTAL_FAIL=$((TOTAL_FAIL + fail))
        echo "GPU $gpu: $pass passed, $fail failed"
    fi
 done
 TOTAL_TESTED=$((TOTAL_PASS + TOTAL_FAIL))
 if [ $TOTAL_TESTED -gt 0 ]; then
    ACCURACY=$(echo "scale=1; $TOTAL_PASS * 100 / $TOTAL_TESTED" | bc)
 else
    ACCURACY="0.0"
 fi
 echo ""
 echo "------------------------------------------------------------"
 echo "Total: $TOTAL_PASS/$TOTAL_TESTED passed ($ACCURACY%)"
 echo "Duration: ${DURATION}s ($(echo "scale=1; $DURATION / 60" | bc) minutes)"
 echo "Throughput: $(echo "scale=2; $TOTAL_TESTED * 60 / $DURATION" | bc) samples/min"
 echo "------------------------------------------------------------"
 # Save detailed results
 {
    echo "RULER NIAH Parallel Test Results"
    echo "================================"
    echo "Date: $(date '+%Y-%m-%d %H:%M:%S')"
    echo "GPUs: $GPUS"
    echo "Total samples: $TOTAL_TESTED"
    echo "Passed: $TOTAL_PASS"
    echo "Failed: $TOTAL_FAIL"
    echo "Accuracy: $ACCURACY%"
    echo "Duration: ${DURATION}s"
    echo ""
    echo "Per-sample results:"
    for i in $(seq 0 $((TOTAL_SAMPLES - 1))); do
        if [ -f "$RESULT_DIR/sample_${i}.result" ]; then
            result=$(cat "$RESULT_DIR/sample_${i}.result")
            echo "Sample $i: $result"
        fi
    done
 } > "$OUTPUT_LOG"
 echo ""
 echo "Detailed results saved to: $OUTPUT_LOG"
 # Cleanup
 # rm -rf "$RESULT_DIR"
 # Exit with appropriate code
 if [ $TOTAL_FAIL -eq 0 ]; then
    echo ""
    echo "test_ruler_niah.sh: ALL PASSED"
    exit 0
 else
    echo ""
    echo "test_ruler_niah.sh: $TOTAL_FAIL FAILED"
    exit 1
 fi
Author	SHA1	Message	Date
Zijie Tian	cf168fd9b9	✅ test: add comprehensive RULER benchmark test suite - Add test_ruler.py supporting all 13 RULER tasks (NIAH, QA, CWE, FWE, VT) - Implement RULER official evaluation metrics (string_match_all/part) - Fix max_model_len to 32896 to prevent decode OOM on long inputs - Add ruler_benchmark_report.md with full test results (92.1% accuracy) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-14 00:51:30 +08:00
Zijie Tian	76af506956	[claudesquad] update from 'multi-request-2' on 13 Jan 26 02:01 CST	2026-01-13 02:01:07 +08:00
Zijie Tian	49519c7ce7	📝 docs: update offload accuracy issue with independent testing results Document key finding: single request inference works correctly (100% accuracy). The 66% accuracy issue in batch mode is due to state accumulation between sequential requests in the same process. - Add comparison table: independent (100%) vs batch (66%) testing modes - Document root cause analysis: state cleanup issue between requests - Add workaround using test_ruler_niah.sh for independent testing - Update next steps to focus on OffloadEngine reset/cleanup logic Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-12 21:08:35 +08:00
Zijie Tian	1424e665e7	✅ test: add parallel multi-GPU RULER NIAH test script Add test_ruler_niah.sh for independent sample testing across multiple GPUs. Each sample runs in a separate Python process to avoid state accumulation issues. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-12 21:08:27 +08:00
Zijie Tian	64971c8e8a	Merge branch 'zijie/fix-dist-3': Fix distributed port conflict - Auto port allocation with _find_free_port() in model_runner.py - Resource management refactor with close() + context manager in llm_engine.py - Add tests/test_port_conflict.py and tests/run_parallel_niah.sh - Remove docs/torch_distributed_port_issue.md (issue fixed) - Ignore tests/data/ directory Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-12 16:27:25 +08:00
Zijie Tian	de6f36bdb2	[docs] Added dist port issue.	2026-01-12 15:16:39 +08:00
Zijie Tian	8e0888c20c	[docs] Added offload_acc issue.	2026-01-12 15:05:55 +08:00
Zijie Tian	a6cc703d73	[tests] Added test_niah_standalone.py.	2026-01-12 00:16:37 +08:00
Zijie Tian	5895de0c97	[docs] Added transformers error desp.	2026-01-11 18:48:50 +08:00
Zijie Tian	2771312565	[docs] Add sparse prefill integration plan from int-minference analysis Consolidated analysis from int-minference-1/2/3 branches into a unified integration plan for MInference, XAttention, and FlexPrefill strategies. Key design decisions: - Backward compatible: Keep existing SparsePolicy interface - Unified BlockMask intermediate representation for new strategies - XAttention/FlexPrefill use block_sparse_attn_func kernel - MInference can optionally use block_sparse_attn (Phase 4) Five-phase implementation plan: 1. BlockMask + block_sparse_attn wrapper 2. XAttention implementation 3. FlexPrefill implementation 4. Optional MInference refactoring 5. Integration and testing Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-10 23:33:09 +08:00
Zijie Tian	de6eae472d	[docs] Update CLAUDE.md with multi-model support documentation - Update overview to reflect Qwen3/Qwen2/Llama support - Add docs/multi_model_support.md to documentation index - Add Llama-3.1-8B-Instruct to model limits Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-10 21:29:39 +08:00
Zijie Tian	e23be2e844	Merge branch 'zijie/add-llama-1': Add multi-model support - Add model registry system for dynamic model loading - Implement LlamaForCausalLM with Llama3 RoPE scaling - Register Qwen3ForCausalLM and Qwen2ForCausalLM - Update ModelRunner to use get_model_class() for dynamic model selection Tested: needle 32k test PASSED Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-10 21:20:53 +08:00
Zijie Tian	24f5ae5fc3	[claudesquad] update from 'add-llama-1' on 10 Jan 26 21:14 CST	2026-01-10 21:14:32 +08:00
Zijie Tian	9377ff63fe	Merge remote-tracking branch 'origin/zijie/fix-bug-2' into tzj/vs_offload	2026-01-09 16:13:38 +08:00
Zijie Tian	067e36f4a2	[claudesquad] update from 'fix-bug-2' on 09 Jan 26 16:10 CST	2026-01-09 16:10:28 +08:00
Zijie Tian	1425510a2e	[claudesquad] update from 'fix-bug-2' on 09 Jan 26 16:05 CST	2026-01-09 16:05:36 +08:00
Zijie Tian	335117bfca	Merge remote-tracking branch 'origin/zijie/fix-bug-2' into tzj/vs_offload	2026-01-09 15:21:48 +08:00
Zijie Tian	5012b11291	[bench] Modify bench_vllm.py	2026-01-09 15:20:37 +08:00
Zijie Tian	ccf04d3917	[claudesquad] update from 'fix-bug-2' on 09 Jan 26 15:16 CST	2026-01-09 15:16:55 +08:00
Zijie Tian	59f8970ed3	[claudesquad] update from 'fix-bug-2' on 09 Jan 26 15:12 CST	2026-01-09 15:12:42 +08:00
Zijie Tian	6378cb4c17	Merge remote-tracking branch 'origin/zijie/fix-ga-perf-2' into tzj/vs_offload	2026-01-09 14:21:00 +08:00
Zijie Tian	47e3e465f0	[claudesquad] update from 'fix-ga-perf-2' on 09 Jan 26 14:08 CST	2026-01-09 14:08:12 +08:00
Zijie Tian	aac94c9481	[claude] Added some commands.	2026-01-09 13:16:23 +08:00
Zijie Tian	79c4df4a27	[claudesquad] update from 'int-minference-1' on 08 Jan 26 23:42 CST	2026-01-08 23:42:30 +08:00
Zijie Tian	ea4e904de0	[claudesquad] update from 'int-minference-1' on 08 Jan 26 23:22 CST	2026-01-08 23:22:38 +08:00
Zijie Tian	0bfe1984ef	[docs] Refine GPU mutex: exclusive for benchmarks, port check for tests Benchmarks (bench*.py) still require exclusive GPU access for accurate measurements. Other scripts (tests, examples) now only check for distributed port 29500 conflicts, allowing parallel GPU sharing. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-08 21:35:08 +08:00
Zijie Tian	105201b902	[claudesquad] update from 'lw-offload-2' on 08 Jan 26 21:19 CST	2026-01-08 21:19:38 +08:00
Zijie Tian	a8c9f0d837	[claudesquad] update from 'lw-offload-2' on 08 Jan 26 20:53 CST	2026-01-08 20:53:08 +08:00
Zijie Tian	85bcca3d17	[claudesquad] update from 'int-offload-1' on 08 Jan 26 19:44 CST	2026-01-08 19:44:29 +08:00
Zijie Tian	b5c0ef3b7a	[docs] Replace chunked prefill docs with layer-wise offload strategy Remove all chunked prefill related documentation (ring buffer, sgDMA, Triton merge kernels, known issues) and replace with layer-wise offload system documentation including: - Design philosophy and benefits - Memory layout and per-layer KV size table - Prefill and decode flow pseudocode - Critical implementation details (sync offload, causal=False for decode) - Helper methods in HybridKVCacheManager 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-08 05:39:26 +08:00
Zijie Tian	bbbfd1e7da	[docs] Simplify multi-instance development with direct PYTHONPATH Replace pip install -e . --prefix=./.local approach with simpler PYTHONPATH method: - No pip install required - Code changes take effect immediately - Each worktree is completely isolated Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-08 04:51:55 +08:00
Zijie Tian	c1ddb44e5d	Merge branch 'zijie/layer-prefill-1' into tzj/vs_offload Adds MInference sparse attention support: - New MInference sparse policy implementation - A-shape, vertical-slash, and block-sparse patterns - Updated bench.py with sparse attention options - test_minference_gpu.py validation test 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>	2026-01-08 03:40:53 +08:00
Zijie Tian	d8a87da1c3	[claudesquad] update from 'layer-prefill-1' on 08 Jan 26 03:36 CST	2026-01-08 03:36:39 +08:00
Zijie Tian	ecd9ae0271	[WIP] changed to layerwise offload.	2026-01-08 00:28:27 +08:00