[ROCm][Quant] Requantize serialized MXFP8 linears to FP8 PTPC - #48427
tanpinsiang wants to merge 1 commit into
Conversation
|
Documentation preview: https://vllm--48427.org.readthedocs.build/en/48427/ |
mergify
Bot
added
documentation
6860540 to
170aa20
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Additional validation: PTPC with AITER sparse PA and shuffled KVexport VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1Frozen comparison: base selected MiniMax
Both comparison arms use the same sparse-attention path. These results measure PTPC against its control with AITER sparse PA and shuffled KV enabled; they do not isolate the standalone uplift of AITER sparse PA over Triton. QualityPaired GSM8K used
C64: does PTPC improve over native MXFP88192 input, 512 output, concurrency 64.
PTPC improved throughput by 4.595%; every reported latency metric also C256: does PTPC improve over native MXFP88192 input, 1024 output, concurrency 256.
PTPC improved throughput by 4.427%. TTFT, mean TPOT, and p99 ITL improved; is PTPC better than forced emulation?
PTPC improved throughput by 2.712% |
170aa20 to
0197453
Compare
tanpinsiang
changed the title
0197453 to
85199f1
Compare
tanpinsiang
changed the title
|
@tjtanaa @hongxiayang I addressed the request to make this generic and rebased onto current
|
|
I tested environment
PTPC startup: blocked by AITER shape coverage
The PTPC flag reaches the generic source-aware path, but kernel selection The packaged AITER tuning probe found no gfx950 preshuffled configuration for any of this model's dense shapes:
Both TP1 and TP4 failed the mandatory ready gate. New code regression checkthe code change does not hurt performance. It also shows that the PTPC activation path is operationally limited by AITER's tuned shape coverage: arbitrary ModelOpt MXFP8 checkpoints cannot use it yet. C64 throughput8192 input, 512 output, concurrency 64, 128 warmup, 64 verification, 512 measured.
Flag-absent quality: exact parityPaired GSM8K used
|
| @@ -6,6 +6,7 @@ | |||
| """ | |||
|
|
|||
| import os | |||
| from types import SimpleNamespace | |||
There was a problem hiding this comment.
Maybe these tests should be in test_online.py?
Can you also add a test that ensures that the MXFP8 -> BF16 -> FP8 PTPC conversion does not double memory requirement during requantization (due to BF16 dequant)? Is it indeed the case?
Maybe similar to
vllm/tests/quantization/test_online.py
Line 188 in d6dbdb9
| if self.requantization_source is not None and current_platform.is_rocm(): | ||
| from vllm.model_executor.kernels.linear import ( | ||
| AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, | ||
| ) | ||
|
|
||
| if not isinstance( | ||
| self.fp8_linear, | ||
| AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, | ||
| ): | ||
| raise RuntimeError( | ||
| "ROCm source-aware FP8 PTPC requantization requires the " | ||
| "AITER preshuffled per-token FP8 kernel, selected " | ||
| f"{type(self.fp8_linear).__name__}." | ||
| ) |
There was a problem hiding this comment.
Why is this done here? Shouldn't choose_scaled_mm_linear_kernel be the source of truth and already handle this:
vllm/vllm/model_executor/kernels/linear/__init__.py
Lines 555 to 577 in d6dbdb9
There was a problem hiding this comment.
Same question for the above isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel actually..
| if self.requantization_source is not None: | ||
| logger.info_once( | ||
| "Requantized serialized weights to FP8 PTPC via %s", | ||
| type(self.fp8_linear).__name__, | ||
| scope="global", | ||
| ) |
There was a problem hiding this comment.
Same comment as #49313 (comment), for an other PR
| get_current_vllm_config().kernel_config.linear_backend | ||
| ) | ||
| if linear_backend not in {"auto", "aiter"}: | ||
| raise ValueError( | ||
| "ModelOpt MXFP8 to FP8 PTPC requantization requires " | ||
| "the AITER linear kernel; use --linear-backend=auto " | ||
| "or --linear-backend=aiter, got " | ||
| f"--linear-backend={linear_backend}." | ||
| ) |
There was a problem hiding this comment.
I don't think this should be done here.
should handle this already| def get_quant_method( | ||
| self, layer: torch.nn.Module, prefix: str | ||
| ) -> "QuantizeMethodBase | None": | ||
| if current_platform.is_rocm() and isinstance(layer, LinearBase): |
There was a problem hiding this comment.
Why the check on current_platform.is_rocm()? IMO such re-quantization logic should be accelerator-agnostic, and fail gracefully in case no requantization implementation / backend implementation is available for a given accelerator.
| target_method = OnlineQuantizationConfig(args).get_quant_method( | ||
| layer, prefix | ||
| ) | ||
| if isinstance(target_method, Fp8PtpcOnlineLinearMethod): |
There was a problem hiding this comment.
else? not supported error?
| @@ -1717,6 +1726,40 @@ def get_name(self) -> QuantizationMethods: | |||
| def get_supported_act_dtypes(self) -> list[torch.dtype]: | |||
| return [torch.bfloat16] | |||
|
|
|||
| def get_quant_method( | |||
There was a problem hiding this comment.
I am not sure about most of the logic here, seems it could belong elsewhere
|
This pull request has merge conflicts that must be resolved before it can be |
Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com> Signed-off-by: Tan Pin Siang <tanpinsiang@gmail.com>
85199f1 to
32f6260
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
|
Related: #51392 -> once this is merged, we could relax the constraint of online quantization not colliding with the base quant method, and allow online quantization to override the base quant method (if implemented, otherwise error) as done in this PR. |
2 participants
Summary
This PR lets ROCm users requantize serialized MXFP8 linear weights to FP8 per-token/per-channel (PTPC) at model startup:
The supported activation flag is:
The implementation is source-format-aware and model-agnostic. It adds no model-specific quantization class or model-name routing.
ModelOptnames vLLM's existing checkpoint-format handler:ModelOptMxFp8ConfigandModelOptMxFp8LinearMethod. MiniMax-M3 checkpoints tagged withquant_method: mxfp8use the same serialized E4M3-value plus E8M0 block-scale format, so vLLM routes them through this handler. This does not mean the checkpoint was created by NVIDIA ModelOpt.The PTPC target is generic, but this PR currently provides source reconstruction for the ModelOpt MXFP8 loader.
Implementation
LinearRequantizationSourceis the small generic interface between a serialized source method and an online quantizer.ModelOptMxFp8LinearMethod.dequantize_weight()reconstructs BF16 from the checkpoint's MXFP8 values and E8M0 scales.Fp8PtpcOnlineLinearMethodperforms PTPC quantization and AITER preshuffling from the reconstructed weight.QuantizationConfigArgs; there is nomodel_typecheck or hard-coded layer-name list.Usage
Without the final flag, the checkpoint keeps its native MXFP8 linear path.
Validation
environment
b6754f536e9cb40a1e117705a44d74218873b43785199f11c15b8860429c2eb8f37ff2b8dec143cfStartup and backend checks
auto, no PTPC flagRocmDotScaledMxfp8LinearKernelauto,fp8_per_channelAiterPreshuffledPerTokenFp8ScaledMMLinearKernelauto, no PTPC flagemulation,fp8_per_channelemulation, no PTPC flagEmulationMxfp8LinearKernelQuality: does PTPC preserve native MXFP8 accuracy?
The required paired GSM8K run used
lm-eval[api]==0.4.12, 25-shot, 200 samples, chat template, temperature 0, max 4096 tokens, concurrency 200, and the same dataset revision.Because this two-sample difference lacked per-sample records, a separate full-dataset run repeated the protocol on all 1,319 samples with sample logging:
C64: does PTPC improve over current native MXFP8?
PTPC improved throughput by 4.149%; every reported request-latency metric also
improved in this run.
C256: does PTPC improve over current native MXFP8?
PTPC improved throughput by 4.738%.
Migration: is PTPC better than forced emulation?
This was a direct concurrent C256 comparison, not a percentage inferred from
other runs. Both arms completed 2,560/2,560 with zero failures and no
measured-phase JIT/autotuning.
PTPC improved throughput by 2.846%