Introduce SCV vectorized masking with graph/adaptive modes #3
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
added 7 commits
October 15, 2025 16:10
…omputation Introduced SCV (Speculative Computation Vectorization) mode to GPUModelRunner to optimize mask computation during decoding. Added SCVGraphExecutor and _SCVGraphEntry classes leveraging CUDA Graphs for efficient repeated mask calculations. The SCV mode supports 'graph' and 'adaptive' operation and falls back gracefully if CUDA graph execution fails. This enhancement improves decoding performance by reusing captured CUDA graphs for mask operations in speculative decoding workflows.
…peculation tokens Introduce an adaptive mode in GPUModelRunner to dynamically compute and adjust the speculation token mask based on recent acceptance ratios during decoding. This update adds the _scv_update_controller method to modify the number of speculative tokens used, aiming to maintain a target acceptance ratio and improve decoding efficiency.
…V adaptive mode Add the test `test_scv_vectorized_mask_matches_reference` to validate the `_build_nwor_acceptance_mask` behavior in GPUModelRunner when SCV adaptive mode is enabled, ensuring the mask output matches the expected reference.
…lization code Relocate the `_scv_enabled` method within GPUModelRunner to sit alongside the initialization logic, improving readability without altering behavior.
yuz207
deleted the
terragon/remove-bot-coauthors-and-scrub-commits-9d2ndb
branch
yuz207
added a commit
that referenced
this pull request
Oct 19, 2025
This commit implements five correctness-preserving optimizations that reduce GPU-CPU synchronization overhead in speculative decoding paths without changing behavior. Estimated total speedup: 5-11ms per decode step. Optimization #1: Batch mask sum operations (⭐⭐⭐) - Before: N GPU-CPU syncs (one per request) via .sum().item() in loop - After: Single batched sync via torch.stack().cpu() for all requests - Impact: Reduces 4-8ms overhead to ~0.5ms for typical batch sizes - Locations: Lines 2712-2740 (SCV path), 2757-2829 (fallback path) - Safety: Guards against empty sum_tensors to prevent stacking errors Optimization #2: Eliminate CPU transfer in SCV cache key (⭐⭐⭐) - Before: cu_int32.cpu().tolist() forces GPU->CPU sync on every SCV call - After: Use itertools.accumulate() to compute cumsum directly on CPU - Impact: Removes 0.5-2ms overhead per SCV call, even for cache hits - Location: Lines 2893-2900 - Safety: Uses spec_decode_metadata.num_draft_tokens (already CPU list) Optimization #3: Combine device/dtype conversions (⭐⭐) - Before: Two sequential .to() calls launch two separate kernels - After: Single .to(device=..., dtype=...) launches one combined kernel - Impact: 2x faster conversions (~0.3ms saved) - Locations: Lines 2749-2750, 2882-2883 - Safety: PyTorch API guarantees identical behavior for combined .to() Optimization #4: Hoist device/dtype checks outside loop (⭐⭐) - Before: Per-request device/dtype checks and conversions inside loop - After: Single conversion before loop (tensor slices inherit properties) - Impact: Eliminates 0.1-0.5ms per-request overhead - Location: Lines 2771-2772 (moved from inside loop at 2782-2785) - Safety: PyTorch guarantees all rows share parent tensor's device/dtype Optimization #5: Cache _nwor_debug lookup (⭐) - Before: Duplicate getattr() calls at lines 2640 and 2644 - After: Single lookup cached in local variable - Impact: Negligible performance, cleaner code - Location: Line 2639 - Safety: Trivial refactor with identical semantics All optimizations maintain exact correctness while eliminating redundant GPU-CPU synchronization points and duplicate kernel launches. No changes to NWOR/SCV algorithms or numerical results.
yuz207
added a commit
that referenced
this pull request
Oct 22, 2025
- Implement commit_draft_kernel copying exact pattern from reshape_and_cache_flash - Support both NHD and HND cache layouts (Flash/Paged) - Full dtype dispatch: fp16/bf16/fp32 source, auto/fp8/fp8_e5m2 cache - Proper quantization with CopyWithScaleOp template - Per-token and scalar scale support - Mask early-return optimization (Issue #3) - TORCH_CHECK validation for all pointers (Issue #7) - Add key_value_dtype to DraftEntry for source dtype tracking This is Phase 3 (CUDA kernel) of the draft commit implementation. Next: PyTorch bindings + integration hooks.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
No description provided.