[EXL3] dispatch prefill-shaped GEMMs to reconstruct+hgemm (+113 % prefill, decode unchanged) - #316
Conversation
…gemm exl3_gemm autotunes with timing launches, so _require_enforce_eager refuses non-eager execution for every dense tensor_storage checkpoint. Decode capture sizes are a known finite set, so prime every serialized shard over them during process_weights_after_loading -- the same pattern _warm_decode_shapes already uses for the online-K6 path -- and permit decode-only capture behind VLLM_EXL3_GRAPH_DECODE. Measured on Qwen3.8-27B EXL3 K4 (RTX PRO 6000, SM120, TP1): decode throughput +92%/+84%/+98% at C1/C4/C8, distribution parity vs eager exactly 0.000000 KLD over 32 contexts.
exl3_gemm decodes the trellis once per output tile, so its cost grows with M. exllamav3's own LinearEXL3.reconstruct_hgemm switches to reconstruct+hgemm above 1024 rows for that reason. Measured crossover with the extension vLLM loads is m=128, and the win at m=2048 is 4.1-5.2x per geometry. Dispatch happens inside the existing custom op, at runtime: a Python-level branch around the call is resolved at trace time, which makes the decode CUDA graphs capture the prefill path and collapses decode throughput. Measured on Qwen3.8-27B EXL3 (RTX PRO 6000, SM120, TP1): prefill 2,369 -> 5,050 tok/s at 2k and 2,362 -> 5,146 at 6k, decode unchanged, +0.43% divergence.
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review (1)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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 |
Same change as on the vllm-project#314 branch, kept in sync so the stacked prefill dispatch PR is testable on its own.
|
Follow-up measurements on this PR, including a negative result and the ceiling arithmetic that says where to stop. 1. Attribution: the MLP kernel was the whole story, the attention overlay is notFour configurations on 1x RTX PRO 6000 Blackwell (SM120), dense
2.
|
| K | N | m | ext.hgemm |
torch.mm |
ratio |
|---|---|---|---|---|---|
| 5120 | 17408 | 256 | 0.154 ms | 0.158 ms | 0.98x |
| 5120 | 17408 | 2048 | 0.903 ms | 0.911 ms | 0.99x |
| 5120 | 17408 | 4096 | 1.760 ms | 1.702 ms | 1.03x |
| 17408 | 5120 | 1024 | 0.465 ms | 0.506 ms | 0.92x |
| 17408 | 5120 | 2048 | 0.968 ms | 0.919 ms | 1.05x |
| 5120 | 32768 | 2048 | 1.585 ms | 1.613 ms | 0.98x |
0.92-1.06x across the range. Swapping in cuBLAS would buy nothing.
3. Larger prefill chunks do not amortise the reconstruct further: no change
--max-num-batched-tokens 8192 --long-prefill-token-threshold 8192: PP 5,043 / 5,120 tok/s versus 5,050 / 5,146 at defaults - within run-to-run noise, because the scheduler already places these prompts in a single chunk, so reconstruct is already paid once per prompt. That lever is closed.
4. The remaining gap is structural, and I would rather state it than chase it
At m=2048 a gate_proj-shaped call costs 1.20 ms of which the GEMM is 0.90 ms; the other ~0.30 ms is Hadamard + reconstruct. gate+up+down = 3.69 ms per layer, x64 layers = 236 ms, an MLP-only ceiling of 8.7k tok/s (we measure 5.6k end to end). A perfect fused kernel making reconstruct free would give 173 ms, i.e. 11.8k tok/s MLP-only, call it 7-8k end to end. On this box official FP8 serves 10,667 tok/s and NVFP4 14,528 tok/s, and neither pays a decode step at all.
So: prefill parity with FP8 is not reachable for a 4-bit-class trellis format by dispatch or tuning changes in this runtime. It needs dequant fused into the GEMM epilogue (Marlin-style) - a new kernel, out of scope for this PR. This PR takes the 2.1-2.3x that is available for a threshold and a scratch buffer, and I am documenting the rest as a measured limit rather than an open TODO.
5. Also pushed here
Priming now uses 0.05 * randn instead of a zero arena (kept in sync with #314). Measured on the eager-vs-graph decode-parity harness: it does not change near-tie drift (25/32 vs 24/32, identical logprob deltas), and unquantised BF16 drifts the same amount (24/32), so that drift is ambient to the graph path rather than caused by either patch - details in #314.
Method and receipts: docs/26-prefill-attribution.md, microbenchmark in tools/gemm_cmp.py, bench driver in tools/bench.py.
1 participant
Stacked on #314 (same branch base). Addresses dense EXL3 prefill throughput.
Problem
exl3_gemmis decode-shaped: it decodes the trellis once per output tile, so cost grows with M.Exl3LinearMethod._apply_oneroutes every non-K6 shard through it at every row count. exllamav3's ownLinearEXL3.reconstruct_hgemmswitches toreconstruct+hgemmabove 1024 rows precisely for this reason.Effect on a real dense checkpoint (Qwen3.8-27B, 193 EXL3 matrices): prefill 2,369 tok/s, against 10,667 for the same model in official FP8 and 14,528 in NVFP4 on the same GPU.
Kernel measurements
Speedup of reconstruct+hgemm over
exl3_gemmon the checkpoint's three real geometries, using the extensionVLLM_EXL3_EXT_PATHloads (which does not export the fusedreconstruct_had_slice, so the unfusedhad_r_128→reconstruct[_slice]→hgemm→had_r_128sequence is timed, reconstruct cost included):lm_head)Crossover is m=128 for all three, hence the default threshold.
Change
Inside the existing
vllm::exl3_gemmcustom op: ifx.shape[0] >= VLLM_EXL3_PREFILL_RECONSTRUCT_M(default 128,0disables), materialise the weight and usehgemm; otherwise callexl3_gemmexactly as today. Scratch is one fp16 buffer per(device, K, N-chunk)reused across all layers of the same geometry, bounded at 336 MB by 32768-column slicing, so peak memory does not scale with layer count.The dispatch must be inside the op. A Python-level
if rows >= thresholdaround the two calls is resolved once at trace time — with a single compiled shape range the profile run bakes in the prefill branch, the decode CUDA graphs then capture reconstruct+hgemm, and decode throughput collapses. Measured while developing this: 56.5 → 22.6 tok/s at C1 with prefill improving, i.e. a silent regression that looks like a win if only prefill is benchmarked. Two related traps:torch.cuda.is_current_stream_capturing()is not Dynamo-traceable (torch.* op returned non-Tensor), and raw pybind extension calls are not either (Attempted to call function marked as skipped: …had_r_128).End-to-end results
Same box,
--max-num-seqs 8, greedy,ignore_eos, 256 output tokens, median of 3 runs (dispersion <1 %), prefill measured with exact token-count prompts:+113 % / +118 % prefill, decode unchanged.
Numerical impact, disclosed
This changes fp16 summation order, so unlike #314 it is not bit-exact. Measured with a full-vocabulary teacher-forced comparison over 32 contexts × 2047 positions:
+0.43 % measured divergence, top-1 unchanged to four decimals. For scale, the same harness measures a 61 % gap between this checkpoint and official FP8, and the patched-vs-unpatched difference is the same order as the difference between vLLM's own logit path and an offline replay of it.
VLLM_EXL3_PREFILL_RECONSTRUCT_M=0restores the previous path exactly, so operators who need bit-identical output to today can opt out.Follow-up, not in this PR
The MLP GEMMs are no longer prefill-dominant, so the remaining gap to NVFP4 is elsewhere. The prime suspect is the online-K6 overlay: 208 attention projections go through
_b12x_trellis_linear, whose dense K6 kernel is documented as a small-M decode kernel, at prefill row counts. The same row-count dispatch likely applies there.