A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][AITER][Quark] Tag per-channel FP8 weights as PER_CHANNEL so AITER pre-shuffled GEMM is selected (#44626)
Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com>
xaguilar-amd committed
Jun 17, 2026 at 16:05 UTC
0b131b16c93308df90d013c5ac24c6fa25bd863b
1 file changed
+7
-3
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
+7
-3
@@ -16,8 +16,8 @@ from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
16
from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
GroupShape,
18
kFp8DynamicTokenSym,
19
+ kFp8StaticChannelSym,
20
kFp8StaticTensorSym,
20
- kFp8StaticTokenSym,
21
)
22
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
23
normalize_e4m3fn_to_e4m3fnuz,
@@ -49,13 +49,17 @@ class QuarkW8A8Fp8(QuarkScheme):
49
per_token_activation = (
50
not self.is_static_input_scheme and self.input_qscheme == "per_channel"
51
)
52
- per_token_weight = self.weight_qscheme == "per_channel"
52
+ per_channel_weight = self.weight_qscheme == "per_channel"
53
54
self.activation_quant_key = (
55
kFp8DynamicTokenSym if per_token_activation else kFp8StaticTensorSym
56
)
57
+ # A per-output-channel weight scale is one fp32 value per weight row
58
+ # (length N). Tag it as ``GroupShape.PER_CHANNEL`` to match the
59
+ # canonical compressed-tensors CHANNEL strategy, so kernel selection
60
+ # (e.g. AITER's pre-shuffled FP8 GEMM) treats it uniformly.
61
self.weight_quant_key = (
58
- kFp8StaticTokenSym if per_token_weight else kFp8StaticTensorSym
62
+ kFp8StaticChannelSym if per_channel_weight else kFp8StaticTensorSym
63
)
64
self.out_dtype = torch.get_default_dtype()
65
self.input_dtype = get_current_vllm_config().model_config.dtype