Repository navigation
[Bugfix][Kernel] silu_and_mul_quant: run the vectorized loop - #61716
Conversation
act_and_mul_quant_kernel computed its vectorized loop bound as round_to_previous_multiple_of(elems_per_128bit_load, block_end), which rounds 8 down to a multiple of block_end and is 0 for any d > 8. The 128-bit loop therefore never ran and every element went through the scalar cleanup loop with 2-byte loads and 1-byte stores. Swap the arguments so the bound is block_end rounded down to the vector width, and take the vectorized path only when d is a multiple of the vector width and both data pointers are 16-byte aligned, since rows of other widths are not 16-byte aligned. For the scalar case the cleanup loop now starts at the block's own range instead of 0. Output is unchanged: both loops evaluate the same per-element expression. Signed-off-by: Monishver Chandrasekaran <monishverchandrasekaran@gmail.com>
|
Gist: the vectorized loop in silu_and_mul_quant never ran because the two arguments of round_to_previous_multiple_of were swapped, so every element went through the scalar path. Output is bit-identical to main. This was the first part of #59514, which also rewrote the kernel. I split it so this PR carries only the fix, keeping it minimal and easy to review. On H100, the kernel is 2.18x faster than main (geomean over 32 configs), and Llama-3.1-8B-FP8 serving goes from 37.76 to 38.75 req/s (+2.6%), 1.0% below the Inductor path (I have tested further changes, a flat grid and an fp32 fast-math chain, that close this gap and get slightly ahead of Inductor; they would be a follow-up PR). Output logs: silu_mul_quant_fixA_logs.tgz @mgoin, @yewentao256, can you pls review this when you get some time? |
yewentao256
left a comment
There was a problem hiding this comment.
Thanks for the work!
Please do the following:
- full reproduce command line, error log in main and the same command line works in current PR
- e2e accuracy using
lm_eval... - e2e benchmark using
vllm bench serve...
When doing reproduce / e2e test, please attach with raw output log in case some agents fake the data. DO NOT use AI generated Table here
For example:
pytest tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py::test_extract_hidden_states_tp2
# main
| (EngineCore pid=1580) File "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/core.py", line 1070, in __init__
| (EngineCore pid=1580) super().__init__(
| (EngineCore pid=1580) File "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/core.py", line 160, in __init__
| (EngineCore pid=1580) self.scheduler: SchedulerInterface = Scheduler(
| (EngineCore pid=1580) ^^^^^^^^^^
| (EngineCore pid=1580) File "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/async_scheduler.py", line 14, in __init__
| (EngineCore pid=1580) super().__init__(*args, **kwargs)
| (EngineCore pid=1580) File "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py", line 277, in __init__
| (EngineCore pid=1580) self.kv_cache_manager = KVCacheManager(
| (EngineCore pid=1580) ^^^^^^^^^^^^^^^
| (EngineCore pid=1580) File "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_manager.py", line 151, in __init__
| (EngineCore pid=1580) self.coordinator = get_kv_cache_coordinator(
| (EngineCore pid=1580) ^^^^^^^^^^^^^^^^^^^^^^^^^
| (EngineCore pid=1580) File "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_coordinator.py", line 891, in get_kv_cache_coordinator
| (EngineCore pid=1580) return HybridKVCacheCoordinator(
| (EngineCore pid=1580) ^^^^^^^^^^^^^^^^^^^^^^^^^
| (EngineCore pid=1580) File "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_coordinator.py", line 563, in __init__
| (EngineCore pid=1580) assert all(
| (EngineCore pid=1580) ^^^^
| (EngineCore pid=1580) AssertionError: Each KV cache group's real block_size must be divisible by hash_block_size. block_sizes=[544, 544, 544, 544, 181], hash_block_size=98464
| FAILED
| v1/kv_connector/extract_hidden_states_integration/test_extraction.py::test_extract_hidden_states_tp2
# Now
================= 1 passed, 23 warnings in 174.57s (0:02:54) =================="lm_eval --model local-completions --model_args "base_url=http://127.0.0.1:9256/v1/completions,model=$MODEL,num_concurrent=1024" --tasks gsm8k --trust_remote_code"
# Now
|Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k| 3|flexible-extract| 5|exact_match|↑ |0.9212|± |0.0074|
| | |strict-match | 5|exact_match|↑ |0.9212|± |0.0074|
# main
|Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k| 3|flexible-extract| 5|exact_match|↑ |0.9151|± |0.0077|
| | |strict-match | 5|exact_match|↑ |0.9158|± |0.0076|
`lm_eval run --model local-chat-completions --model_args base_url=http://127.0.0.1:9256/v1/chat/completions,model=$MODEL,num_concurrent=1024,seed=256 --tasks aime25 --trust_remote_code`
# now
|Tasks |Version|Filter|n-shot| Metric | |Value | |Stderr|
|------|------:|------|-----:|-----------|---|-----:|---|-----:|
|aime25| 0|none | 0|exact_match|↑ |0.7333|± |0.0821|
# main
|Tasks |Version|Filter|n-shot| Metric | |Value | |Stderr|
|------|------:|------|-----:|-----------|---|-----:|---|-----:|
|aime25| 0|none | 0|exact_match|↑ |0.7333|± |0.0821|vllm bench serve --model zai-org/GLM-5.2-FP8 --dataset-name random --host 127.0.0.1 --port 9256 --random-input-len 8192 --random-output-len 1 --request-rate inf --num-prompts 128 --num-warmups 16
# Now
============ Serving Benchmark Result ============
Successful requests: 128
Failed requests: 0
Benchmark duration (s): 3.08
Total input tokens: 1048576
Total generated tokens: 128
Request throughput (req/s): 41.50
Output token throughput (tok/s): 41.50
Peak output token throughput (tok/s): 76.00
Peak concurrent requests: 128.00
Total token throughput (tok/s): 340041.81
---------------Time to First Token----------------
Mean TTFT (ms): 1669.15
Median TTFT (ms): 1669.02
P99 TTFT (ms): 2724.38
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 0.00
Median TPOT (ms): 0.00
P99 TPOT (ms): 0.00
---------------Inter-token Latency----------------
Mean ITL (ms): 0.00
Median ITL (ms): 0.00
P99 ITL (ms): 0.00
==================================================
# main
============ Serving Benchmark Result ============
Successful requests: 128
Failed requests: 0
Benchmark duration (s): 3.23
Total input tokens: 1048576
Total generated tokens: 128
Request throughput (req/s): 39.63
Output token throughput (tok/s): 39.63
Peak output token throughput (tok/s): 50.00
Peak concurrent requests: 128.00
Total token throughput (tok/s): 324671.20
---------------Time to First Token----------------
Mean TTFT (ms): 1668.69
Median TTFT (ms): 1649.00
P99 TTFT (ms): 3126.43
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 0.00
Median TPOT (ms): 0.00
P99 TPOT (ms): 0.00
---------------Inter-token Latency----------------
Mean ITL (ms): 0.00
Median ITL (ms): 0.00
P99 ITL (ms): 0.00
==================================================|
Thanks for the review @yewentao256. Raw logs below (all from the tgz attached in my comment above). Reproduce: the bug is a dead code path, so the probe is scripts/proof_main.py (in the tgz). It runs the op with the input and output pointers deliberately misaligned by 8 bytes. The vectorized loop uses 16-byte loads, and a 16-byte load on a misaligned pointer faults with a misaligned address. So on main, if the op completes, the vectorized loop did not run. It completes for d = 14336 and 10752, and faults only at d = 8, the one shape where the swapped bound is non-zero. On this PR, all three complete, because misaligned pointers now take the scalar loop. In the output, A is the real shapes (d = 14336, 10752) where the swapped bound is 0, and B is d = 8, the one shape where the bound is non-zero, and the loop actually executes. The last line is the script's verdict for the main run (it checks whether d = 8 faults); on this branch the absence of a fault is the intended result. |
yewentao256
left a comment
There was a problem hiding this comment.
LGTM, thanks for the work!
|
/ci run |
|
✅ @Monishver11, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #95155 for commit ℹ️ This PR is 56 commits behind upstream |
|
/ci retry |
|
✅ Queued 2 failed job(s) for retry in Buildkite CI #95155. |
|
Hi @yewentao256, CI is green except |
|
Thanks @yewentao256 for the review and @khluu for the merge. |
…oject#61716) Signed-off-by: Monishver Chandrasekaran <monishverchandrasekaran@gmail.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
Purpose
silu_and_mul_quant(act_and_mul_quant_kernelincsrc/quantization/activation_kernels.cu) is the fused SiLU-and-mul + static per-tensor FP8 quant that the manual activation+quant fusion (#51415,maybe_fused_act_quant) routes static-FP8 models through by default.Its 128-bit vectorized loop has never executed. The loop bound is computed as
round_to_previous_multiple_of(elems_per_128bit_load, block_end), with the two arguments swapped: the helper rounds its first argument down to a multiple of its second, so this rounds 8 down to a multiple ofblock_end, which is 0 for anyd > 8. Withvec_loop_end = 0the vectorized loop has an empty range and every element goes through the scalar cleanup loop with 2-byte loads and 1-byte stores. Present since the kernel landed in #10867. Runtime proof onmain: with the input and output pointers misaligned by 8 bytes the op completes and matches the aligned result for d = 14336 and 10752 (no 16-byte load executed), and faults withmisaligned addressat d = 8, the one shape where the bound is non-zero.The cost: on H100 the kernel runs at 1.9 TB/s effective against 3.1 TB/s for the Inductor-fused kernel it replaced, and Llama-3.1-8B-FP8 serves about 3.5% slower with the fusion wired than on the Inductor path (numbers below).
This PR:
block_endrounded down to the vector width and the 128-bit loop covers the block.dis a multiple of the vector width (8 for bf16/fp16) and both data pointers are 16-byte aligned, via the existingvllm::is_16byte_aligned. Rows of other widths are not 16-byte aligned (d = 781, hidden size 1562 in the unit test, starts odd rows 4 bytes off) and would fault on the fixed loop; they keep the scalar loop, which is what every shape runs today.block_startinstead of 0. With 2 to 4 blocks per row at 32 tokens or fewer, starting at 0 made each block redo the preceding blocks' elements (same values written, wasted bandwidth).Output is bit-identical to
main: both loops evaluate the same per-element expression. Verified below. No change to grid, block size or arithmetic; the remaining gap to the Inductor kernel (0.9x at prefill, about 0.5x at decode) is the subject of a separate follow-up.Not a duplicate. No open PR touches this kernel. #49828 adds a separate per-token dynamic kernel in the same file with its own launch. #59514 is my earlier PR that bundled this fix with a kernel rewrite; it is being split into this fix and a follow-up so each can be reviewed on its own.
Test Plan
mainand of the alignment guard on this branch (misaligned-pointer test, script in the attached logs).main(4b61f30): sha256 of the output bytes over bf16 and fp16, tokens 1 to 16384, d in {8192, 10752, 14336, 28672} plus {8, 24, 781, 2048, 4104, 10760, 14341}, and misaligned-pointer cases; 302 cases.pytest tests/kernels/test_fused_quant_activation.py tests/compile/passes/test_silu_mul_quant_manual_fusion.py tests/compile/passes/test_silu_mul_quant_fusion.py tests/compile/passes/test_functionalization.pydo_bench_cudagraph; the op on both builds and atorch.compileofSiluAndMul.forward_native+QuantFP8(static).forward_nativeas the Inductor column.mainkernel, this fix, and the same tree with thellama.pywiring line reverted (Inductor path):Test Result
main, d = 14336 and 10752 complete on 8-byte-misaligned pointers and match the aligned output, d = 8 faults withCUDA error: misaligned address. On this branch all three complete and match (the guard sends misaligned inputs to the scalar loop).main: 302 identical, 0 mismatched.Microbenchmark, bf16, d = 14336 (us):
Geomean 2.18x vs main over 32 configs (tokens 1 to 16384, four d values), range 1.25x to 5.07x. Against Inductor: 0.34x to 0.66x at 256 tokens and below, 0.65x to 0.96x from 1024 tokens up (the 0.65x is 1024 tokens at d = 8192; the other prefill configs are 0.82x to 0.96x). The wider sweep in the attached logs (bf16 and fp16, 64 cases) gives 1.99x geomean, range 1.26x to 4.67x.
E2e, Llama-3.1-8B-FP8, H100
Identical, as expected from bit-identical kernel output.
vllm bench serve, three interleaved rounds per arm:Per-round req/s: main 37.82 / 37.77 / 37.70; fix 38.68 / 38.55 / 39.02; Inductor 39.01 / 39.32 / 39.12. The kernel was switched by rebuilding between rounds. Every fix round is above every main round and below every Inductor round.
Profile, one
vllm bench latencyiteration (batch 32, 1024/32), 1088 launches of the fused op per arm:Full logs attached.
I used Claude Code for this change. I reviewed every line and ran the tests and benchmarks above myself.
Essential Elements of an Effective PR Description Checklist