[EXL3] Route prefill batches through the planned Trellis MoE (+58-64% prefill)

[davidsyoung]

Stacked on your exl3-backend branch (local-inference-lab#139) — merging here lands the commits inside vllm-project#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.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[davidsyoung]

Companion sparkinfer build-compat PR (independent; needed to build this stack on public PyPI DSL 4.6.0): brandonmmusic-max/b12x#1

[davidsyoung]

Re-targeting to the upstream repo (local-inference-lab/vllm) per review — closing this one.