[GG] Add EXL3 Trellis backend for rank-sliced MoE checkpoints - #190
Conversation
|
Warning Review limit reached
Next review available in: 38 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe changes add mixed-bitrate EXL3 rank-sliced MoE loading and Trellis execution, draft metadata hydration, speculative structured-output token handling, backend-aware sparse-indexer scheduling, DCP cache sharding, defensive KV-transfer handling, and NVFP4 environment registration updates. ChangesEXL3 rank-sliced quantization
Speculative structured-output advancement
Attention indexer and DCP scheduling
Scheduler state validation
NVFP4 environment contract
Estimated code review effort: 5 (Critical) | ~90+ minutes Possibly related issues
Possibly related PRs
🚥 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: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
vllm/v1/structured_output/__init__.py (1)
428-432: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winReset minimax streaming end-state before same-step marker localization
_find_reasoning_end_index()replaysis_reasoning_end_streaming()token-by-token whenreasoner.is_reasoning_end_streaming(all_token_ids, delta)already returnedTrue. That state is latched forminimax_m3_reasoning_parservia_reasoning_ended_streaming, so the localizer can return the first token atstartinstead of the actual marker and trim too little reasoned output. Reset the local streaming end flag before the per-token scan, or avoid mutating the parser state while localizing the index.🤖 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/v1/structured_output/__init__.py` around lines 428 - 432, Update the speculative reasoning-end path around _find_reasoning_end_index so the parser’s latched _reasoning_ended_streaming state is cleared before per-token localization, or ensure localization does not mutate parser streaming state. Preserve the actual reasoning marker index and existing structured_req.reasoning_end_token_index assignment.
🧹 Nitpick comments (4)
tests/quantization/test_exl3.py (1)
160-172: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThese two tests re-implement the production logic they aim to protect.
_key()andresolved()reconstruct the cache key and the min-mdefault locally, so they pass even ifExl3MoEMethod._rank_sliced_runtimestops prefixing the scope or stops branching on_is_draft_layer. Extracting the key/default computation into small module-level helpers inexl3.pyand asserting on those would make the regression coverage real.Also applies to: 229-240
🤖 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/quantization/test_exl3.py` around lines 160 - 172, The tests’ local _key() and resolved() helpers duplicate production cache-key and min-m computation, so they do not verify Exl3MoEMethod._rank_sliced_runtime behavior. Extract the scope-prefixed key construction and _is_draft_layer-dependent default computation into small module-level helpers in exl3.py, update _rank_sliced_runtime to use them, and have the tests assert those shared helpers instead of reimplementing the logic.vllm/model_executor/layers/quantization/exl3.py (1)
68-71: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winModule-global runtime cache is never evicted and its scope counter is mutated unguarded.
_RANK_SLICED_RUNTIMESretains plans plus ~1 GiB prefill arenas and parity staging for the process lifetime, keyed by a monotonically increasing scope id, so sleep/wake or engine re-instantiation in one process leaks the previous model's arenas rather than reusing them._NEXT_RUNTIME_SCOPE_IDis also incremented without synchronization; harmless under one worker thread per process, but worth a note or a lock if planning can ever be reached from more than one thread.Also applies to: 1685-1685
🤖 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/exl3.py` around lines 68 - 71, Update the runtime-cache lifecycle around _RANK_SLICED_RUNTIMES and _NEXT_RUNTIME_SCOPE_ID so plans and their prefill/parity arenas are evicted when a model scope is destroyed or replaced, allowing later sleep/wake or engine instantiation to release and reuse resources. Synchronize allocation of the scope ID and associated cache access so concurrent planning cannot race; apply the same lifecycle and synchronization behavior at the additional cache-use site.vllm/envs.py (1)
2255-2272: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeclare the nine new EXL3 environment variables in the module-level type annotations and route consumers through
envs.py.The new registry entries are registered in
environment_variables, but the correspondingTYPE_CHECKINGannotations are still missing invllm/envs.py, which can leavevllm.envs.*lookup/type-checker support out of sync. The consumers invllm/model_executor/layers/quantization/exl3.pyread these throughos.environ/os.getenvdirectly rather thanenvs.VLLM_EXL3_*, so the registry additions do not enforce a single canonical getter 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/envs.py` around lines 2255 - 2272, Add TYPE_CHECKING module-level annotations in vllm.envs.py for all nine EXL3 variables registered in environment_variables, including the MLA scales variable. Update consumers in exl3.py to read these settings through envs.VLLM_EXL3_* instead of direct os.environ/os.getenv access, preserving existing parsing and default behavior.tests/v1/structured_output/test_reasoning_structured_output.py (1)
291-295: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRuff RUF005: prefer unpacking over list concatenation (same at Lines 338-342).
♻️ Proposed change
- request.all_token_ids = request.prompt_token_ids + [ - 7, - end_token_id, - content_token_id, - ] + request.all_token_ids = [ + *request.prompt_token_ids, + 7, + end_token_id, + content_token_id, + ]🤖 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/structured_output/test_reasoning_structured_output.py` around lines 291 - 295, Update the token ID construction in the affected test setup, including the matching occurrence, to use iterable unpacking instead of concatenating request.prompt_token_ids with a list. Preserve the existing token order and values.Source: Linters/SAST tools
🤖 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 `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 1504-1513: In the initialization validation alongside the
min_trellis_m/max_trellis_m check, validate that the parity capacity can cover
every m below min_trellis_m: reject configurations where min_trellis_m - 1
exceeds the configured prefill capacity represented by chunk (and
max_batched_tokens when applicable). Raise a clear ValueError at plan time,
before the parity path in the prefill execution flow can encounter the oversized
batch.
In `@vllm/model_executor/models/glm4_moe.py`:
- Around line 525-528: Assign the constructor’s quant_config argument to
self.quant_config alongside self.config = config in the relevant GLM model
initializer. Preserve the existing normalize_rank_sliced_weight_name lookup so
weight loading can access the stored configuration.
In `@vllm/model_executor/models/utils.py`:
- Around line 813-830: Restrict the EXL3 hydration block around
quant_config.maybe_update_config to cases where either draft_hf or target_hf
defines hybrid_tr3_tail. Skip the entire update, including any remote
configuration fetch, when neither config is rank-sliced; preserve the existing
draft-versus-target hf_config selection for hybrid cases.
In `@vllm/v1/attention/backends/mla/indexer.py`:
- Around line 413-429: Use a single cached result from use_b12x_sparse_indexer()
in the indexer builder, including self.use_flattening and the later B12X
schedule metadata, active-width setup, and prefill chunking branches around the
relevant builder logic. Replace direct envs.VLLM_USE_B12X_SPARSE_INDEXER checks
with this cached predicate so backend-only B12X selection consistently uses the
native path and metadata.
---
Outside diff comments:
In `@vllm/v1/structured_output/__init__.py`:
- Around line 428-432: Update the speculative reasoning-end path around
_find_reasoning_end_index so the parser’s latched _reasoning_ended_streaming
state is cleared before per-token localization, or ensure localization does not
mutate parser streaming state. Preserve the actual reasoning marker index and
existing structured_req.reasoning_end_token_index assignment.
---
Nitpick comments:
In `@tests/quantization/test_exl3.py`:
- Around line 160-172: The tests’ local _key() and resolved() helpers duplicate
production cache-key and min-m computation, so they do not verify
Exl3MoEMethod._rank_sliced_runtime behavior. Extract the scope-prefixed key
construction and _is_draft_layer-dependent default computation into small
module-level helpers in exl3.py, update _rank_sliced_runtime to use them, and
have the tests assert those shared helpers instead of reimplementing the logic.
In `@tests/v1/structured_output/test_reasoning_structured_output.py`:
- Around line 291-295: Update the token ID construction in the affected test
setup, including the matching occurrence, to use iterable unpacking instead of
concatenating request.prompt_token_ids with a list. Preserve the existing token
order and values.
In `@vllm/envs.py`:
- Around line 2255-2272: Add TYPE_CHECKING module-level annotations in
vllm.envs.py for all nine EXL3 variables registered in environment_variables,
including the MLA scales variable. Update consumers in exl3.py to read these
settings through envs.VLLM_EXL3_* instead of direct os.environ/os.getenv access,
preserving existing parsing and default behavior.
In `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 68-71: Update the runtime-cache lifecycle around
_RANK_SLICED_RUNTIMES and _NEXT_RUNTIME_SCOPE_ID so plans and their
prefill/parity arenas are evicted when a model scope is destroyed or replaced,
allowing later sleep/wake or engine instantiation to release and reuse
resources. Synchronize allocation of the scope ID and associated cache access so
concurrent planning cannot race; apply the same lifecycle and synchronization
behavior at the additional cache-use site.
🪄 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: acd2a1e9-7b08-4d07-a9e1-e38ad968ca18
📒 Files selected for processing (16)
tests/quantization/test_exl3.pytests/quantization/test_exl3_prefill_plan.pytests/v1/structured_output/test_reasoning_structured_output.pyvllm/config/model.pyvllm/envs.pyvllm/model_executor/layers/fused_moe/routed_experts.pyvllm/model_executor/layers/quantization/__init__.pyvllm/model_executor/layers/quantization/exl3.pyvllm/model_executor/models/deepseek_mtp.pyvllm/model_executor/models/deepseek_v2.pyvllm/model_executor/models/glm4_moe.pyvllm/model_executor/models/utils.pyvllm/models/deepseek_v32/nvidia/attention.pyvllm/v1/attention/backends/mla/indexer.pyvllm/v1/core/sched/scheduler.pyvllm/v1/structured_output/__init__.py
Adds --quantization exl3: any exllamav3 EXL3 checkpoint loads through a QuantizationConfig + LinearMethodBase pair that reproduces turboderp's dense no_reconstruct forward argument-for-argument at the exl3_gemm boundary (input suh sign-flip + Hadamard-128, direct trellis GEMM, output svh + Hadamard-128), with mcg/mul1 sentinel semantics, legacy su/sv expansion, reference pad/trim, and tp_import-discipline TP slicing (KV replication; Qwen3.5 [(0,1,2),3] tuple shards). Registry delta is exactly three lines (Literal member, lazy import, mapping entry). RoutedExperts.load_weights: qualify is_fused by the matched mapping entry, not tensor rank alone — per-expert EXL3 trellis tensors are rank-3 [K/16, N/16, 16*bpw] and were mangled by the fused transpose/chunk path (fail-loud downstream, but wrong). Fused entries are exactly those whose weight_name carries no per-expert index; behavior for genuinely fused checkpoints is unchanged. Backend enforces eager execution at build time: exl3_gemm autotunes with timing launches on first call per (m-bucket, k, n, K) hash, which is incompatible with CUDA-graph capture, and m-bucketing defeats warmup coverage. Quantized lm_head with added vocabulary raises (TP slicing would silently misalign). MoE expert loader is marked supports_moe_loading for llama4-style loader paths. GPU validation gates (owner-run): smoke-load Qwen3-0.6B-exl3, generate on Qwen3.6-27B-exl3_3.30bpw, top-logprob comparison vs exllamav3 pinned to no_reconstruct=True. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Modules carrying the MCG codebook marker are prepared once at weight-processing time via b12x prepare_trellis256_dense_weight (zero-copy, bits inferred from native storage) and applied through run_trellis256_dense — the fused CuTeDSL kernel whose output is byte-identical across batch size m. Legacy default-codebook and MUL1 shards, and any shard the fused prepare rejects, fall back per-shard to the bit-faithful exllamav3_ext parity op, which remains the oracle. VLLM_EXL3_FUSED=0 forces the parity path (A/B switch); VLLM_EXL3_LOG_PATH_SELECTION=1 logs the per-shard selection. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A multimodal wrapper offers the quantized head as language_model.lm_head while the checkpoint storage map keys it as top-level lm_head. The segment-collapse loop only considered interior segments, so the head silently fell back to UnquantizedLinearMethod and loading the checkpoint's lm_head.mcg then failed. Collapse leading model/language_model segments too; exact-prefix candidates still take precedence. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The fused dense entry performs its outer rotations through exllamav3_ext.had_r_128, imported by module name inside b12x. Load the extension through the parity path's VLLM_EXL3_ABI_SHIM/VLLM_EXL3_EXT_PATH contract first, so an all-fused model does not import an unshimmed site-packages copy (undefined at::cuda::getCurrentCUDABlasHandle on torch 2.12). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Review follow-up (CodeRabbit): _moe_prefix_is_exl3 hardcoded gate_proj/up_proj/down_proj while _validate_codebooks keys off the layer's ckpt_*_proj_name fields, so a remapped-projection MoE checkpoint would pass validation logic but silently miss EXL3 detection. Thread the RoutedExperts layer through and read its names, defaulting to the standard trio for layer variants that do not carry them. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Brandon Music <brandon.m.music@gmail.com>
…MTP crash Port two upstream vLLM fixes so tool calling works under MTP speculative decoding: - structured_output: advance the grammar from the authoritative new_token_ids step delta instead of the counter-derived window, so the reasoning-end (</think>) marker is detected under async scheduling + spec decode and the tool-call grammar engages instead of emitting unconstrained output. (vllm-project#48516, supersedes vllm-project#44993.) - deepseek_v32 DSA: derive has_indexer from index_k as well, so MTP draft steps 1+ (skip_topk under index_share_for_mtp_iteration) no longer crash fused_norm_rope's index_k assertion. (vllm-project#48528.) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Builds on the rank-sliced EXL3 (ExLlamaV3 trellis) support added in local-inference-lab#139 by @brandonmusic, paired with the trellis3_t256 planned MoE from local-inference-lab/sparkinfer#49. No change to that work's weight loading, formats, or decode path; this commit only extends its batch dispatch. The vllm-project#139 backend serves every batch with m > 32 through the eager exllamav3_ext parity path: fp16-in/fp32-out staging (~25 GB extra HBM traffic per 4096-token step at GLM-5.2 geometry), python chunk loops, and one expert re-stream per 128-row chunk. That covers 100% of prefill, so prefill throughput is capped at ~2.0-2.3k tok/s while the same weights decode at full speed through the planned Trellis window. Build a second trellis_moe plan at max_tokens=max_num_batched_tokens with block_size_m=64 (env VLLM_EXL3_PREFILL_BLOCK_M, allowed {8,16,32,48,64}) beside the existing decode plan (32/8), and dispatch three ways in _apply_rank_sliced: * m in [min, max] -> decode plan (unchanged) * max < m <= capacity -> prefill plan (sparkinfer bind already accepts any tokens in [1, max_tokens]; zero kernel changes needed) * m < min -> parity path, whose persistent staging (xh/out32/token_sorted/weight_sorted) shrinks from capacity rows to one chunk while the prefill plan is live (~110 MiB/GPU returned) VLLM_EXL3_PREFILL_TRELLIS=0 restores the exact single-plan parity behavior (one-variable serving A/B lever). The parity branch now raises during CUDA graph capture instead of silently recording eager ext calls. The prefill arena (~1.05 GiB/GPU at capacity 3072) allocates during the profile pass, so vLLM's memory profiler sees it and the KV pool auto-shrinks instead of risking request-time OOM. Matched A/B serving brandonmusic/GLM-5.2-EXL3-TR3-3.0bpw on the published verdictai/glm52-exl3-sparkinfer runtime image (4x RTX PRO 6000 Blackwell, TP4/DCP4/MTP3, identical boot geometry, only the kill-switch flipped, fresh JIT cache both arms; llm_decode_bench prefill-only, exact /tokenize, C1, 60 s/context): ctx parity trellis delta 8k 2,287 3,742 +63.6% 32k 2,218 3,557 +60.4% 64k 2,078 3,381 +62.7% 128k 1,920 3,032 +57.9% Decode C1 unchanged-to-better; 57k-token needle retrieval exact. (cherry picked from commit d57d775)
Optional companion to the prefill-plan commit; mocked plan/ext APIs, no GPU, sparkinfer, or exllamav3_ext required. Covers plan construction order, the three-way dispatch boundaries, chunk-capped parity staging, the VLLM_EXL3_PREFILL_TRELLIS=0 kill-switch, VLLM_EXL3_PREFILL_BLOCK_M override, elastic re-plan above scheduler capacity, and the parity-path capture guard. Follows the existing CPU-test pattern of tests/quantization/test_exl3.py from local-inference-lab#139. (cherry picked from commit f80a56d)
MTP next_n>2 was force-flattened onto the DeepGEMM next_n<=2 fallback, which corrupts B12X sparse-indexer code generation. The B12X/sparkinfer indexer supports native (B, next_n) on SM120, so keep it on the native path (FULL cudagraphs preserved). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The MTP draft head (deepseek_mtp) could not load a rank-sliced EXL3
(hybrid_tr3_tail) MoE layer -- e.g. a tr3-converted GLM-5.2 MTP layer 78 --
for two reasons:
1. get_draft_quant_config() built the draft Exl3Config via
get_quantization_config() but never called maybe_update_config(), which is
what hydrates rank_sliced_metadata from hybrid_tr3_tail (the target model
receives this call in config/vllm.py). Without it the draft config's
rank_sliced_metadata stayed None and the MTP experts fell back to the stock
FusedMoE path.
2. DeepSeekMTP.load_weights() did not normalize rank-sliced weight names, so a
checkpoint tensor `...experts.0.down_proj.rank0.mcg` mapped to
`...routed_experts.w2_rank0.mcg` while Exl3MoEMethod registers the fused
`w2_mcg` param -> KeyError. The target model (deepseek_v2.load_weights)
already strips `.rank{r}` via normalize_rank_sliced_weight_name; mirror it.
With both fixes a tr3 MTP layer loads and serves. (Running a tr3 MTP layer also
needs VLLM_EXL3_TRELLIS_MIN_M=1 so the MTP-N draft's small-m GEMMs stay inside
the Trellis cudagraph window -- a runtime env, not a code change.)
Validated on GLM-5.2-EXL3-TR3-3.0bpw with an EXL3-TR3 MTP layer-78 head (TP4,
DCP4, MTP-3): loads clean, Estonia 10/10 PASS, LAVD EXACT 5 / NEAR 5 / FAIL 0.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QsVTetPQQc5gP7LveJCWRE
The rank-sliced runtime cache was keyed only on device, dtype, shape, topk and planner settings. The cached value owns mutable Trellis/prefill scratch plus parity staging and sort buffers, so any two layers that hash equal share those allocations. A target MoE layer and a rank-sliced MTP draft layer collide on every component of that key: same hidden/intermediate size, same local expert count, same topk, identical planner env, and both resolve max_num_batched_tokens from the same scheduler config. The draft therefore reused the target's scratch, defeating the target/draft resource isolation that their independently captured CUDA graphs depend on and coupling graph lifetime across the two models. Key on the owning quant config instead. The draft is built with its own Exl3Config while every layer of one model shares a single config, so each model gets exactly one runtime. This is deliberately coarser than per-layer ownership: the prefill arena alone is ~1 GiB, so a per-layer cache would cost tens of GiB per rank on a 75+ layer model. Adds CPU regression tests asserting that two same-shape layers owned by different models produce distinct, stable scopes and cache keys that differ only by that scope.
Three defects in the rank-sliced MoE runtime cache, found while investigating issue vllm-project#183 (EXL3 + MTP corruption reports). 1. Target/draft isolation was defeatable. _runtime_scope_id() derives identity from the Exl3Config object, which only separates a target MoE layer from a rank-sliced MTP draft layer if the draft model file builds its own config. deepseek_mtp does, but several siblings (glm4_moe_mtp, qwen3_next_mtp, deepseek_eagle, ...) pass vllm_config.quant_config straight through, so the draft reuses the target's mutable Trellis/prefill scratch -- the exact shared-scratch corruption this scoping exists to prevent. Scope is now keyed on (config scope, is_draft), with the role derived from the layer prefix, so isolation no longer depends on how a model file happens to construct configs. 2. The runtime cache key was batch-dependent: max_batched_tokens folded in x.shape[0], so any m above the planned capacity produced a cache MISS and silently planned a fresh ~1 GiB arena mid-serve, which also made the capacity guard in _apply_rank_sliced unreachable. Capacity is a property of the layer, not of one forward pass. 3. max_num_batched_tokens silently defaulted to 4096 when scheduler_config was unavailable, which can put target and draft on different plans with no error -- a mismatch that only manifests at scale. Now raises. Note on vllm-project#183 specifically: the reported mechanism (partial-block masking under VLLM_EXL3_TRELLIS_MIN_M=1) did not reproduce. The fused Trellis MoE is bitwise correct at m=1,2,3 at production geometry (tile 64x256x64x256, BLOCK_M=8, capacity 32, topk=8) with the scratch arena NaN-poisoned, and needle-in-a- haystack passes 15/15 to ~200k tokens on TP4/DCP4 with MIN_M=1 and MTP-3. Also worth stating: MIN_M=1 widens the Trellis window so the draft's m=1..3 GEMMs stay on the fused path and remain graph-capturable; it is not a per-token capture and not a throughput tax.
…ut a workaround Fixes the boot failure reported in vllm-project#183. With the Trellis window at its historical default of 4, CUDA-graph capture of an EXL3 rank-sliced MTP draft reaches the eager parity path at m=1,2,3 -- illegal during capture -- and the engine cannot start: RuntimeError: EXL3 eager parity path entered during CUDA graph capture (m=3); capture sizes must lie inside the Trellis window [4, 32] fired from both profile_cudagraph_memory() and the real capture. The reporter established it is invariant to num_speculative_tokens (1 -> m=2, 3 -> m=3, 4 -> m=3) and to cudagraph_capture_sizes ([4,8,..,32] and [8,16,24,32] alike), because m here is the draft's row count per step, not a target batch size. Until now the only remedy was asking every operator to set VLLM_EXL3_TRELLIS_MIN_M=1 by hand. That is a workaround, not a fix: it is undiscoverable from the error, and anyone running a rank-sliced draft without it simply cannot boot. A backend should satisfy its own capture contract. This exports MIN_CAPTURABLE_TRELLIS_M -- the smallest row count the Trellis path can service, and therefore the smallest m an EXL3 rank-sliced MoE layer can be captured at -- and defaults draft layers to it. Target layers keep the previous default, and an explicit VLLM_EXL3_TRELLIS_MIN_M still wins in both directions so the kill switch is preserved. The constant is module-level so a capture-size selector can align its sizes with the backend rather than discovering the mismatch at capture time. The capture-time error, if it is ever still reached, now names the layer, says whether it was classified as draft or target, and states the value to set. Numerical safety of the widened window was established separately: the fused Trellis MoE is bitwise correct at m=1,2,3 at production geometry with the scratch arena NaN-poisoned, capture and replay below plan capacity match eager bit-for-bit (SparkInfer vllm-project#49), and needle-in-a-haystack recovers 42/42 to ~480k tokens across both nvfp4_ds_mla and fp8 KV.
… with no workaround Completes the vllm-project#183 boot-failure fix. The previous commit exported MIN_CAPTURABLE_TRELLIS_M and defaulted draft layers' Trellis window to it, but the role check it relied on could not actually classify a GLM-5.2-style MTP head, for two reasons established by direct testing: 1. The head is an extra decoder layer named exactly like a target layer -- with num_hidden_layers=78 the draft MoE is 'model.layers.78.mlp.experts' -- so no name substring can identify it. 2. Index-vs-num_hidden_layers comparison fails at plan/capture time because set_current_vllm_config is construction-scoped: by the time the quant method runs, get_current_vllm_config_or_none() returns None, the index branch is silently skipped, and the layer classifies as target. The boot error then fires exactly as reported. The role is therefore stamped at construction, where the config context is provably live: create_weights sets layer.exl3_is_draft from model_config.runner_type == "draft", the first-class signal minted for every model-backed speculator draft (SpeculativeConfig.__post_init__ -> ModelConfig(runner="draft")). This also covers speculators that bypass load_eagle_model (dspark, dflash). _is_draft_layer consults the stamp first; the name heuristic remains only as a fallback for layers built outside a speculator context. An immutable stamp also keeps the runtime-cache owner key phase-stable, which a phase-dependent heuristic would not. Also hardens _positive_env_int: an env var that is present but blank now means "unset" instead of aborting startup with a bare int('') ValueError -- compose and Kubernetes both render unset variables as empty strings -- and a non-integer value now gets an error naming the variable. Validated on 4x RTX PRO 6000 (SM120), TP4/DCP4, GLM-5.2-EXL3-TR3 with the tr3 MTP-78 draft, MTP-3: with VLLM_EXL3_TRELLIS_MIN_M entirely unset the engine boots, captures, and serves (previously: guaranteed startup failure). Zero occurrences of the capture-time error in the logs; inference verified.
Startup was flagging every EXL3 knob as unknown: WARNING [envs.py] Unknown vLLM environment variable detected: VLLM_EXL3_TRELLIS_MIN_M WARNING [envs.py] Unknown vLLM environment variable detected: VLLM_EXL3_EXT_PATH ... (9 in total) The variables are read directly from os.environ by the EXL3 quantization backend, so they take effect despite the warning -- but the spray reads as "your configuration is being ignored" and just sent a deployment down a false debugging trail. envs.py also has a strict mode where unknown VLLM_* variables raise instead of warn, which would turn this cosmetic noise into a hard startup failure. Registers all nine as passthrough entries (consuming code applies its own context-dependent defaults, e.g. the draft-vs-target Trellis window minimum). The remaining warned variables on a full deployment (VLLM_PCIE_DMA_FP8, VLLM_RTX6K_*, VLLM_CPP_AR_*, VLLM_CACHE_DIR, VLLM_B12X_MLA_SPEC_*) belong to the base runtime, not this PR, and are left to their owners.
d35dd8f to
89bfbb8
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
vllm/v1/attention/backends/mla/indexer.py (1)
293-293: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winUse a Google-style docstring for this parameterized helper.
Add
Args:andReturns:sections for the four inputs and computed block count.Proposed fix
- """Size indexer metadata for the cache group's effective DCP layout.""" + """Size indexer metadata for the cache group's effective DCP layout. + + Args: + max_model_len: Maximum model context length. + block_size: KV cache block size. + configured_cp_world_size: Configured context-parallel world size. + dcp_replicated: Whether the DCP cache is fully replicated. + dcp_kv_shard_count: Optional number of unique KV cache shards. + + Returns: + Maximum number of cache blocks required per request. + """As per coding guidelines, Python docstrings must use Google-style
Args:andReturns:sections.🤖 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/v1/attention/backends/mla/indexer.py` at line 293, Update the docstring for the parameterized helper near “Size indexer metadata for the cache group's effective DCP layout” to use Google style, adding an Args section documenting all four inputs and a Returns section documenting the computed block count. Preserve the helper’s behavior and existing summary.Source: Coding guidelines
🧹 Nitpick comments (3)
tests/quantization/test_exl3_prefill_plan.py (2)
255-272: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
pytest.raisesovertry/except/else.The other tests in this file already use
pytest.raises(...)with amatch=; using it here keeps failure messages consistent.🤖 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/quantization/test_exl3_prefill_plan.py` around lines 255 - 272, The test_parity_path_guarded_against_capture test should use pytest.raises with a match assertion for the expected RuntimeError instead of manual try/except/else handling. Wrap the _apply(method, layer, 2) call in pytest.raises and match the existing “capture” message text.
275-282: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDefine
planned_capson_Harnessdirectly.Monkey-patching the method onto the class at import time works only because attribute lookup is deferred; a plain method keeps the harness self-contained and greppable.
♻️ Proposed refactor
-def _install_planned_caps_helper(): - def planned_caps(self): - return self.api.planned - - _Harness.planned_caps = planned_caps - - -_install_planned_caps_helper()add inside
_Harness:def planned_caps(self): return self.api.planned🤖 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/quantization/test_exl3_prefill_plan.py` around lines 275 - 282, Move the planned_caps method definition directly into the _Harness class, returning self.api.planned. Remove the _install_planned_caps_helper function and its invocation so the harness no longer relies on import-time monkey-patching.vllm/model_executor/layers/quantization/exl3.py (1)
1247-1249: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead trailing
return.The
ifblock already returns; the bare fallthrough makes the earlyreturnat Line 1249 pointless.♻️ Tidy-up
if self.quant_config.rank_sliced_metadata is not None: self._prepare_rank_sliced_weights(layer) - return🤖 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/exl3.py` around lines 1247 - 1249, Remove the redundant trailing return in the rank-sliced metadata branch of the relevant layer method, leaving _prepare_rank_sliced_weights(layer) as the final operation in that conditional while preserving the existing fallthrough behavior.
🤖 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 `@vllm/envs.py`:
- Around line 2296-2297: Remove the duplicate VLLM_NVFP4_MLA_SCALES_FILE entry
near the calibrated MLA scales configuration, preserving the earlier
registration that strips the environment value and defaults to an empty string.
Ensure the dictionary contains only one key so the declared string behavior and
Ruff validation remain intact.
---
Outside diff comments:
In `@vllm/v1/attention/backends/mla/indexer.py`:
- Line 293: Update the docstring for the parameterized helper near “Size indexer
metadata for the cache group's effective DCP layout” to use Google style, adding
an Args section documenting all four inputs and a Returns section documenting
the computed block count. Preserve the helper’s behavior and existing summary.
---
Nitpick comments:
In `@tests/quantization/test_exl3_prefill_plan.py`:
- Around line 255-272: The test_parity_path_guarded_against_capture test should
use pytest.raises with a match assertion for the expected RuntimeError instead
of manual try/except/else handling. Wrap the _apply(method, layer, 2) call in
pytest.raises and match the existing “capture” message text.
- Around line 275-282: Move the planned_caps method definition directly into the
_Harness class, returning self.api.planned. Remove the
_install_planned_caps_helper function and its invocation so the harness no
longer relies on import-time monkey-patching.
In `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 1247-1249: Remove the redundant trailing return in the rank-sliced
metadata branch of the relevant layer method, leaving
_prepare_rank_sliced_weights(layer) as the final operation in that conditional
while preserving the existing fallthrough behavior.
🪄 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: dbee0fbc-2141-4f06-85de-54bdded08d15
📒 Files selected for processing (18)
tests/model_executor/layers/test_sparse_attn_indexer_b12x.pytests/model_executor/test_eagle_quantization.pytests/quantization/test_exl3.pytests/quantization/test_exl3_prefill_plan.pytests/v1/structured_output/test_reasoning_structured_output.pyvllm/config/model.pyvllm/envs.pyvllm/model_executor/layers/fused_moe/routed_experts.pyvllm/model_executor/layers/quantization/__init__.pyvllm/model_executor/layers/quantization/exl3.pyvllm/model_executor/models/deepseek_mtp.pyvllm/model_executor/models/deepseek_v2.pyvllm/model_executor/models/glm4_moe.pyvllm/model_executor/models/utils.pyvllm/models/deepseek_v32/nvidia/attention.pyvllm/v1/attention/backends/mla/indexer.pyvllm/v1/core/sched/scheduler.pyvllm/v1/structured_output/__init__.py
🚧 Files skipped from review as they are similar to previous changes (1)
- vllm/config/model.py
|
Full-checkpoint E2E validation with the unified SparkInfer API is complete.
SparkInfer implementation and full details: local-inference-lab/sparkinfer#90 |
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/quantization/test_exl3.py (1)
293-308: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the
monkeypatchfixture in Google-style format.The modified docstring describes a function with an argument but lacks an
Args:section.Proposed fix
Target profiling can also produce m=3, so role-dependent defaults merely move the same failure from the draft to MTP0. The backend declares one capability floor and uses it for both roles. + + Args: + monkeypatch: Pytest fixture used to control environment variables. """As per coding guidelines, Python code must use Google-style docstrings with
Args:/Returns:/Raises:sections.🤖 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/quantization/test_exl3.py` around lines 293 - 308, Update the docstring for test_rank_sliced_window_defaults_to_min_capturable_m to add a Google-style Args: section documenting the monkeypatch fixture, while preserving the existing regression description.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@tests/quantization/test_exl3.py`:
- Around line 293-308: Update the docstring for
test_rank_sliced_window_defaults_to_min_capturable_m to add a Google-style Args:
section documenting the monkeypatch fixture, while preserving the existing
regression description.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 1067cd97-6b4e-46c0-b8f8-5ed2db30c784
📒 Files selected for processing (5)
tests/model_executor/layers/test_mla_cache_format.pytests/quantization/test_exl3.pytests/quantization/test_exl3_prefill_plan.pyvllm/envs.pyvllm/model_executor/layers/quantization/exl3.py
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
vllm/model_executor/layers/quantization/exl3.py (3)
1693-1693: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePass
strict=Truewhen zipping tier bits with tier IDs. Both sites pair the bitrate tuple with the tier-ID tuple and silently truncate on a length mismatch; ruff flags both under B905.
vllm/model_executor/layers/quantization/exl3.py#L1693-L1693:zip(tiers, tier_ids, strict=True).vllm/model_executor/layers/quantization/exl3.py#L1822-L1822:zip(mixed["tier_bits"], mixed["tier_ids"], strict=True).🤖 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/exl3.py` at line 1693, Update both zip calls in vllm/model_executor/layers/quantization/exl3.py at lines 1693-1693 and 1822-1822 to pass strict=True, including the calls pairing tiers with tier_ids and mixed["tier_bits"] with mixed["tier_ids"], so length mismatches are detected instead of silently truncated.Source: Linters/SAST tools
1857-1867: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueName the mixed-path route block size.
The
8at Line 1857 and the(route_slots + 7) // 8at Line 1867 are the same route block size expressed twice; a named module constant (or an env override mirroringVLLM_EXL3_TRELLIS_BLOCK_M) keeps the packing andmax_m_blocksderivations from drifting apart.🤖 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/exl3.py` around lines 1857 - 1867, Define a named module-level constant for the mixed-path route block size and use it in both api.max_packed_route_slots and the max_m_blocks ceiling calculation within compile_mixed_trellis setup. If the module already supports an environment override pattern, mirror it for this constant; otherwise keep the existing value as the default.
216-239: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAvoid mutating the imported SparkInfer module; wrap it in a local facade.
Lines 235-236 write
prepare_weights/max_packed_route_slotsonto the third-party module object, so the patch is visible to every other importer ofsparkinfer.moe._shared.kernels.w4a16.mixed_trellisand will silently shadow same-named symbols if upstream ever adds them.♻️ Facade instead of module mutation
- module.prepare_weights = prepare.prepare_trellis256_moe_weights - module.max_packed_route_slots = host.max_packed_route_slots - _SPARKINFER_MIXED_TRELLIS_API = module - return module + api = SimpleNamespace( + compile_mixed_trellis=module.compile_mixed_trellis, + make_mixed_trellis_buffers=module.make_mixed_trellis_buffers, + run_mixed_trellis=module.run_mixed_trellis, + build_tiered_maps=module.build_tiered_maps, + combine_trellis_rotations=module.combine_trellis_rotations, + prepare_weights=prepare.prepare_trellis256_moe_weights, + max_packed_route_slots=host.max_packed_route_slots, + ) + _SPARKINFER_MIXED_TRELLIS_API = api + return api🤖 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/exl3.py` around lines 216 - 239, Update _load_sparkinfer_mixed_trellis to return a local facade that exposes the imported mixed_trellis implementation together with prepare_weights and max_packed_route_slots, rather than assigning those attributes onto the third-party module. Preserve lazy caching through _SPARKINFER_MIXED_TRELLIS_API and ensure callers can access the same required symbols through the facade.tests/quantization/test_exl3.py (1)
356-362: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the per-tier slab geometry, not just tier counts.
trellis_bits/num_experts/tier_idsall pass even if the wrong experts' tensors were stacked into a tier, since both tiers have two experts here. Checking the stacked shapes (last dim16 * bits) pins the partitioning that_prepare_mixed_rank_sliced_weightsis responsible for.💚 Additional assertions
assert [entry["trellis_bits"] for entry in api.prepared] == [3, 4] assert [entry["num_experts"] for entry in api.prepared] == [2, 2] + for entry in api.prepared: + bits = entry["trellis_bits"] + assert tuple(entry["w13"].shape) == ( + 2, + 2, + hidden // 16, + intermediate // 16, + 16 * bits, + ) + assert tuple(entry["w2"].shape) == ( + 2, + intermediate // 16, + hidden // 16, + 16 * bits, + ) + assert layer.exl3_mixed_trellis["tier_bits"] == (3, 4)🤖 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/quantization/test_exl3.py` around lines 356 - 362, Extend the mixed-rank preparation assertions around _prepare_mixed_rank_sliced_weights to verify each tier’s stacked slab geometry, including the final dimension being 16 multiplied by that tier’s trellis bits. Keep the existing trellis_bits, num_experts, tile_config, and tier_ids assertions, and assert shapes that distinguish which experts were assigned to each tier.
🤖 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 `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 1898-1906: Update the buffer accounting loop in the plan path to
process only tensor-valued attributes from vars(state["buffers"]).values()
before calling untyped_storage(). Ignore scalar, None, list, and other
non-tensor fields while preserving the existing deduplicated storage-byte
calculation for tensors.
---
Nitpick comments:
In `@tests/quantization/test_exl3.py`:
- Around line 356-362: Extend the mixed-rank preparation assertions around
_prepare_mixed_rank_sliced_weights to verify each tier’s stacked slab geometry,
including the final dimension being 16 multiplied by that tier’s trellis bits.
Keep the existing trellis_bits, num_experts, tile_config, and tier_ids
assertions, and assert shapes that distinguish which experts were assigned to
each tier.
In `@vllm/model_executor/layers/quantization/exl3.py`:
- Line 1693: Update both zip calls in
vllm/model_executor/layers/quantization/exl3.py at lines 1693-1693 and 1822-1822
to pass strict=True, including the calls pairing tiers with tier_ids and
mixed["tier_bits"] with mixed["tier_ids"], so length mismatches are detected
instead of silently truncated.
- Around line 1857-1867: Define a named module-level constant for the mixed-path
route block size and use it in both api.max_packed_route_slots and the
max_m_blocks ceiling calculation within compile_mixed_trellis setup. If the
module already supports an environment override pattern, mirror it for this
constant; otherwise keep the existing value as the default.
- Around line 216-239: Update _load_sparkinfer_mixed_trellis to return a local
facade that exposes the imported mixed_trellis implementation together with
prepare_weights and max_packed_route_slots, rather than assigning those
attributes onto the third-party module. Preserve lazy caching through
_SPARKINFER_MIXED_TRELLIS_API and ensure callers can access the same required
symbols through the facade.
🪄 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: a9ff4a2c-5827-42e5-8bed-cf2111d9c04a
📒 Files selected for processing (2)
tests/quantization/test_exl3.pyvllm/model_executor/layers/quantization/exl3.py
|
@coderabbitai review |
✅ Action performedReview finished.
|
2d69f95
into
local-inference-lab:dev/gilded-gnosis
Summary
Add native EXL3 checkpoint loading and execution without reconstructing or requantizing routed-expert weights at load time:
--quantization exl3configuration and packed tensor loadingThis supersedes #139 with the implementation replayed directly on the current
dev/gilded-gnosishead. Original commit authorship is preserved.Unified SparkInfer API
The homogeneous consumer uses SparkInfer #90 rather than the separate
trellis_moeAPI from superseded SparkInfer #49:fused_moe.plan_weights(source_format="exl3_trellis_mcg", quant_modes="w4a16", ...)fused_moe.prepare_weights(...)Caps/plan/ caller-ownedscratch_specs()/bind/runThis gives EXL3 Trellis and NF3 the same prepared-weight, scratch-arena, activation-dispatch, cache, and CUDA graph lifecycle. The consumer preserves the checkpoint's BF16 or FP16 parameter dtype rather than hard-coding BF16.
Mixed K3/K4 checkpoints use the paired internal API from SparkInfer #104. vLLM interprets the checkpoint metadata and creates two native packed weight tiers; SparkInfer packs global routes once, rotates activations once, and dispatches each CTA to its native bitrate inside one cooperative FC1/activation/FC2 grid. No BF16 reconstruction or requantization is performed.
Clean-up from #139
dev/gilded-gnosis, not an integration/build branch.Dependencies
Validation
Pinned CUDA 13.2 / PyTorch 2.12 / SM120 runtime:
18 passedafter the mixed integration update4 passed89 passedM=3072willfalco/GLM-5.2-EXL3-TR3-3.25bpwcheckpoint passed TP4/DCP4 with MTP0 and MTP3, CUDA graph capture, eight concurrent requests, and a 65,536-token prompt plus 256 generated tokensMixed-checkpoint performance:
Profiler validation reduced the mixed decode kernel from
82.38to54.11 us/layerwithout adding a host synchronization or copy. The safe GLM tile geometry reduced production-shape relative L2 error from8.32e-3to approximately4.1e-8.Ruff,
py_compile, andgit diff --checkpass for all mixed-path changes.Review notes
The largest review surface is
vllm/model_executor/layers/quantization/exl3.py. The shared GLM sparse-attention changes are replayed on current GG rather than copied from an older full-file overlay.No canonical branch is modified by this PR.
Summary by CodeRabbit
New Features
exl3quantization method, including rank-sliced EXL3 draft/target metadata handling.Bug Fixes
1when unset) and improved related env override behavior.