Proof of concept #2
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
* position ids for rope * cleanup * no need for mask * no mask * more cleanup * add back filtering * more cleanup * revert the signature of llm's generate and forward * use self.decoder.lm_use_tokens * use torch inference_mode * add back comment * fix bug * add back comments * add back comments
Co-authored-by: Kashif Rasul <[email protected]>
Co-authored-by: Kashif Rasul <[email protected]>
always use
Implementing KV Cache for inference
Multi-node training (and a few other things, should have created more PRs!)
- Add export_executorch.py with dynamic shapes and int8 quantization - Add test_executorch_export.py for end-to-end inference testing - Add test_executorch_accuracy.py for numerical accuracy validation - Add ONNX export scripts in onnx_export/ - Fix language_model.py for export compatibility: - Remove dynamic RoPE scaling (data-dependent control flow) - Add position_ids parameter to forward() - Fix SDPA to use explicit masks instead of is_causal ExecuTorch export produces 6.0GB (unquantized) or 2.3GB (quantized) models that generate accurate descriptions. Both exports tested and working. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
This proof of concept implements full C++ inference for nanoVLM using ExecuTorch .pte models with multi-image preprocessing via Rust. Major Components Added: - cpp-inference/: Complete C++ ExecuTorch inference pipeline - Multi-image support with grid position tokens - Proper JSON config loading using nlohmann/json - KV cache management for autoregressive generation - Embedding replacement for vision-language fusion - rust-preprocessor/: Rust preprocessing library with C FFI - Image splitting (global view + grid patches) - Dynamic resize matching Python implementation - Tokenization with special tokens - Text decoding - Documentation: - CPP_INFERENCE_STATUS.md: Comprehensive status report - BUILD_LOG.md: Build configuration details - EXPORT_NOTES.md: Model export documentation Key Features: ✅ All 6 models load and execute (vision, projector, prefill, decode, embedding, lm_head) ✅ Image splitting: 17 images (1 global + 16 patches in 4x4 grid) ✅ Grid token generation: <|global_image|>, <row_X_col_Y> ✅ Embedding replacement verified working ✅ KV cache: 60 tensors (30 blocks × 2) ✅ Greedy sampling with EOS detection Known Issues:⚠️ Tokenization mismatch: C++ generates 1252 tokens vs Python's 1118 This causes garbled output despite correct infrastructure Export Updates: - Added --use-xnnpack flag to export_executorch.py for backend control - Supports both portable ops and XNNPack delegation Build Requirements: - ExecuTorch with custom_ops (for SDPA) - Rust toolchain for preprocessing library - nlohmann/json (included in ExecuTorch) 🤖 Generated with Claude Code (https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Previously, grid position tokens like <|global_image|> and <row_X_col_Y> were not registered as special tokens in the Rust tokenizer. This caused them to be tokenized as multiple regular tokens instead of single token IDs. Changes: - Add <|global_image|> as special token - Add all 64 <row_X_col_Y> tokens (8x8 grid) as special tokens - Keep existing <|image|> token Impact: - Token count reduced from 1252 to 1118 (matching Python) - C++ inference now generates correct token sequences - Each grid token now encoded as 1 token instead of 7-8 tokens 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Root cause: Rust tokenizer only tracks <|image|> tokens (49152), not <|global_image|> tokens (49153). Python replaces BOTH types with image embeddings, but C++ was only replacing tokenizer-tracked positions. This caused position 3 (the <|global_image|> token) to have a text embedding instead of the global image embedding, making the entire prefill output incorrect. Fix: Modified cpp-inference/main.cpp to manually scan ALL tokens and replace both <|image|> and <|global_image|> tokens with embeddings, matching Python's behavior. Result: - Combined embeddings now match Python (max diff: 0.000153) - Prefill hidden states match Python perfectly - First token prediction matches Python (token ID: 49) Added: - TOKEN_REPLACEMENT_FIX.md: Detailed documentation of the bug and fix - test_prefill_with_python_inputs.cpp: Test to verify prefill with Python inputs - Debug output in test_executorch_pte.py to track token replacement See TOKEN_REPLACEMENT_FIX.md for full debugging process and verification. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
After fixing the image token replacement bug, the decode loop still produces incorrect output. This commit adds comprehensive debug instrumentation and documentation to investigate the root cause. Root Cause Identified: KV cache reference invalidation - EValue references to Module-owned tensors become invalid when forward() is called again on subsequent decode steps. The Module likely reuses/overwrites internal buffers, causing stored references to point to stale data. Evidence: - Prefill stage: ✅ Perfect match with Python - Decode step 1: ✅ Correct (49 → 2800) - Decode step 2 onward: ❌ Hidden states diverge, causing wrong predictions Changes: 1. Added debug output to cpp-inference/main.cpp for first 3 decode steps: - Token embeddings, KV cache shapes, hidden states, logits 2. Added matching debug output to test_executorch_pte.py: - Allows direct Python vs C++ comparison at each step 3. Added TODO comments marking the KV cache bug locations 4. Created DECODE_LOOP_INVESTIGATION.md: - Detailed analysis of the KV cache reference invalidation issue - Evidence comparing Python vs C++ outputs - Attempted fix (copying KV cache data) and why it failed - Potential solutions to explore 5. Rust preprocessor improvements: - Fixed token ordering to match Python (image_token, global_image_token, grid tokens) - Added bicubic interpolation for global view resizing (matches Python's BICUBIC) - Clarified comments about which tokens get replaced with embeddings Next Steps: See DECODE_LOOP_INVESTIGATION.md for potential solutions including: - ExecuTorch Module output ownership investigation - Persistent output buffers approach - Checking Python ExecuTorch implementation for comparison 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
This fixes the token repetition bug where the decode loop would generate the same token repeatedly (e.g., "The image't't't't't't't't"). Problem: - EValue references to KV cache tensors became stale when Module reused internal buffers between forward() calls - Caused decode loop to use incorrect/corrupted KV cache data - Manifested as repeated tokens starting at decode step 3 Solution: - Use ExecuTorch's clone_tensor_ptr() API to create deep copies of KV cache tensors with owned data - Store cloned tensors in std::vector<TensorPtr> - Convert back to EValue when passing to forward() Changes: - Added TensorPtr and clone_tensor_ptr imports - After prefill: clone all KV cache tensors from outputs - In decode loop: convert TensorPtr to EValue for forward() inputs - After decode: clone updated KV cache tensors from outputs - Removed verbose debug output from decode loop Result: - Token repetition bug completely fixed - Generates diverse, coherent text without crashes - Tested with 50+ token generation successfully 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
We export to Executorch and build some Python and C++ runners while we manage our own KV cache due to perf issues with the self-managed KV cache in Executorch itself..