A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 11:15 UTC
[Bugfix] Enable attn quantization of Llama-4 by correctly permuting scales for rope (int8, fp8) (#34243)
Signed-off-by: Your Name <you@example.com> Co-authored-by: Your Name <you@example.com>
Eldar Kurtić committed
Feb 11, 2026 at 19:24 UTC
11c7ace340610e0be376d531b677bcee1ae84ad4
1 file changed
+29
-5
vllm/model_executor/models/llama4.py
+29
-5
@@ -44,6 +44,9 @@ from vllm.model_executor.layers.linear import (
44
RowParallelLinear,
45
)
46
from vllm.model_executor.layers.quantization import QuantizationConfig
47
+from vllm.model_executor.layers.quantization.compressed_tensors import (
48
+ compressed_tensors as ct,
49
+)
50
from vllm.model_executor.layers.rotary_embedding import get_rope
51
from vllm.model_executor.model_loader.weight_utils import (
52
default_weight_loader,
@@ -829,11 +832,20 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
832
loaded_weight: torch.Tensor,
833
) -> tuple[str, torch.Tensor]:
834
# Helper function to permute the weight's channels
832
- def permute(w: torch.Tensor, n_heads: int, is_weight_scale: bool):
835
+ def permute(
836
+ w: torch.Tensor,
837
+ n_heads: int,
838
+ is_nvfp4_weight_scale: bool,
839
+ is_ct_int8_or_fp8_weight_scale: bool,
840
+ ):
841
# Calculate the expected shape of the weight.
842
# Do not rely on w's shape, as it may be in another layout.
843
attn_in = self.config.head_dim * n_heads
836
- attn_out = self.config.hidden_size
844
+ attn_out = (
845
+ self.config.hidden_size
846
+ if not is_ct_int8_or_fp8_weight_scale
847
+ else w.shape[-1]
848
+ )
849
850
# If the weight is FP4 packed as uint8, we need to divide attn_out
851
# by 2.
@@ -844,7 +856,7 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
856
# block size, which is currently 16.
857
elif (
858
w.dtype == torch.float8_e4m3fn
847
- and is_weight_scale
859
+ and is_nvfp4_weight_scale
860
and w.shape[1] * 16 == attn_out
861
):
862
attn_out = attn_out // 16
@@ -862,19 +874,31 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
874
is_nvfp4_weight_scale = (
875
modules[-1] == "weight_scale" and loaded_weight.dtype == torch.float8_e4m3fn
876
)
865
-
866
- if is_weight or is_nvfp4_weight_scale:
877
+ is_ct_int8_or_fp8_weight_scale = False
878
+ if modules[-1] == "weight_scale" and isinstance(
879
+ self.model.quant_config, ct.CompressedTensorsConfig
880
+ ):
881
+ from compressed_tensors import CompressionFormat
882
+
883
+ is_ct_int8_or_fp8_weight_scale = self.model.quant_config.quant_format in [
884
+ CompressionFormat.int_quantized.value,
885
+ CompressionFormat.float_quantized.value,
886
+ ] and loaded_weight.dtype in [torch.float16, torch.bfloat16, torch.float32]
887
+
888
+ if is_weight or is_nvfp4_weight_scale or is_ct_int8_or_fp8_weight_scale:
889
if "wk" in modules or "k_proj" in modules:
890
loaded_weight = permute(
891
loaded_weight,
892
self.config.num_key_value_heads,
893
is_nvfp4_weight_scale,
894
+ is_ct_int8_or_fp8_weight_scale,
895
)
896
elif "wq" in modules or "q_proj" in modules:
897
loaded_weight = permute(
898
loaded_weight,
899
self.config.num_attention_heads,
900
is_nvfp4_weight_scale,
901
+ is_ct_int8_or_fp8_weight_scale,
902
)
903
904
return name, loaded_weight