[NVIDIA] Fix Llama4 Scout FP4 functionality issues

Fix the weight loading issues and accuray issues when using the NVIDIA
ModelOpt Llama4 Scout FP4 model.

Signed-off-by: Po-Han Huang <[email protected]>

Author: nvpohanh

vllm/model_executor/layers/fused_moe/layer.py

@@ -874,6 +874,14 @@ def _load_per_tensor_weight_scale(self, shard_id: str,
         elif shard_id == "w2":
             param_data[expert_id] = loaded_weight
 
+    def _load_w13_weight_scale(self, shard_dim: int,
+                               loaded_weight: torch.Tensor,
+                               param: torch.Tensor, tp_rank: int):
+        shard_size = param.shape[shard_dim]
+        loaded_weight = loaded_weight.narrow(shard_dim, shard_size * tp_rank,
+                                             shard_size)
+        param.copy_(loaded_weight)
+
     def _load_model_weight_or_group_weight_scale(self,
                                                  shard_dim: int,
                                                  expert_data: torch.Tensor,
@@ -1123,7 +1131,12 @@ def weight_loader(self,
                 "weight_scale_2" in weight_name if uses_weight_scale_2 else
                 "weight_scale" in weight_name) or "input_scale" in weight_name
 
-            if per_tensor_conditions:
+            if "w13_weight_scale" in weight_name:
+                self._load_w13_weight_scale(shard_dim=shard_dim,
+                                            loaded_weight=loaded_weight,
+                                            param=param,
+                                            tp_rank=self.tp_rank)
+            elif per_tensor_conditions:
                 self._load_per_tensor_weight_scale(
                     shard_id=shard_id,
                     param=param,

vllm/model_executor/layers/quantization/modelopt.py

@@ -778,8 +778,6 @@ def process_weights_after_loading(self, layer: Module) -> None:
         # Swizzle the weight blockscale.
         # contracting dimension is input dimension
         # block_size = 16;
-        assert (layer.weight_scale.shape[1] % 16 == 0), (
-            "Expected weight_scale.dim(1) to be divisible by 16")
         assert (layer.weight_scale.dtype == torch.float8_e4m3fn), (
             "Weight Block scale must be represented as FP8-E4M3")
         swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)

vllm/model_executor/models/llama4.py

@@ -363,11 +363,23 @@ def load_moe_expert_weights(
                 continue
             param = params_dict[full_param_name]
             weight_loader = param.weight_loader
+
+            # Helper function to check if the weight is FP4.
+            # We use uint8 to store FP4 weights for now.
+            def is_fp4_weight(weight):
+                return weight.dtype == torch.uint8
+
             if fused:
                 if "w13" in full_param_name:
                     shard_idx = 0 if shard_id == "w1" else 1
                     new_loaded_weight = new_loaded_weight[shard_idx]
-                new_loaded_weight = new_loaded_weight.transpose(-1, -2)
+
+                # Only transpose for non-FP4 weights
+                # FP4 weights are already in the correct format and
+                # shouldn't be transposed here.
+                if not is_fp4_weight(new_loaded_weight):
+                    new_loaded_weight = new_loaded_weight.transpose(-1, -2)
+
                 layer_idx = extract_layer_index(name)
                 # EP mapping
                 expert_map = self.layers[
@@ -382,6 +394,11 @@ def load_moe_expert_weights(
             else:
                 # TODO: add EP support for non fused weights
                 pass
+
+            # Only transpose for FP4 weights
+            if is_fp4_weight(new_loaded_weight):
+                new_loaded_weight = new_loaded_weight.transpose(-1, -2)
+
             weight_loader(param,
                           new_loaded_weight,
                           full_param_name,
@@ -447,6 +464,7 @@ def load_weights(self, weights: Iterable[tuple[str,
                 param = params_dict[name]
                 weight_loader = getattr(param, "weight_loader",
                                         default_weight_loader)
+
                 if weight_loader == default_weight_loader:
                     weight_loader(param, loaded_weight)
                 else:
@@ -491,11 +509,17 @@ def load_weights(self, weights: Iterable[tuple[str,
                             else:
                                 shard_id = "w1"
 
+                            # Transpose if the weights are FP8 or FP4.
+                            if loaded_weight.dtype == torch.uint8 \
+                                or loaded_weight.dtype == torch.float8_e4m3fn:
+                                loaded_weight = loaded_weight.transpose(-1, -2)
+
                             weight_loader(param,
                                           loaded_weight,
                                           name,
                                           shard_id=shard_id,
                                           expert_id=0)
+
                         else:
                             # Regular weight loader (handles both
                             # param.weight_loader and default_weight_loader)
@@ -560,23 +584,32 @@ def permute_qk_weight_for_rotary(
         loaded_weight: torch.Tensor,
     ) -> tuple[str, torch.Tensor]:
 
+        # Helper function to permute the weight's channels
         def permute(w: torch.Tensor, n_heads: int):
             attn_in = self.config.head_dim * n_heads
             attn_out = self.config.hidden_size
 
+            # If the weight is FP4 packed as uint8, we need to divide attn_out
+            # by 2.
+            if w.dtype == torch.uint8 and w.shape[1] * 2 == attn_out:
+                attn_out = attn_out // 2
+
             return w.view(n_heads, attn_in // n_heads // 2, 2,
                           attn_out).transpose(1, 2).reshape(attn_in, attn_out)
 
         modules = name.split(".")
 
-        # rotary embeds should be sliced
-        if ("wk" in modules or "k_proj" in modules) \
-           and modules[-1] == "weight":
-            loaded_weight = permute(loaded_weight,
-                                    self.config.num_key_value_heads)
-        elif ("wq" in modules or "q_proj" in modules) \
-                and modules[-1] == "weight":
-            loaded_weight = permute(loaded_weight,
-                                    self.config.num_attention_heads)
+        # Permute Q/K weights and weight block scales for rotary embedding
+        is_weight = modules[-1] == "weight"
+        is_nvfp4_weight_scale = (modules[-1] == "weight_scale" and
+                                 loaded_weight.dtype == torch.float8_e4m3fn)
+
+        if is_weight or is_nvfp4_weight_scale:
+            if ("wk" in modules or "k_proj" in modules):
+                loaded_weight = permute(loaded_weight,
+                                        self.config.num_key_value_heads)
+            elif ("wq" in modules or "q_proj" in modules):
+                loaded_weight = permute(loaded_weight,
+                                        self.config.num_attention_heads)
 
         return name, loaded_weight