[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,17 @@ 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 +1134,14 @@ 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

@@ -810,8 +810,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 = self.swizzle_blockscale(layer.weight_scale)

vllm/model_executor/models/llama4.py

@@ -363,11 +363,22 @@ 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 +393,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,
@@ -491,11 +507,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 +582,28 @@ 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
-
-            return w.view(n_heads, attn_in // n_heads // 2, 2,
-                          attn_out).transpose(1, 2).reshape(attn_in, attn_out)
+            head_dim = w.shape[0] // n_heads
+            return (
+                w.view(n_heads, head_dim // 2, 2, w.shape[1])
+                .transpose(1, 2)
+                .reshape(w.shape[0], w.shape[1])
+            )
 
         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