[GG] Add EXL3 Trellis backend for rank-sliced MoE checkpoints - #190
Open
voipmonitor wants to merge 19 commits into
Open
[GG] Add EXL3 Trellis backend for rank-sliced MoE checkpoints#190voipmonitor wants to merge 19 commits into
voipmonitor wants to merge 19 commits into
Conversation
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.
📝 WalkthroughWalkthroughChangesEXL3 rank-sliced quantization
Speculative structured-output advancement
Attention indexer guards
Estimated code review effort: 5 (Critical) | ~90+ minutes 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
Comment on lines
+1504
to
+1513
| min_trellis_m = _positive_env_int("VLLM_EXL3_TRELLIS_MIN_M", default_min_m) | ||
| max_trellis_m = _positive_env_int("VLLM_EXL3_TRELLIS_MAX_M", 32) | ||
| block_m = _positive_env_int("VLLM_EXL3_TRELLIS_BLOCK_M", 8) | ||
| chunk = _positive_env_int("VLLM_EXL3_PREFILL_CHUNK", 128) | ||
| prefill_trellis = os.environ.get("VLLM_EXL3_PREFILL_TRELLIS", "1") == "1" | ||
| prefill_block_m = _positive_env_int("VLLM_EXL3_PREFILL_BLOCK_M", 64) | ||
| if min_trellis_m > max_trellis_m: | ||
| raise ValueError( | ||
| "VLLM_EXL3_TRELLIS_MIN_M cannot exceed VLLM_EXL3_TRELLIS_MAX_M" | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Validate the parity capacity at plan time, not on the first oversized batch.
parity_rows is min(chunk, max_batched_tokens) once a prefill plan exists (Line 1606), while the parity path serves every m < min_trellis_m. With a raised VLLM_EXL3_TRELLIS_MIN_M and/or a lowered VLLM_EXL3_PREFILL_CHUNK, min_trellis_m - 1 > parity_rows is reachable and the failure surfaces mid-serve as the ValueError at Line 1759 instead of at startup, next to the existing min/max check.
🛡️ Proposed plan-time check
if min_trellis_m > max_trellis_m:
raise ValueError(
"VLLM_EXL3_TRELLIS_MIN_M cannot exceed VLLM_EXL3_TRELLIS_MAX_M"
)
+ if min_trellis_m - 1 > chunk:
+ raise ValueError(
+ "VLLM_EXL3_PREFILL_CHUNK must cover the parity window: "
+ f"chunk={chunk}, parity rows up to {min_trellis_m - 1}"
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| min_trellis_m = _positive_env_int("VLLM_EXL3_TRELLIS_MIN_M", default_min_m) | |
| max_trellis_m = _positive_env_int("VLLM_EXL3_TRELLIS_MAX_M", 32) | |
| block_m = _positive_env_int("VLLM_EXL3_TRELLIS_BLOCK_M", 8) | |
| chunk = _positive_env_int("VLLM_EXL3_PREFILL_CHUNK", 128) | |
| prefill_trellis = os.environ.get("VLLM_EXL3_PREFILL_TRELLIS", "1") == "1" | |
| prefill_block_m = _positive_env_int("VLLM_EXL3_PREFILL_BLOCK_M", 64) | |
| if min_trellis_m > max_trellis_m: | |
| raise ValueError( | |
| "VLLM_EXL3_TRELLIS_MIN_M cannot exceed VLLM_EXL3_TRELLIS_MAX_M" | |
| ) | |
| min_trellis_m = _positive_env_int("VLLM_EXL3_TRELLIS_MIN_M", default_min_m) | |
| max_trellis_m = _positive_env_int("VLLM_EXL3_TRELLIS_MAX_M", 32) | |
| block_m = _positive_env_int("VLLM_EXL3_TRELLIS_BLOCK_M", 8) | |
| chunk = _positive_env_int("VLLM_EXL3_PREFILL_CHUNK", 128) | |
| prefill_trellis = os.environ.get("VLLM_EXL3_PREFILL_TRELLIS", "1") == "1" | |
| prefill_block_m = _positive_env_int("VLLM_EXL3_PREFILL_BLOCK_M", 64) | |
| if min_trellis_m > max_trellis_m: | |
| raise ValueError( | |
| "VLLM_EXL3_TRELLIS_MIN_M cannot exceed VLLM_EXL3_TRELLIS_MAX_M" | |
| ) | |
| if min_trellis_m - 1 > chunk: | |
| raise ValueError( | |
| "VLLM_EXL3_PREFILL_CHUNK must cover the parity window: " | |
| f"chunk={chunk}, parity rows up to {min_trellis_m - 1}" | |
| ) |
🤖 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 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.
Comment on lines
+525
to
+528
| rank_sliced_name = getattr( | ||
| self.quant_config, | ||
| "normalize_rank_sliced_weight_name", | ||
| None, |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Persist quant_config before using it.
Line 526 accesses an attribute never assigned in __init__; GLM weight loading now raises AttributeError. Add self.quant_config = quant_config alongside self.config = config.
Proposed fix
self.config = config
+ self.quant_config = quant_config📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| rank_sliced_name = getattr( | |
| self.quant_config, | |
| "normalize_rank_sliced_weight_name", | |
| None, | |
| self.quant_config = quant_config | |
| rank_sliced_name = getattr( | |
| self.quant_config, | |
| "normalize_rank_sliced_weight_name", | |
| None, |
🤖 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/models/glm4_moe.py` around lines 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.
Comment on lines
+813
to
+830
| if ( | ||
| quant_config is not None | ||
| and getattr(quant_config, "get_name", lambda: None)() == "exl3" | ||
| ): | ||
| draft_hf = getattr(draft_model_config, "hf_config", None) | ||
| target_hf = getattr( | ||
| getattr(vllm_config, "model_config", None), "hf_config", None | ||
| ) | ||
| hf_config = draft_hf | ||
| if ( | ||
| getattr(draft_hf, "hybrid_tr3_tail", None) is None | ||
| and getattr(target_hf, "hybrid_tr3_tail", None) is not None | ||
| ): | ||
| hf_config = target_hf | ||
| quant_config.maybe_update_config( | ||
| draft_model_config.model, | ||
| hf_config=hf_config, | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Scope the hydration call to the rank-sliced case.
The comment describes rank-sliced (hybrid_tr3_tail) hydration, but the branch fires for any EXL3 draft. When neither config carries hybrid_tr3_tail, Exl3Config.maybe_update_config takes the dense path: a remote quantization_config.json fetch during draft construction, raising ValueError when tensor_storage is empty. That is a new hard failure/network call for dense EXL3 drafts that previously skipped this entirely.
🛡️ Proposed guard
draft_hf = getattr(draft_model_config, "hf_config", None)
target_hf = getattr(
getattr(vllm_config, "model_config", None), "hf_config", None
)
- hf_config = draft_hf
- if (
- getattr(draft_hf, "hybrid_tr3_tail", None) is None
- and getattr(target_hf, "hybrid_tr3_tail", None) is not None
- ):
- hf_config = target_hf
- quant_config.maybe_update_config(
- draft_model_config.model,
- hf_config=hf_config,
- )
+ hf_config = draft_hf
+ if getattr(draft_hf, "hybrid_tr3_tail", None) is None:
+ hf_config = target_hf
+ if getattr(hf_config, "hybrid_tr3_tail", None) is not None:
+ quant_config.maybe_update_config(
+ draft_model_config.model,
+ hf_config=hf_config,
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if ( | |
| quant_config is not None | |
| and getattr(quant_config, "get_name", lambda: None)() == "exl3" | |
| ): | |
| draft_hf = getattr(draft_model_config, "hf_config", None) | |
| target_hf = getattr( | |
| getattr(vllm_config, "model_config", None), "hf_config", None | |
| ) | |
| hf_config = draft_hf | |
| if ( | |
| getattr(draft_hf, "hybrid_tr3_tail", None) is None | |
| and getattr(target_hf, "hybrid_tr3_tail", None) is not None | |
| ): | |
| hf_config = target_hf | |
| quant_config.maybe_update_config( | |
| draft_model_config.model, | |
| hf_config=hf_config, | |
| ) | |
| if ( | |
| quant_config is not None | |
| and getattr(quant_config, "get_name", lambda: None)() == "exl3" | |
| ): | |
| draft_hf = getattr(draft_model_config, "hf_config", None) | |
| target_hf = getattr( | |
| getattr(vllm_config, "model_config", None), "hf_config", None | |
| ) | |
| hf_config = draft_hf | |
| if getattr(draft_hf, "hybrid_tr3_tail", None) is None: | |
| hf_config = target_hf | |
| if getattr(hf_config, "hybrid_tr3_tail", None) is not None: | |
| quant_config.maybe_update_config( | |
| draft_model_config.model, | |
| hf_config=hf_config, | |
| ) |
🤖 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/models/utils.py` around lines 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.
Comment on lines
+413
to
+429
| # The B12X / sparkinfer sparse indexer handles native next_n>2 on SM120 | ||
| # directly (see sparse_attn_indexer warmup q_rows 1, 2, 4), so it must | ||
| # NOT be forced onto the DeepGEMM next_n<=2 flatten fallback. Flattening | ||
| # MTP-2/MTP-3 verification into rank-1 rows produced subtly wrong accepted | ||
| # tokens (code-gen syntax errors); the native (B, next_n) path keeps MTP | ||
| # correct. Use the canonical backend-aware predicate so this also holds | ||
| # when B12X is selected via --attention-backend B12X_MLA_SPARSE with the | ||
| # VLLM_USE_B12X_SPARSE_INDEXER env var unset (it also asserts SM120). | ||
| from vllm.model_executor.layers.sparse_attn_indexer import ( | ||
| use_b12x_sparse_indexer, | ||
| ) | ||
|
|
||
| self.use_flattening = ( | ||
| not current_platform.is_device_capability_family(100) | ||
| and next_n not in (1, 2) | ||
| and not use_b12x_sparse_indexer() | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Use one B12X predicate for flattening and metadata construction.
use_b12x_sparse_indexer() disables flattening for backend-only selection, but the builder still gates B12X schedule metadata, active-width setup, and prefill chunking on envs.VLLM_USE_B12X_SPARSE_INDEXER (for example, later at Lines 977, 1156, and 1203). With --attention-backend B12X_MLA_SPARSE and the environment variable unset, native MTP is selected while the DeepGEMM metadata path remains active, risking incorrect decode scheduling. Cache the helper result and reuse it across all B12X branches.
🤖 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` around lines 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.
15 tasks
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Summary
This supersedes #139 with the same EXL3 implementation replayed directly on the current
dev/gilded-gnosishead. Original commit authorship is preserved.The backend adds native EXL3 checkpoint loading and execution without reconstructing or requantizing routed-expert weights at load time:
--quantization exl3configuration and packed tensor loadingClean-up from #139
dev/gilded-gnosishead, not an integration/build branch.docker/DockerfileFlashInfer-JIT change was dropped; the shared Blackwell image already builds FlashInfer from a pinned source revision.The required SparkInfer API/kernel dependency remains #49.
Validation
tests/quantization/test_exl3.pytests/quantization/test_exl3_prefill_plan.pyResult on the Gilded Gnosis CUDA 13.2 / PyTorch 2.12 runtime:
15 passed.SparkInfer #49 CPU contract tests:
6 passed, 1 skipped.The release image builds the EXL3 CUDA extension from the public, immutable SM120 source pin rather than copying the opaque prebuilt
.soand Torch ABI shim used by the earlier standalone image. Full model validation is intentionally delegated to the EXL3 community test round for the resulting v20 r5 candidate.Review notes
The largest review surface is
vllm/model_executor/layers/quantization/exl3.py. The only shared GLM sparse-attention change is the B12X flattening gate invllm/v1/attention/backends/mla/indexer.py; it has been replayed on the current GG implementation rather than copied from an older full-file overlay.No canonical branch is modified by this PR.
Summary by CodeRabbit
New Features
Bug Fixes
Tests