[EXL3] Route prefill batches through the planned Trellis MoE (+58-64% prefill, stacked on #139)

[davidsyoung]

Status (2026-07-22): incorporated into #139 (commits a59770430/f83fc162f, byte-identical, authorship preserved) as of its v20 rebase. Review #139; this PR is the isolated pre-rebase review scope and closes when #139 merges.

Stacked on #139 (@brandonmusic's rank-sliced EXL3 backend); this branch is based on the #139 head (2fdc6d3676), so the Files-changed tab includes its commits until #139 merges. Review scope — exactly this PR's changes: compare 2fdc6d3676...exl3-prefill-trellis-plan · commit d57d77534 (the fix, exl3.py only, +91/−28) · commit f80a56df1 (optional CPU tests — happy to drop). Will rebase after #139 lands so the PR contains only its own commits.

Stacked on your exl3-backend branch (#139) — merging here lands the commits inside #139 with no history rewrite. No change to weight loading, formats, or the decode path; this only extends the rank-sliced batch dispatch.

Problem

_apply_rank_sliced 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, capping prefill at ~2.0–2.3k tok/s while the same weights decode at full speed through the planned Trellis window.

Change

Build a second trellis_moe plan at max_tokens=max_num_batched_tokens with block_size_m=64 (VLLM_EXL3_PREFILL_BLOCK_M, allowed {8,16,32,48,64}) beside the decode plan (32/8), and dispatch three ways:

• m ∈ [min, max] → decode plan (unchanged)
• max < m ≤ capacity → prefill plan — sparkinfer bind already accepts any tokens ∈ [1, max_tokens], so zero kernel changes are 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 the memory profiler sees it and the KV pool auto-shrinks instead of risking request-time OOM.

Measured

Matched A/B serving brandonmusic/GLM-5.2-EXL3-TR3-3.0bpw on your published verdictai/glm52-exl3-sparkinfer runtime image (4× 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 (99.7→101.8 aggregate across our runs); 57k-token needle retrieval exact; zero request errors.

Second commit adds CPU contract tests (mocked plan/ext APIs, no GPU/sparkinfer/ext needed) following the test_exl3.py pattern — happy to drop it if you'd rather keep the branch minimal.

Related build-compat PR on the sparkinfer side (independent, for PyPI DSL 4.6.0 users): see the companion PR on exl3-trellis-fused.

Summary by CodeRabbit

• New Features

  • Added support for EXL3-quantized models, including rank-sliced MoE checkpoints.
  • Improved KV-cache handling for contiguous layouts and replicated attention workloads.
  • Added safer workspace separation for speculative decoding.

• Bug Fixes

  • Improved structured-output reasoning boundary handling during speculative decoding.
  • Corrected sparse attention metadata and DeepSeek shared-layer behavior.
  • Refined PCIe communication buffer sizing and defaults for improved reliability.

• Tests

  • Expanded coverage for quantization, attention, KV-cache layouts, distributed execution, structured output, and workspace management.

[davidsyoung]

Companion sparkinfer build-compat PR (independent; needed to build this stack on public PyPI DSL 4.6.0): local-inference-lab/sparkinfer#70

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds EXL3 quantization for dense and rank-sliced MoE models, including metadata validation, loading, Trellis planning, parity fallback, and model integration. It also updates workspace lanes, distributed capture, KV-cache planning, attention paths, structured-output advancement, and optional FlashInfer JIT-cache installation.

Changes

EXL3 quantization and rank-sliced MoE

Layer / File(s)	Summary
EXL3 contracts and registration vllm/config/model.py, vllm/model_executor/layers/quantization/*, tests/quantization/test_exl3.py	Registers EXL3, validates metadata and tensor storage, and adds the custom GEMM operation.
Dense EXL3 weights and GEMM vllm/model_executor/layers/quantization/exl3.py	Adds dense parameter handling, tensor-parallel slicing, validation, and EXL3 GEMM execution.
EXL3 MoE storage and runtime vllm/model_executor/layers/quantization/exl3.py, tests/quantization/test_exl3*.py	Adds expert storage, slab preallocation, Trellis preparation, decode/prefill plans, parity fallback, and capture guards.
Rank-sliced checkpoint loading vllm/model_executor/models/*.py, vllm/model_executor/layers/fused_moe/routed_experts.py, tests/quantization/test_exl3.py	Normalizes local rank names, skips nonlocal shards, and distinguishes indexed expert tensors from fused tensors.

Workspace and distributed capture

Layer / File(s)	Summary
Independent workspace lanes vllm/v1/worker/workspace.py, vllm/v1/worker/gpu*.py, tests/v1/worker/test_workspace.py	Adds lane-scoped workspace allocation and uses a separate lane for speculative execution.
B12X DCP and PCIe capture vllm/v1/attention/ops/dcp_alltoall.py, vllm/distributed/*, vllm/envs.py, tests/distributed/*	Adds active-group DCP capture, independent channels, parsed PCIe oneshot limits, and updated defaults.

KV-cache and attention behavior

Layer / File(s)	Summary
KV-cache layout and replication vllm/v1/core/kv_cache_utils.py, vllm/v1/kv_cache_interface.py, tests/v1/core/*	Adds native contiguous-page planning, layout validation, replicated MLA sizing, grouping, and memory calculations.
Sparse MLA and indexer paths vllm/v1/attention/backends/mla/*, vllm/models/deepseek_v32/nvidia/attention.py, tests/v1/attention/*	Uses GPU positions for DCP metadata, simplifies decode selection, and skips indexer work for shared iterations.

Structured-output advancement

Layer / File(s)	Summary
Token-aware reasoning advancement vllm/v1/structured_output/*, vllm/v1/core/sched/scheduler.py, tests/v1/structured_output/*	Passes current output tokens into reasoning-end detection and adjusts speculative advancement and trimming.

FlashInfer JIT cache installation

Layer / File(s)	Summary
Conditional Docker installation docker/Dockerfile	Adds INSTALL_FLASHINFER_JIT_CACHE and installs the package only when enabled.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Checkpoint
  participant ModelLoader
  participant Exl3MoEMethod
  participant Trellis
  participant Exl3Extension
  Checkpoint->>ModelLoader: provide rank-sliced weight names
  ModelLoader->>Exl3MoEMethod: normalize and load local shards
  Exl3MoEMethod->>Trellis: prepare weights and plans
  Exl3MoEMethod->>Trellis: bind runtime inputs
  Trellis->>Exl3Extension: execute rank-sliced MoE

Loading

Possibly related PRs

• local-inference-lab/vllm#64: Modifies the same B12X sparse MLA decode-selection path.
• local-inference-lab/vllm#139: Overlaps on EXL3 quantization, Trellis execution, and expert-weight loading.
• local-inference-lab/vllm#144: Overlaps on B12X/DCP capture and PCIe oneshot environment plumbing.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 21.14% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: routing EXL3 prefill batches through a planned Trellis MoE path.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 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/models/glm4_moe.py`:
- Around line 525-534: Assign the constructor’s quant_config parameter to
self.quant_config alongside self.config in Glm4MoeModel.__init__, so
load_weights can safely access it for both quantized and unquantized models.

🪄 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: 97bf2307-72f1-4fad-acca-5d9d9e7fd581

📥 Commits

Reviewing files that changed from the base of the PR and between b07bef7 and b801b62.

📒 Files selected for processing (9)

• docker/Dockerfile
• tests/quantization/test_exl3.py
• tests/quantization/test_exl3_prefill_plan.py
• vllm/config/model.py
• vllm/model_executor/layers/fused_moe/routed_experts.py
• vllm/model_executor/layers/quantization/__init__.py
• vllm/model_executor/layers/quantization/exl3.py
• vllm/model_executor/models/deepseek_v2.py
• vllm/model_executor/models/glm4_moe.py

[coderabbitai]

🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win

Persist quant_config on Glm4MoeModel before using it here.

__init__ only defines the local quant_config (Line 418) and never assigns self.quant_config. Therefore every load_weights call now raises AttributeError, including unquantized models. Add self.quant_config = quant_config alongside self.config in __init__.

🤖 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 - 534, Assign the
constructor’s quant_config parameter to self.quant_config alongside self.config
in Glm4MoeModel.__init__, so load_weights can safely access it for both
quantized and unquantized models.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

vllm/v1/attention/backends/mla/b12x_mla_sparse.py (1)

2192-2198: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Dead branch after gating simplification.

use_decode_kernel is now attn_metadata.max_query_len <= 1, so inside this block the ternary condition attn_metadata.max_query_len <= 1 is always true and the cache_seq_lens_per_token else branch is unreachable. Collapse to the direct assignment for clarity.

♻️ Proposed simplification

         use_decode_kernel = attn_metadata.max_query_len <= 1
         if use_decode_kernel:
-            cache_seqlens = (
-                attn_metadata.cache_seq_lens_per_req
-                if attn_metadata.max_query_len <= 1
-                else attn_metadata.cache_seq_lens_per_token[:num_actual_toks]
-            )
+            cache_seqlens = attn_metadata.cache_seq_lens_per_req

🤖 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/b12x_mla_sparse.py` around lines 2192 - 2198,
In the use_decode_kernel block, simplify the cache_seqlens assignment to
directly use attn_metadata.cache_seq_lens_per_req, removing the redundant
max_query_len ternary and unreachable cache_seq_lens_per_token branch.

vllm/distributed/device_communicators/custom_all_reduce.py (1)

83-95: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider a test for the negative-size validation branch.

_b12x_pcie_oneshot_limits() raises ValueError for negative parsed sizes, but only the valid-size path is covered by test_b12x_oneshot_buffer_tracks_dispatch_limits. A quick test setting a negative size (e.g. "-5") would close this gap for a reachable misconfiguration.

🤖 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/distributed/device_communicators/custom_all_reduce.py` around lines 83 -
95, The test suite should cover the negative-size validation branch in
_b12x_pcie_oneshot_limits. Add a focused test alongside
test_b12x_oneshot_buffer_tracks_dispatch_limits that configures a parsed size
such as "-5" and asserts ValueError is raised, preserving the existing
valid-size coverage.

🤖 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/v1/worker/gpu_worker.py`:
- Around line 400-409: Update the workspace sizing in the initialization flow
around init_workspace_manager to use parallel_config.num_ubatches directly
rather than deriving num_ubatches solely from enable_dbo. Preserve the existing
num_workspace_lanes calculation and pass the configured ubatch count so non-DBO
micro-batching allocates sufficient workspace slots.

---

Nitpick comments:
In `@vllm/distributed/device_communicators/custom_all_reduce.py`:
- Around line 83-95: The test suite should cover the negative-size validation
branch in _b12x_pcie_oneshot_limits. Add a focused test alongside
test_b12x_oneshot_buffer_tracks_dispatch_limits that configures a parsed size
such as "-5" and asserts ValueError is raised, preserving the existing
valid-size coverage.

In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py`:
- Around line 2192-2198: In the use_decode_kernel block, simplify the
cache_seqlens assignment to directly use attn_metadata.cache_seq_lens_per_req,
removing the redundant max_query_len ternary and unreachable
cache_seq_lens_per_token branch.

🪄 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: 5ea1ab24-7efe-4044-965e-31aa2874ab22

📥 Commits

Reviewing files that changed from the base of the PR and between b801b62 and f80a56d.

📒 Files selected for processing (25)

• tests/distributed/test_b12x_fused_all_reduce.py
• tests/distributed/test_dcp_a2a.py
• tests/quantization/test_exl3_prefill_plan.py
• tests/v1/attention/test_sparse_mla_backends.py
• tests/v1/core/test_contiguous_kv_packing.py
• tests/v1/core/test_kv_cache_utils.py
• tests/v1/structured_output/test_reasoning_structured_output.py
• tests/v1/worker/test_workspace.py
• vllm/distributed/device_communicators/custom_all_reduce.py
• vllm/distributed/parallel_state.py
• vllm/envs.py
• vllm/model_executor/layers/quantization/exl3.py
• vllm/models/deepseek_v32/nvidia/attention.py
• vllm/v1/attention/backends/mla/b12x_mla_sparse.py
• vllm/v1/attention/backends/mla/indexer.py
• vllm/v1/attention/ops/dcp_alltoall.py
• vllm/v1/core/kv_cache_utils.py
• vllm/v1/core/sched/scheduler.py
• vllm/v1/kv_cache_interface.py
• vllm/v1/structured_output/__init__.py
• vllm/v1/worker/gpu/model_runner.py
• vllm/v1/worker/gpu/spec_decode/speculator.py
• vllm/v1/worker/gpu/warmup.py
• vllm/v1/worker/gpu_worker.py
• vllm/v1/worker/workspace.py

🚧 Files skipped from review as they are similar to previous changes (2)

• vllm/model_executor/layers/quantization/exl3.py
• tests/quantization/test_exl3_prefill_plan.py

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Check whether ubatch_size can differ from 1 when enable_dbo is False, and how it's used elsewhere.
rg -nP -C3 'ubatch_size' vllm/config/parallel.py
rg -nP 'dbo_current_ubatch_id' vllm/v1 --type=py -g '!**/test/**'

Repository: local-inference-lab/vllm

Length of output: 975

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== vllm/v1/worker/gpu_worker.py ==\n'
sed -n '390,420p' vllm/v1/worker/gpu_worker.py

printf '\n== vllm/v1/worker/workspace.py ==\n'
sed -n '1,260p' vllm/v1/worker/workspace.py

printf '\n== vllm/v1/worker/ubatching.py ==\n'
sed -n '1,240p' vllm/v1/worker/ubatching.py

printf '\n== vllm/config/parallel.py relevant section ==\n'
sed -n '180,240p' vllm/config/parallel.py
sed -n '510,540p' vllm/config/parallel.py

Repository: local-inference-lab/vllm

Length of output: 23682

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
for path, spans in {
    "vllm/v1/worker/workspace.py": [(1, 260)],
    "vllm/v1/worker/ubatching.py": [(1, 240)],
}.items():
    print(f"\n== {path} ==")
    text = Path(path).read_text().splitlines()
    for start, end in spans:
        for i in range(start, min(end, len(text)) + 1):
            print(f"{i}:{text[i-1]}")
PY

Repository: local-inference-lab/vllm

Length of output: 19861

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path

for path, ranges in {
    "vllm/v1/worker/workspace.py": [(130, 210)],
    "vllm/v1/worker/ubatching.py": [(130, 190)],
}.items():
    print(f"\n== {path} ==")
    lines = Path(path).read_text().splitlines()
    for start, end in ranges:
        for i in range(start, min(end, len(lines)) + 1):
            print(f"{i}:{lines[i-1]}")
PY

Repository: local-inference-lab/vllm

Length of output: 6609

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

rg -n "use_ubatching|num_ubatches|make_ubatch_contexts|enable_dbo|ubatch_size" vllm -g '*.py'

Repository: local-inference-lab/vllm

Length of output: 5833

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path

paths = [
    "vllm/v1/worker/gpu_worker.py",
    "vllm/v1/worker/ubatching.py",
    "vllm/config/parallel.py",
]
for path in paths:
    print(f"\n== {path} ==")
    lines = Path(path).read_text().splitlines()
    for i, line in enumerate(lines, 1):
        if "use_ubatching" in line or "num_ubatches" in line or "make_ubatch_contexts" in line or "enable_dbo" in line or "ubatch_size" in line:
            start = max(1, i - 8)
            end = min(len(lines), i + 16)
            for j in range(start, end + 1):
                print(f"{j}:{lines[j-1]}")
            print("---")
PY

Repository: local-inference-lab/vllm

Length of output: 11169

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

rg -n -C3 "use_ubatching|make_ubatch_contexts|enable_dbo|ubatch_size" vllm -g '*.py'

Repository: local-inference-lab/vllm

Length of output: 18892

---

Use parallel_config.num_ubatches for workspace sizing
use_ubatching also enables ubatch_size > 1 without DBO, but this path allocates only one ubatch slot unless DBO is on. That under-sizes WorkspaceManager for supported non-DBO micro-batching and can hit an out-of-range workspace 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/worker/gpu_worker.py` around lines 400 - 409, Update the workspace
sizing in the initialization flow around init_workspace_manager to use
parallel_config.num_ubatches directly rather than deriving num_ubatches solely
from enable_dbo. Preserve the existing num_workspace_lanes calculation and pass
the configured ubatch count so non-DBO micro-batching allocates sufficient
workspace slots.

[davidsyoung]

Prebuilt test image (base = the published verdictai/glm52-exl3-sparkinfer v1 digest + exactly this PR's exl3.py, nothing else):

davidyoung/glm52-exl3-sparkinfer:v1-prefill-trellis-plan-20260722
@sha256:95b7e715e7aca733c44ee6477b2b2abcbed7bfa2bb06acf17b386463be5c0adb

Drop-in with the model repo's compose: IMAGE=davidyoung/glm52-exl3-sparkinfer:v1-prefill-trellis-plan-20260722 docker compose up -d. The fix defaults on; VLLM_EXL3_PREFILL_TRELLIS=0 restores upstream parity-path behavior for an A/B on the same image. CPU contract tests ship at /opt/exl3-tests/test_exl3_prefill_plan.py (python3 /opt/exl3-tests/test_exl3_prefill_plan.py → EXL3_PREFILL_PLAN_TESTS_OK, verified pre-push).

[davidsyoung]

Status: incorporated into #139 as of its v20 rebase.

d57d77534 and f80a56df1 are now in #139 as a59770430/f83fc162f — byte-identical, authorship preserved (verified by diff against this branch). If you're reviewing the EXL3 stack, review #139; this PR remains open only as the isolated pre-rebase review scope (base 2fdc6d3676) and will be closed as soon as #139 merges.

Also note @brandonmusic's 31ff5a2a4 in #139 (exl3(mtp): gate use_flattening off for SM120+B12X native next_n>2 path) fixes a tool-calling issue that exists on the older image this PR was benchmarked on; the prefill numbers here (+58–64%) are unaffected, but anyone deploying MTP3 should take the #139 branch or the v20-based image rather than stacking this PR alone.