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, 00:00 UTC
Revert "[Llama4,Quantization] Simplify and generalize logic for Q/K permutations in quantized self-attn layers " (#34997)
Lucas Wilkinson committed
Feb 20, 2026 at 20:19 UTC
0e22cd618b5da36404365518aad5a522aea008e7
1 file changed
+68
-29
vllm/model_executor/models/llama4.py
+68
-29
@@ -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,
@@ -828,38 +831,74 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
831
name: str,
832
loaded_weight: torch.Tensor,
833
) -> tuple[str, torch.Tensor]:
831
- modules = name.split(".")
832
- # Permute Q/K weights and corresponding scales for rotary embedding.
833
- # This pathway is validated against modelopt and compressed-tensors ckpts,
834
- # and for per-tensor, per-group (e.g. GPTQ), and per-channel quant schemes.
835
- # Note: permutations are not feasible only for per-block (e.g. DeepSeek 128x128)
836
- # For per-block quantization, consider not quantizing q/k_proj.
837
- is_weight = modules[-1] in ("weight", "weight_packed")
838
- is_weight_scale = (
839
- modules[-1] == "weight_scale"
840
- and loaded_weight.numel() > 1 # no need to permute per-tensor scales
841
- )
842
- is_k_proj = "wk" in modules or "k_proj" in modules
843
- is_q_proj = "wq" in modules or "q_proj" in modules
844
-
845
- if (is_weight or is_weight_scale) and (is_k_proj or is_q_proj):
846
- original_ndim = loaded_weight.ndim
847
- if original_ndim == 1:
848
- loaded_weight = loaded_weight.unsqueeze(-1)
849
-
850
- f_out, f_in = loaded_weight.shape
851
- n_heads = (
852
- self.config.num_key_value_heads
853
- if is_k_proj
854
- else self.config.num_attention_heads
834
+ # Helper function to permute the weight's channels
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
844
+ attn_out = (
845
+ self.config.hidden_size
846
+ if not is_ct_int8_or_fp8_weight_scale
847
+ else w.shape[-1]
848
)
856
- loaded_weight = (
857
- loaded_weight.view(n_heads, f_out // n_heads // 2, 2, f_in)
849
+
850
+ # If the weight is FP4 packed as uint8, we need to divide attn_out
851
+ # by 2.
852
+ if w.dtype == torch.uint8 and w.shape[1] * 2 == attn_out:
853
+ attn_out = attn_out // 2
854
+
855
+ # If the weight is a weight scale, we need to divide attn_out by
856
+ # block size, which is currently 16.
857
+ elif (
858
+ w.dtype == torch.float8_e4m3fn
859
+ and is_nvfp4_weight_scale
860
+ and w.shape[1] * 16 == attn_out
861
+ ):
862
+ attn_out = attn_out // 16
863
+
864
+ return (
865
+ w.view(n_heads, attn_in // n_heads // 2, 2, attn_out)
866
.transpose(1, 2)
859
- .reshape(f_out, f_in)
867
+ .reshape(attn_in, attn_out)
868
)
869
862
- if original_ndim == 1:
863
- loaded_weight = loaded_weight.squeeze(-1)
870
+ modules = name.split(".")
871
+
872
+ # Permute Q/K weights and weight block scales for rotary embedding
873
+ is_weight = modules[-1] == "weight"
874
+ is_nvfp4_weight_scale = (
875
+ modules[-1] == "weight_scale" and loaded_weight.dtype == torch.float8_e4m3fn
876
+ )
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