[GG] consolidate EXL3 runtime and prewarm mixed-Trellis routes - #228
[GG] consolidate EXL3 runtime and prewarm mixed-Trellis routes#228voipmonitor wants to merge 17 commits into
Conversation
📝 WalkthroughWalkthroughThe change generalizes online quantization overlays to EXL3, adds persistent encoding caches, extends shared-H rank-sliced MoE support, and adds mixed-bitrate decode/prefill planning with configurable capacity. ChangesEXL3 online quantization and runtime
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Exl3OnlineLinearMethod
participant Exl3OnlineCache
participant ExLlamaV3Encoder
participant SparkInfer
Exl3OnlineLinearMethod->>Exl3OnlineCache: Request cached online quantization
Exl3OnlineCache->>ExLlamaV3Encoder: Encode weights on cache miss
ExLlamaV3Encoder-->>Exl3OnlineCache: Return quantized tensors
Exl3OnlineCache-->>Exl3OnlineLinearMethod: Return cached or new tensors
Exl3OnlineLinearMethod->>SparkInfer: Execute Trellis GEMM
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 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: 7
🧹 Nitpick comments (7)
tests/quantization/test_exl3_online_cache.py (2)
94-106: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for
cache_modeparsing.The tests set
VLLM_EXL3_ONLINE_CACHE_MODEtoreadwriteandreadonlyonly. They do not coveroff, the aliases (read-only,rw,none), or theValueErrorfor an unrecognized value. Theoffmode is the operator escape hatch, so a regression there would silently keep writing cache files.🤖 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_online_cache.py` around lines 94 - 106, Add tests alongside test_readonly_miss_does_not_publish covering cache_mode parsing for off, read-only, rw, and none, plus an unrecognized value raising ValueError. Verify off disables cache reads and writes, while each alias maps to its documented mode behavior.
141-145: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAdd coverage for an unresolved Hub revision.
test_hub_model_identity_tracks_resolved_revisionalways passes an explicit revision. It does not exerciseresolve_model_identity("org/model")with no revision and nohf_config, which is the path that currently substitutes the literal"unresolved"and makes two different revisions share one identity. Add a case for it once the source behavior is settled.Do you want me to generate the test?
🤖 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_online_cache.py` around lines 141 - 145, Add coverage in test_hub_model_identity_tracks_resolved_revision for calling resolve_model_identity("org/model") without revision or hf_config, and assert the expected identity behavior after the source fix so unresolved revisions do not collapse distinct model identities into the literal "unresolved" value.vllm/model_executor/layers/quantization/exl3.py (4)
110-141: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider recording the block-32 policy inputs in a named constant set.
_resolve_mixed_trellis_prefill_block_mencodes one model's geometry as inline literals. The function is well documented and tested, so this is optional. If more qualified geometries appear, move the tuple(device_major, hidden_size, intermediate_size, tier_signature, topk, prefill_tile_config)into a module-level frozenset of qualified signatures.🤖 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 110 - 141, Optionally refactor _resolve_mixed_trellis_prefill_block_m to use a module-level frozenset of qualified policy signatures containing (device_major, hidden_size, intermediate_size, tier_signature, topk, prefill_tile_config), and check the current input tuple against it while preserving the explicit_override guard and existing fallback behavior.
213-217: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReport a clear error when the encoder module lacks
quantize_exl3.The function validates that
quantize.pyexists but not that it exportsquantize_exl3. An incompatible ExLlamaV3 revision then fails with a bareAttributeErrorfar from the configuration mistake. Every other failure in this loader raises a descriptiveRuntimeError.🛡️ Proposed fix
quantize = importlib.import_module( f"{package_name}.modules.quant.exl3_lib.quantize" ) + if not hasattr(quantize, "quantize_exl3"): + raise RuntimeError( + "VLLM_EXL3_ENCODER_SOURCE points to an incompatible ExLlamaV3 " + f"encoder: {package_root} has no quantize_exl3 entry point" + ) _EXL3_ONLINE_QUANTIZER = quantize.quantize_exl3 return _EXL3_ONLINE_QUANTIZER🤖 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 213 - 217, Update the loader that assigns _EXL3_ONLINE_QUANTIZER from the imported quantize module to validate that the module exports quantize_exl3 before accessing it. If the symbol is missing, raise a descriptive RuntimeError consistent with the loader’s other failure paths, then preserve the existing assignment and return behavior for compatible modules.
2397-2407: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_mixed_trellis_prefill_tile_configcurrently duplicates_mixed_trellis_tile_config.The function forwards both arguments unchanged. It adds an indirection with no behavior. Keep it only if a divergent prefill geometry is planned; the comment suggests that intent. Otherwise call
_mixed_trellis_tile_configdirectly at the two call sites.🤖 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 2397 - 2407, Remove the redundant _mixed_trellis_prefill_tile_config wrapper and update both callers to invoke Exl3MoEMethod._mixed_trellis_tile_config directly with the same arguments. Preserve the existing tile geometry and behavior unless divergent prefill geometry is intentionally required.
2731-2769: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winEvery mixed-bitrate forward re-reads environment variables and queries device properties.
_mixed_rank_sliced_runtimeruns on each_apply_mixed_rank_slicedcall, including decode steps. Before the_MIXED_TRELLIS_RUNTIMESlookup it performs fouros.environreads, atorch.cuda.get_device_propertiescall, and the block-policy resolution. A GLM-5.2 style model repeats this for every MoE layer on every step.These values are batch-invariant. Cache them on the layer, or compute the runtime key from a per-layer memoized settings object.
Line 2761 also re-reads
VLLM_EXL3_PREFILL_BLOCK_Monly to log the configured value that line 2735 already read intoprefill_block_mbefore the reassignment. Keep the first read in a separateconfigured_block_mvariable.♻️ Proposed fix for the duplicate read
prefill_block_raw = os.environ.get("VLLM_EXL3_PREFILL_BLOCK_M") - prefill_block_m = _positive_env_int("VLLM_EXL3_PREFILL_BLOCK_M", 64) + configured_block_m = _positive_env_int("VLLM_EXL3_PREFILL_BLOCK_M", 64) @@ prefill_block_m = _resolve_mixed_trellis_prefill_block_m( - configured_block_m=prefill_block_m, + configured_block_m=configured_block_m, @@ prefill_block_m, - _positive_env_int("VLLM_EXL3_PREFILL_BLOCK_M", 64), + configured_block_m,🤖 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 2731 - 2769, Cache the batch-invariant mixed Trellis settings used by _mixed_rank_sliced_runtime, including environment-derived limits, device properties, and resolved prefill block policy, on the layer or in a per-layer memoized settings object so repeated _apply_mixed_rank_sliced calls reuse them before the _MIXED_TRELLIS_RUNTIMES lookup. In the initialization flow, preserve the original VLLM_EXL3_PREFILL_BLOCK_M value in a separate configured_block_m variable, derive the effective prefill_block_m from it, and use configured_block_m in the log instead of reading the environment again.vllm/model_executor/layers/quantization/exl3_online_cache.py (1)
224-233: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueRead the cache file once instead of twice.
_loadopens the file withsafe_openfor metadata, then callsload_filewhich reopens and re-reads it.safe_opencan supply both the metadata and the tensors in one pass. This runs once per layer at startup, so the saving is modest.♻️ Proposed refactor
with safe_open(path, framework="pt", device="cpu") as handle: metadata = handle.metadata() or {} - if metadata.get("cache_key") != key.canonical_json(): - raise ValueError("online EXL3 cache key metadata does not match") - tensors = load_file(path, device="cpu") + if metadata.get("cache_key") != key.canonical_json(): + raise ValueError("online EXL3 cache key metadata does not match") + tensors = {name: handle.get_tensor(name) for name in handle.keys()} _validate_tensors(tensors, key)🤖 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_online_cache.py` around lines 224 - 233, Update _load to obtain both metadata and tensors from the existing safe_open handle, eliminating the separate load_file(path, device="cpu") call and second file read. Preserve the cache-key validation, tensor validation, proxy_error conversion, and Exl3OnlineCacheResult construction using the tensors loaded from that single handle.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/quantization/test_exl3_prefill_plan.py`:
- Around line 116-119: Update run_mixed_trellis to consume and record the tier
arguments from args[1:3] along with each input slice’s row count, while
preserving the existing routing capture and float32 return. Extend the mixed
prefill test assertions to require exactly two calls with row counts [128, 72]
and verify every call used the configured prefill_tiers rather than decode
tiers.
In `@tests/quantization/test_exl3.py`:
- Around line 818-824: Move the calls list from the class body of FakeMixedApi
into an __init__ method, initializing self.calls separately for each instance.
Keep run_mixed_trellis appending to self.calls so assertions continue to inspect
only the current instance’s calls.
In `@vllm/model_executor/layers/quantization/exl3_online_cache.py`:
- Around line 176-194: Document in the cache-directory configuration or helper
around cache_root that VLLM_EXL3_ONLINE_CACHE_DIR must be writable only by the
serving user and trusted at the same level as the model checkpoint. In the lock
creation logic near cache_path/load_or_quantize, remove the explicit 0o666 mode
and rely on the process umask so shared-cache lock files are not world-writable.
- Around line 301-302: Update the readonly branch in the cache lookup flow to
return the quantized result with path=None on a miss, rather than propagating
the requested path. Keep path populated only when a real cache entry exists.
- Around line 82-88: Update the unresolved Hub-revision branch in the model
identity logic to fail closed instead of hashing a payload containing the
literal "unresolved". Signal the unresolved identity through the existing
return/error mechanism so load_or_quantize can disable caching for that model,
while preserving normal hashing when a revision resolves.
In `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 1078-1095: Update the docstring for the method shown to document
both possible method outcomes: the existing MXFP8 method and
Exl3OnlineLinearMethod when VLLM_EXL3_ONLINE_TRELLIS_BITS is configured. Keep
the existing eligibility and ValueError documentation, and use Google-style
wording consistent with the current Args, Returns, and Raises sections.
- Around line 524-551: Update _sparkinfer_trellis_weight to cache prepared
weights by the owning trellis tensor rather than raw trellis.data_ptr() in the
process-global dictionary, using an inner key of (suh.data_ptr(),
svh.data_ptr(), dtype). Ensure the cache retains the trellis tensor reference so
entries remain tied to the source allocation and are released when that tensor
is freed, preventing stale address reuse and indefinite retention.
---
Nitpick comments:
In `@tests/quantization/test_exl3_online_cache.py`:
- Around line 94-106: Add tests alongside test_readonly_miss_does_not_publish
covering cache_mode parsing for off, read-only, rw, and none, plus an
unrecognized value raising ValueError. Verify off disables cache reads and
writes, while each alias maps to its documented mode behavior.
- Around line 141-145: Add coverage in
test_hub_model_identity_tracks_resolved_revision for calling
resolve_model_identity("org/model") without revision or hf_config, and assert
the expected identity behavior after the source fix so unresolved revisions do
not collapse distinct model identities into the literal "unresolved" value.
In `@vllm/model_executor/layers/quantization/exl3_online_cache.py`:
- Around line 224-233: Update _load to obtain both metadata and tensors from the
existing safe_open handle, eliminating the separate load_file(path,
device="cpu") call and second file read. Preserve the cache-key validation,
tensor validation, proxy_error conversion, and Exl3OnlineCacheResult
construction using the tensors loaded from that single handle.
In `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 110-141: Optionally refactor
_resolve_mixed_trellis_prefill_block_m to use a module-level frozenset of
qualified policy signatures containing (device_major, hidden_size,
intermediate_size, tier_signature, topk, prefill_tile_config), and check the
current input tuple against it while preserving the explicit_override guard and
existing fallback behavior.
- Around line 213-217: Update the loader that assigns _EXL3_ONLINE_QUANTIZER
from the imported quantize module to validate that the module exports
quantize_exl3 before accessing it. If the symbol is missing, raise a descriptive
RuntimeError consistent with the loader’s other failure paths, then preserve the
existing assignment and return behavior for compatible modules.
- Around line 2397-2407: Remove the redundant _mixed_trellis_prefill_tile_config
wrapper and update both callers to invoke
Exl3MoEMethod._mixed_trellis_tile_config directly with the same arguments.
Preserve the existing tile geometry and behavior unless divergent prefill
geometry is intentionally required.
- Around line 2731-2769: Cache the batch-invariant mixed Trellis settings used
by _mixed_rank_sliced_runtime, including environment-derived limits, device
properties, and resolved prefill block policy, on the layer or in a per-layer
memoized settings object so repeated _apply_mixed_rank_sliced calls reuse them
before the _MIXED_TRELLIS_RUNTIMES lookup. In the initialization flow, preserve
the original VLLM_EXL3_PREFILL_BLOCK_M value in a separate configured_block_m
variable, derive the effective prefill_block_m from it, and use
configured_block_m in the log instead of reading the environment again.
🪄 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: ed2b97fd-8330-4d2f-9025-d83403b66eb8
📒 Files selected for processing (9)
docs/features/quantization/online.mdtests/quantization/test_exl3.pytests/quantization/test_exl3_online_cache.pytests/quantization/test_exl3_prefill_plan.pytests/quantization/test_quantization_config_args.pyvllm/config/quantization.pyvllm/envs.pyvllm/model_executor/layers/quantization/exl3.pyvllm/model_executor/layers/quantization/exl3_online_cache.py
| if not model_path.exists(): | ||
| payload = { | ||
| "kind": "hub", | ||
| "model": model_name, | ||
| "revision": resolved_revision or "unresolved", | ||
| } | ||
| return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick win
An unresolved Hub revision produces a colliding cache identity.
When model_name is not a local path and neither revision nor hf_config._commit_hash resolves, the payload records the literal string "unresolved". Every revision of the same repository then hashes to the same model_identity. load_or_quantize treats the resulting file as a valid hit and returns trellis weights that were encoded from a different revision. The engine serves corrupted weights with no error.
Fail closed instead. Signal that the identity is unresolved and let the caller disable caching for that model.
🛡️ Proposed fix
resolved_revision = revision or getattr(hf_config, "_commit_hash", None)
model_path = Path(model_name).expanduser()
if not model_path.exists():
+ if not resolved_revision:
+ raise ValueError(
+ "online EXL3 caching requires a resolved revision for Hub "
+ f"checkpoint {model_name!r}; pass --revision or set "
+ "VLLM_EXL3_ONLINE_CACHE_MODE=off"
+ )
payload = {
"kind": "hub",
"model": model_name,
- "revision": resolved_revision or "unresolved",
+ "revision": resolved_revision,
}📝 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.
| if not model_path.exists(): | |
| payload = { | |
| "kind": "hub", | |
| "model": model_name, | |
| "revision": resolved_revision or "unresolved", | |
| } | |
| return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() | |
| if not model_path.exists(): | |
| if not resolved_revision: | |
| raise ValueError( | |
| "online EXL3 caching requires a resolved revision for Hub " | |
| f"checkpoint {model_name!r}; pass --revision or set " | |
| "VLLM_EXL3_ONLINE_CACHE_MODE=off" | |
| ) | |
| payload = { | |
| "kind": "hub", | |
| "model": model_name, | |
| "revision": resolved_revision, | |
| } | |
| return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() |
🧰 Tools
🪛 ast-grep (0.45.0)
[info] 87-87: use jsonify instead of json.dumps for JSON output
Context: json.dumps(payload, sort_keys=True)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🤖 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_online_cache.py` around lines 82
- 88, Update the unresolved Hub-revision branch in the model identity logic to
fail closed instead of hashing a payload containing the literal "unresolved".
Signal the unresolved identity through the existing return/error mechanism so
load_or_quantize can disable caching for that model, while preserving normal
hashing when a revision resolves.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
vllm/model_executor/layers/quantization/exl3.py (1)
540-554: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winThe cache key can alias after
id()reuse of the rotation tensors.The cache is now owned by
trellis, which fixes the previous address-reuse hazard for the trellis tensor. The residual hazard iskey = (id(suh), id(svh), dtype). The entry stores onlyweightand holds no reference tosuhorsvh. If a caller passes a temporary rotation tensor, CPython can free it and reuse the sameidfor a different tensor. A later lookup then returns a prepared weight that was built from the previous rotations, and the GEMM produces wrong results with no error.Store the source tensors in the entry so their identities stay valid for the life of the cache.
🛡️ Proposed fix
key = (id(suh), id(svh), dtype) - weight = cache.get(key) + entry = cache.get(key) + weight = None if entry is None else entry[2] if weight is None: weight = api.prepare_weight( trellis, suh, svh, codebook="mcg", params_dtype=dtype, ) - cache[key] = weight + # Retain the rotation tensors so their ``id`` cannot be reused by a + # different tensor while this entry is live. + cache[key] = (suh, svh, weight) return weight🤖 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 540 - 554, Update the cache entries in the trellis-owned cache around the key lookup and api.prepare_weight call to retain references to both suh and svh alongside the prepared weight, preventing their ids from being reused while cached. On cache hits, continue returning the stored weight, and on misses, store the weight together with the source rotation tensors.
🧹 Nitpick comments (1)
vllm/model_executor/layers/quantization/exl3.py (1)
197-219: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueFailed loads leave the synthetic packages registered in
sys.modules.The loop at lines 204-210 registers
_vllm_exl3_encoder*before validation. If the import fails orquantize_exl3is absent, the entries stay. A retry in the same process with a correctedVLLM_EXL3_ENCODER_SOURCEskips re-registration because of theif name in sys.modules: continueguard, so the stale__path__still points at the previous root.Remove the entries the function created when it raises.
♻️ Proposed refactor
- for name, path in packages.items(): - if name in sys.modules: - continue - module = ModuleType(name) - module.__path__ = [str(path)] - module.__package__ = name - sys.modules[name] = module - - quantize = importlib.import_module( - f"{package_name}.modules.quant.exl3_lib.quantize" - ) - if not hasattr(quantize, "quantize_exl3"): - raise RuntimeError( - "VLLM_EXL3_ENCODER_SOURCE points to an incompatible ExLlamaV3 " - "encoder: modules.quant.exl3_lib.quantize has no quantize_exl3" - ) + created: list[str] = [] + for name, path in packages.items(): + if name in sys.modules: + continue + module = ModuleType(name) + module.__path__ = [str(path)] + module.__package__ = name + sys.modules[name] = module + created.append(name) + + try: + quantize = importlib.import_module( + f"{package_name}.modules.quant.exl3_lib.quantize" + ) + if not hasattr(quantize, "quantize_exl3"): + raise RuntimeError( + "VLLM_EXL3_ENCODER_SOURCE points to an incompatible ExLlamaV3 " + "encoder: modules.quant.exl3_lib.quantize has no quantize_exl3" + ) + except Exception: + for name in created: + sys.modules.pop(name, None) + raise🤖 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 197 - 219, Update the synthetic package setup around the packages loop and quantize import to track the _vllm_exl3_encoder entries created by this load attempt, then remove those entries from sys.modules whenever importlib.import_module or the quantize_exl3 validation raises. Preserve pre-existing module entries and successful-load 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 `@tests/quantization/test_exl3.py`:
- Around line 229-233: Update the cleanup around
exl3_module._load_exl3_online_quantizer() to wrap the pytest.raises assertion in
a try/finally block, and remove all _vllm_exl3_encoder and _vllm_exl3_encoder.*
entries with sys.modules.pop directly in the finally block. Ensure cleanup runs
even when the assertion fails and does not use monkeypatch.delitem.
In `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 2407-2413: Update the tier preparation flow to reuse the already
prepared final mixed tier for prefill when _mixed_trellis_prefill_tile_config
produces the same configuration as _mixed_trellis_tile_config. In the relevant
prepare_tier caller, assign prefill_tiers from prepared_tiers[-1] instead of
invoking prepare_weights again, while preserving separate preparation when the
configurations differ.
---
Outside diff comments:
In `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 540-554: Update the cache entries in the trellis-owned cache
around the key lookup and api.prepare_weight call to retain references to both
suh and svh alongside the prepared weight, preventing their ids from being
reused while cached. On cache hits, continue returning the stored weight, and on
misses, store the weight together with the source rotation tensors.
---
Nitpick comments:
In `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 197-219: Update the synthetic package setup around the packages
loop and quantize import to track the _vllm_exl3_encoder entries created by this
load attempt, then remove those entries from sys.modules whenever
importlib.import_module or the quantize_exl3 validation raises. Preserve
pre-existing module entries and successful-load 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: d71184d7-26f8-4e43-8b73-45fdf9ac89d2
📒 Files selected for processing (5)
tests/quantization/test_exl3.pytests/quantization/test_exl3_online_cache.pytests/quantization/test_exl3_prefill_plan.pyvllm/model_executor/layers/quantization/exl3.pyvllm/model_executor/layers/quantization/exl3_online_cache.py
🚧 Files skipped from review as they are similar to previous changes (2)
- vllm/model_executor/layers/quantization/exl3_online_cache.py
- tests/quantization/test_exl3_prefill_plan.py
| with pytest.raises(RuntimeError, match="has no quantize_exl3"): | ||
| exl3_module._load_exl3_online_quantizer() | ||
| for name in tuple(exl3_module.sys.modules): | ||
| if name == "_vllm_exl3_encoder" or name.startswith("_vllm_exl3_encoder."): | ||
| monkeypatch.delitem(exl3_module.sys.modules, name, raising=False) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
pytest monkeypatch delitem restores value at teardown
💡 Result:
Yes, pytest's monkeypatch fixture restores the state of objects, dictionaries, and environment variables modified by its methods (including delitem) at the end of the test or fixture teardown [1][2][3]. However, it is important to understand exactly what "restoring state" means in this context: 1. Reversing Monkeypatch Operations: The monkeypatch fixture tracks the changes it performs and explicitly reverses them during teardown [1][3]. For monkeypatch.delitem(obj, name), pytest records that it deleted the specified key and restores it by re-adding the original value (if it existed) or ensuring the key is removed if that was the tracked action [1][4]. 2. Limitations: Monkeypatch does not track or "undo" arbitrary modifications made to objects outside of the monkeypatch API [5]. If your test code modifies an object directly (e.g., manually deleting a key from a dictionary) rather than using monkeypatch.delitem, the monkeypatch fixture will not know about that change and will not restore the object to its original state [5]. 3. Expected Behavior: The fixture is designed to undo exactly what it did [1][5]. If you need to ensure a key is gone regardless of whether it existed previously, or if you need to "pin" a state to restore later, you should rely on the documented behavior of monkeypatch methods and avoid manual modifications to the target objects [2][5]. For complex scenarios where you need granular control over the cleanup, you can use monkeypatch.context() to create a scope where patches are applied and automatically reverted upon exiting the block [2][3].
Citations:
- 1: https://docs.pytest.org/en/latest/_modules/_pytest/monkeypatch.html
- 2: https://docs.pytest.org/en/stable/how-to/monkeypatch.html
- 3: https://docs.pytest.org/en/latest/reference/reference.html?highlight=monkeypatch
- 4: https://github.com/pytest-dev/pytest/blob/main/src/_pytest/monkeypatch.py
- 5: Monkeypatch.delitem/delattr does not keep track of non-existing properties pytest-dev/pytest#14094
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Repository files matching test_quantization/test_exl3.py:"
git ls-files | grep -F 'tests/quantization/test_exl3.py' || true
if [ -f tests/quantization/test_exl3.py ]; then
echo
echo "Relevant test body (lines 200-245):"
sed -n '200,245p' tests/quantization/test_exl3.py | cat -n | sed 's/^/ /'
echo
echo "Search for _load_exl3_online_quantizer and encoder sys.modules handling:"
rg -n "_load_exl3_online_quantizer|_vllm_exl3_encoder|sys\.modules|importlib\.import_module" tests/quantization/test_exl3.py -C 3
fiRepository: local-inference-lab/vllm
Length of output: 3388
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Locate _load_exl3_online_quantizer implementation:"
rg -n "def _load_exl3_online_quantizer|_vllm_exl3_encoder|sys\.modules|importlib\.import_module" vllm tests/quantization/test_exl3.py -C 4
echo
echo "Candidate files containing implementation:"
python3 - <<'PY'
from pathlib import Path
for p in Path("vllm").rglob("*exl3*"):
if p.is_file():
print(p)
PYRepository: local-inference-lab/vllm
Length of output: 50380
Remove the synthetic encoder packages directly before assertion failure.
_load_exl3_online_quantizer registers _vllm_exl3_encoder* entries in sys.modules before raising. monkeypatch.delitem records the deletion and puts the entry back at teardown, so the synthetic module persists. If the pytest.raises checks fail, the cleanup code never runs. Use sys.modules.pop in a try/finally block.
💚 Proposed fix
- with pytest.raises(RuntimeError, match="has no quantize_exl3"):
- exl3_module._load_exl3_online_quantizer()
- for name in tuple(exl3_module.sys.modules):
- if name == "_vllm_exl3_encoder" or name.startswith("_vllm_exl3_encoder."):
- monkeypatch.delitem(exl3_module.sys.modules, name, raising=False)
+ try:
+ with pytest.raises(RuntimeError, match="has no quantize_exl3"):
+ exl3_module._load_exl3_online_quantizer()
+ finally:
+ for name in tuple(exl3_module.sys.modules):
+ if name == "_vllm_exl3_encoder" or name.startswith("_vllm_exl3_encoder."):
+ exl3_module.sys.modules.pop(name, None)📝 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.
| with pytest.raises(RuntimeError, match="has no quantize_exl3"): | |
| exl3_module._load_exl3_online_quantizer() | |
| for name in tuple(exl3_module.sys.modules): | |
| if name == "_vllm_exl3_encoder" or name.startswith("_vllm_exl3_encoder."): | |
| monkeypatch.delitem(exl3_module.sys.modules, name, raising=False) | |
| try: | |
| with pytest.raises(RuntimeError, match="has no quantize_exl3"): | |
| exl3_module._load_exl3_online_quantizer() | |
| finally: | |
| for name in tuple(exl3_module.sys.modules): | |
| if name == "_vllm_exl3_encoder" or name.startswith("_vllm_exl3_encoder."): | |
| exl3_module.sys.modules.pop(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 `@tests/quantization/test_exl3.py` around lines 229 - 233, Update the cleanup
around exl3_module._load_exl3_online_quantizer() to wrap the pytest.raises
assertion in a try/finally block, and remove all _vllm_exl3_encoder and
_vllm_exl3_encoder.* entries with sys.modules.pop directly in the finally block.
Ensure cleanup runs even when the assertion fails and does not use
monkeypatch.delitem.
| @staticmethod | ||
| def _mixed_trellis_prefill_tile_config(hidden_size: int, intermediate_size: int): | ||
| # Route packing stays block-64, while SparkInfer's mixed-kernel ABI v2 | ||
| # executes FC2 as arithmetic-equivalent 8-row subtiles. Reuse the | ||
| # checkpoint's original one-grid tile geometry so prefill has the same | ||
| # reduction order as the stock-r16 quality reference. | ||
| return Exl3MoEMethod._mixed_trellis_tile_config(hidden_size, intermediate_size) |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Locate the SparkInfer mixed-trellis prepare_weights contract.
set -euo pipefail
fd -t f -e py . | xargs rg -n -C 5 '\bdef prepare_weights\s*\(' || true
rg -n -C 5 '_load_sparkinfer_mixed_trellis' --type=pyRepository: local-inference-lab/vllm
Length of output: 2400
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== candidate files =="
fd -t f 'exl3\.py$|exl3.*' . | sed 's#^\./##'
echo
echo "== line counts =="
for f in $(fd -t f 'exl3\.py$' .); do
wc -l "$f"
done
echo
echo "== relevant exl3 symbols =="
ast-grep outline $(fd -t f 'exl3\.py$' . | head -n 1) --match '_mixed_trellis_prefill_tile_config' --view expanded || true
ast-grep outline $(fd -t f 'exl3\..py$' . | head -n 1) --match 'prepare_tier' --view expanded || trueRepository: local-inference-lab/vllm
Length of output: 619
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Use exact filename/path from fd output: vllm/model_executor/layers/quantization/exl3.py
f='vllm/model_executor/layers/quantization/exl3.py'
echo "== line counts for $f =="
wc -l "$f"
echo
echo "== symbol outlines =="
ast-grep outline "$f" --match '_mixed_trellis_prefill_tile_config' --view expanded || true
ast-grep outline "$f" --match 'prepare_tier' --view expanded || true
echo
echo "== relevant sections around tile configs and tier prep =="
sed -n '110,150p' "$f"
echo "----"
sed -n '2380,2425p' "$f"
echo "----"
sed -n '2550,2590p' "$f"Repository: local-inference-lab/vllm
Length of output: 6201
🏁 Script executed:
#!/bin/bash
set -euo pipefail
f='vllm/model_executor/layers/quantization/exl3.py'
echo "== imports and exl3_mixed_trellis usage candidates =="
rg -n -C 3 'exl3_mixed_trellis|prepared_tiers|prefill_tiers|prepare_tier|prefill_tile_config|mixed_tile_config' "$f"
echo
echo "== section 2425-2490 =="
sed -n '2425,2490p' "$f"
echo
echo "== section 2490-2578 =="
sed -n '2490,2578p' "$f"
echo
echo "== read-only structural verifier for prefill_vs_decode equality and duplicated calls =="
python3 - <<'PY'
from pathlib import Path
p = Path('vllm/model_executor/layers/quantization/exl3.py')
s = p.read_text()
checks = {
'prefill_delegates_to_mixed': 'return Exl3MoEMethod._mixed_trellis_tile_config(hidden_size, intermediate_size)' in s,
'prefill_tile_config_assigned_from_mixed': 'prefill_tile_config = Exl3MoEMethod._mixed_trellis_prefill_tile_config(hidden_size, intermediate_size)' in s,
'prepare_tier_def_contains_only_tile_config_fc1_fc2_tile_n_params': all(x in s for x in [
'fc1_tile_n=tile_config[1],',
'fc2_tile_n=tile_config[3],',
'mixed_api.prepare_weights',
]),
'prepare_tier_call_args_identical_except_tile_config': s.count('prepare_tier(mixed_tile_config)') > 0 and s.count('prepare_tier(prefill_tile_config)') > 0,
'prepared_and_prefill_lists_both_appended': 'prepared_tiers.append(prepare_tier(mixed_tile_config))' in s and 'prefill_tiers.append(prepare_tier(prefill_tile_config))' in s,
'comment_explain_unchanged': 'So prefills use the same one-grid tile config' in s,
}
for k, v in checks.items():
print(f'{k}: {v}')
python3 - <<'PY'
# Extract the relevant function body from raw source text to report whether its
# return arguments depend on tile_config.
text = open('vllm/model_executor/layers/quantization/exl3.py').read()
idx = text.index('def prepare_tier(')
end = text.index('\n prepared_tiers.append', idx)
body = text[idx:end]
print('--- prepare_tier body ---')
print(body)
PY
PYRepository: local-inference-lab/vllm
Length of output: 12139
🌐 Web query:
SparkInfer mixed trellis prepare_weights trellis3_t256_proj allocate compact tier-ordered copy
💡 Result:
The terms mentioned relate to the internal weight preparation, memory management, and kernel execution flow of the SparkInfer runtime (also referred to in some contexts as b12x), a library specialized for high-performance Mixture-of-Experts (MoE) and LLM inference on NVIDIA Blackwell (SM120/SM121) architectures [1][2][3]. The sequence described corresponds to the workflow used to prepare, allocate, and execute fused MoE layers: 1. SparkInfer (Trellis Path): The system uses a Trellis-based API for handling weight preparation and kernel execution [4]. Specifically, trellis3_t256_proj refers to a specific projection kernel or tile configuration (Trellis256) designed for efficient MoE decoding [4]. 2. prepare_weights: This is a public API method in the SparkInfer/b12x library used to validate, wrap, and prepare model weights (e.g., W4A16 quantization) before they are bound for inference [4][1]. 3. allocate & compact: These operations are part of the memory management and scratch-space planning. allocate refers to reserving GPU memory (scratch buffers) necessary for the kernel's execution, while compact often refers to optimizing these memory layouts or metadata buffers (e.g., ensuring contiguous or efficiently packed data) to avoid fragmentation and ensure CUDA graph compatibility [1][2][5]. 4. tier-ordered copy: This describes an optimized data transfer or arrangement pattern, likely involving a multi-tier memory strategy where weights or expert data are arranged based on their tier (e.g., expert size or compute capacity) to ensure efficient access during fused MoE dispatch [6][7]. In summary, this sequence is a low-level routine in the SparkInfer runtime that manages the lifecycle of expert weights—from initial preparation through memory allocation and compacting, ending in a tier-ordered copy—to ensure that expert data is perfectly aligned for the Blackwell-native CUDA kernels used in MoE decoding [4][2][5].
Citations:
- 1: https://github.com/local-inference-lab/sparkinfer
- 2: https://github.com/lukealonso/b12x
- 3: https://github.com/gittensor-ai-lab/sparkinfer
- 4: moe: add planned EXL3 Trellis path for fused Blackwell MoE b12x#49
- 5: https://github.com/gittensor-ai-lab/sparkinfer-moe/blob/main/README.md
- 6: Second part of refactoring the routing part flashinfer-ai/flashinfer#2993
- 7: https://github.com/flashinfer-ai/flashinfer/blob/2b150b39/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cu
Avoid preparing identical mixed tiers twice.
_mixed_trellis_prefill_tile_config always returns _mixed_trellis_tile_config, so mixed_tile_config and prefill_tile_config are identical. prepare_tier only changes tile_config, fc1_tile_n, and fc2_tile_n, but those are identical for both calls. Reuse prepared_tiers[-1] for prefill_tiers when the configs match instead of calling prepare_weights again.
🤖 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 2407 - 2413,
Update the tier preparation flow to reuse the already prepared final mixed tier
for prefill when _mixed_trellis_prefill_tile_config produces the same
configuration as _mixed_trellis_tile_config. In the relevant prepare_tier
caller, assign prefill_tiers from prepared_tiers[-1] instead of invoking
prepare_weights again, while preserving separate preparation when the
configurations differ.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
vllm/model_executor/layers/quantization/exl3.py (1)
128-135: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd Google-style sections to
_resolve_mixed_trellis_prefill_block_m.The docstring describes the policy but does not document its keyword-only inputs or return value. Add
Args:andReturns:sections for all eight parameters and the selected block size.Proposed docstring update
tuning and every other model geometry retain the configured value. + + Args: + configured_block_m: Configured prefill block size. + explicit_override: Whether the operator explicitly set the block size. + hidden_size: Model hidden size. + intermediate_size: Per-partition intermediate size. + tier_signature: Mixed-bitrate tier signature. + topk: Router top-k value. + device_major: CUDA device major compute capability. + prefill_tile_config: SparkInfer prefill tile configuration. + + Returns: + The selected prefill block size. """As per coding guidelines: "Use Google-style docstrings in Python code, with
Args:/Returns:/Raises:sections instead of reStructuredText/Sphinx fields."🤖 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 128 - 135, Update the docstring for _resolve_mixed_trellis_prefill_block_m to add Google-style Args and Returns sections. Document all eight keyword-only parameters and describe the selected prefill block size returned, while preserving the existing policy 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.
Outside diff comments:
In `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 128-135: Update the docstring for
_resolve_mixed_trellis_prefill_block_m to add Google-style Args and Returns
sections. Document all eight keyword-only parameters and describe the selected
prefill block size returned, while preserving the existing policy description.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 9be15411-155e-4b5d-bed5-aceee5a0bc53
📒 Files selected for processing (2)
tests/quantization/test_exl3.pyvllm/model_executor/layers/quantization/exl3.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/quantization/test_exl3.py
Assisted-by: OpenAI Codex Signed-off-by: Michel Belleau <michel.belleau@malaiwah.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
Assisted-by: OpenAI Codex
Signed-off-by: Martin Vit <martin@voipmonitor.org>
25ddfa5 to
d3b2d69
Compare
Co-authored-by: OpenAI Codex <codex@openai.com>
voipmonitor
changed the title
|
Paired TP4/DCP4 qualification is complete, with one correction to the original CI diagnosis. The original exact-head field gate passed for vLLM Correction: the 24-hour self-hosted pre-commit timeout did mask real changed-file violations; it was not solely an infrastructure failure. Exact local replay found and fixed Ruff formatting, EXL3 technical-term spelling, SPDX, mypy narrowing/signatures, a forbidden stdlib The repaired runtime was then applied onto the current Release disposition: the #228 + #306 feature/runtime composition passes. Remaining merge work is repository-side: merge #306 into this head branch, then explicitly dispose of the repository-wide |
3 participants
Summary
Provide one self-contained GG integration for the r20 EXL3 runtime. The branch
targets
dev/gilded-gnosisdirectly and includes the required behavior from#225 and #226, plus support for shared-H mixed K3/K4 checkpoints.
Changes
192/64, 206/50, and 148/108 partitions;
expert weights;
TP rank/size, quantization parameters, seed, and schema;
off,readonly, andreadwritecache modes.A warm cache hit does not import the encoder or materialize the FP32 source
transpose. The Docker release binds the cache to
/cache/exl3-online, so itsurvives restarts while the content key prevents stale reuse.
Scope
online-quant policy.
behavior.
Validation
tests/quantization/test_exl3.py: 44 passed.py_compile, andgit diff --check: pass.willfalco/GLM-5.2-EXL3-TR3-3.42bpw@ae68c65947efa90bea37308e15421872f124c46dverified.
MTP3.
The exact 3.42 checkpoint layer-3 benchmark is numerically exact against the
serial K3/K4 oracle and the mixed kernel is 1.567x faster than two serial tier
launches.
Relationship to earlier PRs
This PR supersedes the release-integration role of #225 and #226 without a
stacked PR chain. It does not merge or alter the companion B12X runtime
ABI change in local-inference-lab/b12x#117.
B12X package port validation
The current head uses the renamed
b12xPython package throughout the EXL3 runtime and tests. On RTX PRO 6000 Blackwell GPU 7, the vLLM EXL3 suites passed 74/74 and the paired B12X Trellis packaging, dense-linear, and mixed-MoE suites passed 43/43. The port changes names/imports only; kernel geometry, online K6 caching, and quantization policy are unchanged. This PR owns the EXL3 portion of the package rename paired with #246.Mixed-Trellis route-pack prewarm
The current head also preserves the single required integration commit from superseded #250. During the existing pre-KV kernel-warmup phase, vLLM asks B12X #126 to materialize the finite route-pack specialization set reachable by mixed-Trellis targets and native MTP drafts. This moves approximately 2 MiB/rank of module residency ahead of KV sizing and prevents first-request JIT allocation failures. It does not change model math or steady-state dispatch.
Focused integration tests against the current r30 runtime passed 3/3. The companion B12X head passed its six focused tests and GPU qualification found no post-start route-pack JIT, byte-identical output/acceptance, and a -0.010% or smaller logical KV change in the measured profiles. The original #250 is closed as superseded; release manifests should pin this PR head plus B12X #126, not #250 separately.