[EXL3] route K6/MCG shards away from the B12X decode kernel at prefill row counts - #318
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.
Same change as on the vllm-project#314 branch, kept in sync so the stacked prefill dispatch PR is testable on its own.
…l row counts Also adds an opt-in int8 input-embedding overlay (VLLM_EXL3_EMBED_BITS=8).
|
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 |
|
Hardware-validation request for the embedding overlay (the model-file half is #319): I have narrowed the public claim after auditing the memory receipt. The current evidence is a 96 GB SM120 with vLLM capped to 30.24 GiB; that is an engine-budget proof, not a hard physical 32 GB proof. If someone has an RTX 5090, the closing test is:
The capped run allocated 279,007 KV tokens and passed 3/3 retrievals. The overlay now quantizes in 4,096-row chunks, avoiding the previous 5.1 GB float32 load-time temporary. A report of total VRAM, startup accounting, exact commit SHAs, and request outcome would close the remaining hardware caveat. |
|
Production data from the single-5090 Qwen3.8-27B deployment, independently converging on this PR's dispatch-by-row-count conclusion — with one asymmetry worth recording. Decode side (confirms this PR): we shipped an env-gated M-threshold ( Prefill side (contradicts the m=2048 column, with a suspected cause): our end-to-end control measured b12x winning K6 prefill decisively — all-recon PP 1504 vs b12x-for-K6 1967 at m=3072 on RTX 5090 — where this PR's table has recon 1.08–1.4× ahead at m=2048. Difference in the recon implementations: ours allocates a fresh full-size One caution for adopters: our dispatch is a Python-level branch, safe only because decode CUDA graphs are shape-specialised and |
1 participant
Stacked on #316.
The problem
Exl3LinearMethod._apply_onesends any K6/MCG shard with 128-divisible dimensions to_b12x_trellis_linearbefore the reconstruct dispatch from #316 is consulted. On a denseQwen3.8-27B EXL3 checkpoint with serialized attention that is 208 attention projections, 64
down_projand the head — the majority of the model — running a decode-shaped kernel throughprefill.
Exl3OnlineLinearMethod.applydoes the same unconditionally.B12X is the right kernel at decode and the wrong one at prefill. Measured on SM120
(RTX PRO 6000, real geometries, median of 10):
down_proj17408x5120 K6lm_head5120x248320 K6in_proj_qkvz5120x16384 K6The change
Route by row count inside
vllm::b12x_trellis_linear_out, which is already an opaquecustom op. A Python-level
if rows >= Naround the two calls is resolved once at trace timeand would bake the prefill branch into the decode graphs — the same trap documented in #316,
where it cost 56.5 -> 22.6 tok/s at C1 while prefill improved.
Measured end to end
Dense Qwen3.8-27B EXL3 K5/K6 with serialized K5 attention,
--max-num-seqs 8, graphs on,median of 3 runs:
+3.4 % prefill, decode unchanged. Fidelity cost measured on 136 held-out contexts with
exact full-vocabulary KL against BF16: +0.0000377, 95 % CI [−0.00001, +0.00009], 59/136
contexts — a coin flip, i.e. free.
Also in this branch: an opt-in int8 input-embedding overlay
VLLM_EXL3_EMBED_BITS=8narrows the input embedding table to per-row symmetric int8. On thismodel that is 2.543 GB -> 1.272 GB resident. It starts at native 262,144 under a 30.24 GiB
vLLM engine budget (279,007 KV tokens), for +0.000065 mean KLD and unchanged multimodal
scoring. The proof host was a capped 96 GB SM120, not a physically constrained RTX 5090; a
hard-limit 5090 run remains required. Quantization is row-chunked so a full 5.1 GB float32
temporary cannot defeat that use case. It needs the companion model-file change (separate PR)
because
VocabParallelEmbeddingis constructed without a quant config.Negative result recorded so nobody repeats it: emitting the reconstructed weight as FP8 E4M3
and using
torch._scaled_mmgives +31 % prefill but costs +0.0141 mean KLD — worsethan official FP8. Row-wise scaling does not help; the loss is FP8 activations. Left behind
VLLM_EXL3_PREFILL_FP8=1, off by default.