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 7, 21:30 UTC
[Llama4,CI] Bring back Llama-4 bug fixes, and also fix Maverick tests (#35033)
Signed-off-by: Eldar Kurtic <you@example.com> Co-authored-by: Eldar Kurtic <you@example.com>
Eldar Kurtić committed
Feb 23, 2026 at 15:04 UTC
1e8438a89a6453a6b1ba28798bb9c51d6364ed96
2 files changed
+31
-70
tests/models/multimodal/generation/test_maverick.py
+2
-2
@@ -305,10 +305,10 @@ def create_text_model_weights(text_config: dict[str, Any]) -> dict[str, torch.Te
305
306
# Self-attention weights (separate q, k, v projections)
307
weights[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(
308
- hidden_size, num_attention_heads * head_dim, dtype=torch.bfloat16
308
+ num_attention_heads * head_dim, hidden_size, dtype=torch.bfloat16
309
)
310
weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(
311
- hidden_size, num_key_value_heads * head_dim, dtype=torch.bfloat16
311
+ num_key_value_heads * head_dim, hidden_size, dtype=torch.bfloat16
312
)
313
weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(
314
num_key_value_heads * head_dim, hidden_size, dtype=torch.bfloat16
vllm/model_executor/models/llama4.py
+29
-68
@@ -44,9 +44,6 @@ 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
-)
47
from vllm.model_executor.layers.rotary_embedding import get_rope
48
from vllm.model_executor.model_loader.weight_utils import (
49
default_weight_loader,
@@ -831,74 +828,38 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
828
name: str,
829
loaded_weight: torch.Tensor,
830
) -> tuple[str, torch.Tensor]:
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]
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
855
)
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)
856
+ loaded_weight = (
857
+ loaded_weight.view(n_heads, f_out // n_heads // 2, 2, f_in)
858
.transpose(1, 2)
867
- .reshape(attn_in, attn_out)
859
+ .reshape(f_out, f_in)
860
)
861
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
- )
862
+ if original_ndim == 1:
863
+ loaded_weight = loaded_weight.squeeze(-1)
864
865
return name, loaded_weight