[GG] perf(nf3): integrate Grid188 hybrid decode - #113
Conversation
Co-authored-by: David Young <david@davidsyoung.com>
Co-authored-by: David Young <david@davidsyoung.com>
📝 WalkthroughWalkthroughThe PR adds an exact Grid188 fused decode path for hybrid NVFP4/MXFP4 and NF3 layers, recognizes the hybrid online overlay, and adds tests for descriptor encoding. It also changes DCP rank localization and speculative decoding attention metadata bounds. ChangesGrid188 hybrid decode
DCP attention localization
Speculative decode metadata
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant HybridLayer
participant HybridRuntime
participant Grid188Kernel
HybridLayer->>HybridRuntime: initialize runtime
HybridRuntime->>HybridRuntime: prepare eligible Grid188 state
HybridLayer->>HybridRuntime: apply four-token batch
HybridRuntime->>Grid188Kernel: launch fused hybrid decode
Grid188Kernel-->>HybridLayer: return output
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
vllm/model_executor/layers/quantization/nvfp4_nf3_hybrid.py (1)
1029-1029: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winUse vLLM’s cached
current_stream()helper in this decode hot path.
torch.cuda.current_stream()constructs a stream wrapper on every eligible MoE layer and token, adding avoidable host overhead to the new fast path.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/model_executor/layers/quantization/nvfp4_nf3_hybrid.py` at line 1029, Update the decode hot path around the current_stream argument to use vLLM’s cached current_stream() helper instead of torch.cuda.current_stream(). Preserve the existing device selection and integer CUDA stream handle passed to the downstream operation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/v1/attention/test_indexer_dcp_localize.py`:
- Around line 309-319: Add a CUDA-gated variant of
test_get_dcp_local_seq_lens_rank_does_not_allocate_cuda_scalar that creates
seq_lens on CUDA and verifies the same localized result, preserving the
torch.tensor allocation guard; alternatively, rename the existing test to
explicitly describe its CPU-only guarantee.
---
Nitpick comments:
In `@vllm/model_executor/layers/quantization/nvfp4_nf3_hybrid.py`:
- Line 1029: Update the decode hot path around the current_stream argument to
use vLLM’s cached current_stream() helper instead of
torch.cuda.current_stream(). Preserve the existing device selection and integer
CUDA stream handle passed to the downstream operation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 593ec3e4-04c4-48d2-ac97-8c669dcb87af
📒 Files selected for processing (7)
tests/quantization/test_nvfp4_nf3_hybrid.pytests/v1/attention/test_indexer_dcp_localize.pyvllm/config/quantization.pyvllm/envs.pyvllm/model_executor/layers/quantization/nvfp4_nf3_hybrid.pyvllm/v1/attention/backends/utils.pyvllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
| def test_get_dcp_local_seq_lens_rank_does_not_allocate_cuda_scalar(monkeypatch): | ||
| seq_lens = torch.tensor([17, 33, 65], dtype=torch.int32) | ||
|
|
||
| def fail_tensor_allocation(*args, **kwargs): | ||
| raise AssertionError("rank localization must not allocate a scalar tensor") | ||
|
|
||
| monkeypatch.setattr(torch, "tensor", fail_tensor_allocation) | ||
|
|
||
| actual = get_dcp_local_seq_lens(seq_lens, dcp_size=4, dcp_rank=2) | ||
|
|
||
| assert actual.tolist() == [4, 8, 16] |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win
Exercise the CUDA behavior this regression test claims to cover.
Because seq_lens is CPU-backed, the test only proves that torch.tensor is not called; it cannot detect a CUDA-only scalar conversion or synchronization regression. Add a CUDA-gated case with seq_lens on CUDA, or rename the test to state its narrower guarantee.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/v1/attention/test_indexer_dcp_localize.py` around lines 309 - 319, Add
a CUDA-gated variant of
test_get_dcp_local_seq_lens_rank_does_not_allocate_cuda_scalar that creates
seq_lens on CUDA and verifies the same localized result, preserving the
torch.tensor allocation guard; alternatively, rename the existing test to
explicitly describe its CPU-only guarantee.
|
Final baked-image validation on
This matches the clean v17 MTP0 result within run variance and retains the measured Grid188 MTP3 decode gain. |
|
Superseded by the canonical dev/gilded-gnosis consolidation. The corresponding implementation is present in commit(s): 45837b4,1dc20c2b7f. Closing the old-base PR so future work targets the canonical GG branch. |
Summary
Requires local-inference-lab/b12x#36.
Grid188 is enabled by default for its exact geometry and can be disabled with VLLM_NF3_GRID188_DECODE=0 for A/B comparison.
Validation
Benchmark
Checkpoint: madeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid with the exact 546-module MXFP8 selector, A16, FP8 KV for the matched old-v18 overlay test, CC1, three 60-second runs, no concurrent model loading.
The Grid188 median gain over the serial fallback is +10.9%. The additional host-sync median gain is +1.16%; its throughput effect is smaller and partially coupled to probabilistic MTP acceptance variance, so it is kept as a separate commit.
Provenance
The Grid188 design and the two host-sync observations were published in davidsyoung/vllm-glm52#2. This PR ports the required symbols onto GG rather than replacing GG files with release overlays.
Summary by CodeRabbit
New Features
nvfp4_nf3_hybridModelOpt online quantization configuration.VLLM_NF3_GRID188_DECODEenvironment option to control the specialized decode path.Performance
Bug Fixes