Repository navigation
[ROCm][DSV4][Perf] Use FP8 WO_A output projection - #54894
Conversation
mergify
Bot
added
deepseek
c9628e5 to
a9cf3bc
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
please do a full 1319 gsm8k test, current result seems to be lower than the baseline. We expect the full test score near 0.96 |
|
shen-shanshan
added
the
verified
shen-shanshan
left a comment
shen-shanshan
left a comment
There was a problem hiding this comment.
Overall LGTM, left some minor suggestions. Btw, could you please add some performance comparison about TPOT, not only TTFT and throughput?
| if ( | ||
| weight is None | ||
| or scale is None | ||
| or weight.dim() != 2 | ||
| or scale.dim() != 2 | ||
| or weight.dtype not in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) | ||
| ): | ||
| return | ||
|
|
||
| groups = self.n_local_groups | ||
| out_per_group = self.o_lora_rank | ||
| out_features, in_features = weight.shape | ||
| if ( | ||
| out_features != groups * out_per_group | ||
| or out_per_group % 128 != 0 | ||
| or in_features % 128 != 0 | ||
| or scale.shape != (out_features // 128, in_features // 128) | ||
| ): | ||
| return | ||
|
|
||
| e8m0_scale = _wo_a_block_scale_to_e8m0(scale) | ||
| if e8m0_scale is None: | ||
| return |
There was a problem hiding this comment.
We should add some info about the failure reason, instead of directly return them.
There was a problem hiding this comment.
| Concurrency | BF16 wo_a output tok/s | FP8 wo_a output tok/s | Throughput change | BF16 TPOT | FP8 TPOT | TPOT change | BF16 TTFT | FP8 TTFT |
|---|---|---|---|---|---|---|---|---|
| 1 | 63.35 | 64.55 | +1.89% | 15.43 ms | 15.15 ms | -1.85% | 375.74 ms | 367.52 ms |
| 4 | 213.53 | 219.99 | +3.02% | 17.76 ms | 17.23 ms | -2.97% | 1009.75 ms | 986.39 ms |
| 16 | 662.89 | 677.29 | +2.17% | 22.27 ms | 21.75 ms | -2.34% | 1899.74 ms | 1908.81 ms |
| 64 | 1386.37 | 1431.21 | +3.23% | 42.15 ms | 40.72 ms | -3.39% | 3999.62 ms | 3984.62 ms |
| def _wo_a_block_scale_to_e8m0(scale: torch.Tensor) -> torch.Tensor | None: | ||
| """Return raw E8M0 exponent bytes for a WO_A block-scale tensor.""" | ||
| if scale.dtype == torch.float8_e8m0fnu: | ||
| return scale.view(torch.uint8).contiguous() | ||
| if scale.dtype == torch.uint8: | ||
| return scale.contiguous() | ||
| if not scale.dtype.is_floating_point: | ||
| return None | ||
|
|
||
| scale_f32 = scale.detach().float() | ||
| if not bool(torch.isfinite(scale_f32).all()) or bool((scale_f32 <= 0).any()): | ||
| return None | ||
| exponent = torch.round(torch.log2(scale_f32)) | ||
| if not torch.equal(torch.exp2(exponent), scale_f32): | ||
| return None | ||
| biased = exponent.to(torch.int32) + 127 | ||
| if int(biased.min()) < 0 or int(biased.max()) > 255: | ||
| return None | ||
| return biased.to(torch.uint8).contiguous() |
There was a problem hiding this comment.
Could you please add some explanations or references of why the conversion rules look like this?
There was a problem hiding this comment.
There was a problem hiding this comment.
vLLM loader compatibility / validation
↓
raw uint8 E8M0 normalization
↓
inverse_rope_group_quant -> activation raw E8M0
↓
batched_gemm_a8w8_mxscale
↓
AITER tuned dispatch -> OPUS kernel
a9cf3bc to
8a03995
Compare
shen-shanshan
left a comment
shen-shanshan
left a comment
There was a problem hiding this comment.
Overall LGTM, leaving one final suggestion to make the code more robust. Ready label added, you can trigger the CI yourself after modification.
| del mxscale_op, inverse_quant_op | ||
|
|
||
| weight = getattr(self.wo_a, "weight", None) | ||
| scale = getattr(self.wo_a, "weight_scale_inv", None) |
There was a problem hiding this comment.
Can you align this with:
vllm/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
Lines 1632 to 1636 in 708c249
Use:
vllm/vllm/model_executor/layers/quantization/utils/fp8_utils.py
Lines 53 to 65 in 708c249
shen-shanshan
added
the
ready
|
✅ @LiuYinfeng01, CI is now available for this PR.
|
Signed-off-by: Liuyinfeng01 <yinfeliu@amd.com>
Resolve the wo_a block scale through get_fp8_block_weight_scale() with the ModelOpt MXFP8 weight_scale fallback, matching rocm_aiter_mla_sparse, and use is_fp8() for the weight dtype check. Log the concrete reason whenever the FP8 WO_A fast path falls back to BF16. Signed-off-by: Liuyinfeng01 <yinfeliu@amd.com>
37a5796 to
d910c96
Compare
|
/ci run |
|
✅ Triggered Buildkite CI #90044 for commit |
Signed-off-by: Liuyinfeng01 <yinfeliu@amd.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
Summary
Replace the ROCm DeepSeek V4 BF16
wo_aoutput-projection path on gfx950 with:inverse_rope_group_quant, which fuses inverse RoPE with per-tokenE8M0 group quantization.
batched_gemm_a8w8_mxscale, which consumes the quantized activationand the checkpoint's native FP8
wo_aweight.This removes the BF16 grouped einsum and keeps
wo_ain its checkpoint FP8form. Unsupported weights, scales, devices, and older AITER installations
retain the existing BF16 fallback.
Dependency
This fast path requires AITER >= 0.1.20, which added the MX-scale batched GEMM.
It therefore depends on #52826. Current
mainremains safe: the capabilityprobe logs once and falls back to BF16 when the op is unavailable.
Implementation
VLLM_ROCM_USE_AITER_FP8BMMgate and restrict the path togfx950.
float8_e8m0fnu/uint8scales, or losslessly convert positivepower-of-two FP32 scales to biased E8M0 exponent bytes.
rocm_inv_rope_einsumpath as the fallback.Performance
Measured on 8 x MI355X, TP1/PP8, one 100,000-token prefill request, 8K prefill
chunks, prefix caching disabled, three runs:
wo_aDelta: TTFT -7.35% and input throughput +7.93%.
TP8/PP1 prefill concurrency
Measured on 8 x MI355X with TP8/PP1, fixed 8,192-token prompts and one output
token, prefix caching disabled, and 10 requests per concurrency slot. Both arms
use the same image and differ only in the FP8
wo_agate.wo_atotal tok/swo_atotal tok/sAll 1,110 requests in each arm completed without errors.
Revisions and runtime:
a9cf3bc3498eb71bb2d6f5d44c4b0049b62feb7asha256:cef549da00e0efaeadd9338ac8f351f2b96ff71a5ab8651a99bf989458bf1684a6d2b564fd671724a3720b8edf70e8d674e4d6940.3.1sha256:0f851f268bc6bdb41c5ed22663a106339b3b269a88afbef17dbda156f2fb7893Triton sparse prefill, 8K max batched tokens, no prefix cache
VLLM_ROCM_USE_AITER_FP8BMM=0/1The enabled-arm logs load
module_inverse_rope_group_quantandopus_bmm_a8w8_mxscale; neither appears in the disabled-arm logs.TP1/PP8 prefill concurrency
Measured with the same fixed 8,192-token prompts, one output token, disabled
prefix caching, and 10 requests per concurrency slot. The image, revisions,
runtime settings, and A/B gate are identical to the TP8/PP1 measurements above;
only the topology changes to TP1/PP8.
wo_atotal tok/swo_atotal tok/sAll 1,110 requests in each arm completed without errors. The enabled-arm logs
load
module_inverse_rope_group_quantandopus_bmm_a8w8_mxscaleon all eightpipeline stages; neither appears in the disabled-arm logs.
Correctness
Greedy GSM8K, zero invalid responses:
TP8/PP1 GSM8K
Topology-matched greedy 5-shot evaluation on the TP8 stack above:
wo_awo_aThe FP8 candidate first passed the 100-question gate at 93/100 with zero
invalid responses. Full-set accuracy was +0.99 percentage points versus the
topology-matched BF16 reference; one response did not contain a parseable
answer.
Validation
git diff --check: passedtests/models/test_deepseek_v4_rocm_wo_a.py: 7 passedRelated work
#45103 fused inverse RoPE and cached a BF16
wo_a; this PR retains that path asfallback and uses the native FP8 weight on supported gfx950 configurations.
AI assistance
Cursor assisted with implementation, testing, benchmark analysis, and drafting.
The human submitter reviewed the resulting change and is responsible for it.
Re-measurement on the older
DeepSeek-V4-ProcheckpointSetup
vllm/vllm-openai-rocm:nightly-d9105ea8001e0a6d77a96327d17515bb5791fb360.28.1rc1.dev472+gd9105ea80+ this PR applied tovllm/models/deepseek_v4/amd/rocm.pyDeepSeek-V4-Pro(older snapshot), FP4VLLM_ROCM_USE_AITER_FP8BMM=0vs=1, same image and same patched fileAITER_CONFIG_GEMM_A8W8_BLOCKSCALEpointed at a merged a8w8-blockscale tuning tableFast-path confirmation: the enabled arm's server log loads
module_inverse_rope_group_quantand themxscaleBMM; the disabled arm's logcontains neither (verified per run).
TP8/PP1 prefill concurrency
wo_atotal tok/swo_atotal tok/sMean gain across the sweep: +1.29%. All 1,110 requests per arm completed successfully.
TP1/PP8 prefill concurrency
wo_atotal tok/swo_atotal tok/sMean gain across the sweep: +6.09%. All 1,110 requests per arm completed successfully.
TP8/PP1 decode performance (8K input / 1K output)
Re-measured for the reviewer-requested TPOT comparison on the old checkpoint,
8 x MI355X, TP8/PP1, prefix caching disabled. Each concurrency uses four
measured waves (at least 16 requests) after a full unmeasured wave with a
different random seed. The only A/B variable is
VLLM_ROCM_USE_AITER_FP8BMM=0/1. The enabled log loadsmodule_inverse_rope_group_quantandopus_bmm_a8w8_mxscale.wo_aoutput tok/swo_aoutput tok/sThe FP8
wo_apath improves mean TPOT by 1.85%-3.39% and outputthroughput by 1.89%-3.23% across the measured decode concurrency sweep.
TP1/PP8 single 100,000-token prefill
Three runs per arm, 8K prefill chunks, prefix caching disabled. The first run of
each arm carries a cold-start spike (~13.8 s on both arms) and is excluded from
the median.
wo_awo_aDelta: TTFT -2.89%, input throughput +2.98%.
GSM8K (full 1,319, greedy, 5-shot)
Re-run after the E8M0 conversion documentation update on the exact old
checkpoint and the same TP8/PP1 stack used for the decode table. Prefix caching
was disabled; both arms used seed 42 and concurrency 64. The only A/B variable
was
VLLM_ROCM_USE_AITER_FP8BMM=0/1.wo_awo_aThe FP8 WO_A path is accuracy-neutral in this full-set run, with zero invalid
responses in both arms. Output token throughput improves by 1.11%.
Summary on this checkpoint
The direction and magnitude match the 0813 numbers above: a small gain at
TP8/PP1 and a larger one at TP1/PP8, with accuracy at or above the BF16
reference.