[NVIDIA] Fix BF16 FlashInfer TRTLLM MoE RL weight update

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -301,6 +301,9 @@ def __init__(
 
         assert intermediate_size % self.moe_tp_size == 0
         self.intermediate_size_per_partition = intermediate_size // self.moe_tp_size
+        self.intermediate_size_per_partition_unpadded = (
+            self.intermediate_size_per_partition
+        )
         self.reduce_results = reduce_results
         self.use_presharded_weights = use_presharded_weights
 
@@ -612,14 +615,21 @@ def _load_w13(
 
         if shard_id in {"w1", "w3"} and self.moe_runner_config.is_gated:
             # non-fused version
-            shard_size = expert_data.shape[shard_dim] // 2
+            param_shard_size = expert_data.shape[shard_dim] // 2
         elif shard_id in {"w13"} or (
             shard_id in {"w1", "w3"} and not self.moe_runner_config.is_gated
         ):
             # fused version
-            shard_size = expert_data.shape[shard_dim]
+            param_shard_size = expert_data.shape[shard_dim]
         else:
             raise NotImplementedError
+        weight_shard_size = param_shard_size
+        if self.use_padded_loading and not self.use_presharded_weights and not is_bias:
+            if not (shard_id == "w13" and self.moe_runner_config.is_gated):
+                # FlashInfer TRTLLM may pad the destination shard to 128, while
+                # the checkpoint tensor is still unpadded. Slice the source by
+                # its real per-rank size, then copy into the padded destination.
+                weight_shard_size = loaded_weight.shape[shard_dim] // self.moe_tp_size
 
         # Narrow parameter and load.
         # w1, gate_proj: Load into first logical weight of w13.
@@ -629,32 +639,77 @@ def _load_w13(
         if (
             (switch_w13 and shard_id == "w1") or (not switch_w13 and shard_id == "w3")
         ) and self.moe_runner_config.is_gated:
-            start = shard_size
+            start = param_shard_size
         else:
             start = 0
 
+        # Hot-update can load a fused W13 tensor, while FlashInfer TRTLLM stores
+        # the destination as two separately padded halves, e.g.
+        # [W3][W3 pad][W1][W1 pad]. Copy each half separately so W1 data does
+        # not spill into W3 padding.
+        if (
+            self.use_flashinfer_trtllm_moe
+            and shard_id == "w13"
+            and self.moe_runner_config.is_gated
+            and not is_bias
+        ):
+            # A fused W13 source is laid out as two logical halves, W1 then W3,
+            # while FlashInfer TRTLLM stores W13 as W3 then W1. Shard each
+            # source half independently before copying into the padded layout.
+            param_half_size = expert_data.shape[shard_dim] // 2
+            source_half_size = loaded_weight.shape[shard_dim] // 2
+            weight_half_shard_size = source_half_size
+            if not self.use_presharded_weights:
+                weight_half_shard_size //= self.moe_tp_size
+            weight_start = (
+                0 if self.use_presharded_weights else weight_half_shard_size * tp_rank
+            )
+
+            for param_start, source_start in (
+                (0, source_half_size + weight_start),
+                (param_half_size, weight_start),
+            ):
+                half_expert_data, half_loaded_weight = (
+                    narrow_padded_param_and_loaded_weight(
+                        expert_data,
+                        loaded_weight,
+                        param_start,
+                        source_start,
+                        shard_dim,
+                        param_half_size,
+                        narrow_weight=True,
+                        weight_shard_size=weight_half_shard_size,
+                    )
+                )
+                half_loaded_weight = _maybe_copy_weight_view_before_h2d(
+                    half_loaded_weight
+                )
+                half_expert_data.copy_(half_loaded_weight)
+            return
+
         if self.use_padded_loading:
             if _is_cpu and is_bias:
                 shard_dim = 1
             expert_data, loaded_weight = narrow_padded_param_and_loaded_weight(
                 expert_data,
                 loaded_weight,
                 start,
-                shard_size * tp_rank,
+                weight_shard_size * tp_rank,
                 shard_dim,
-                shard_size,
-                not self.use_presharded_weights,
+                param_shard_size,
+                narrow_weight=not self.use_presharded_weights,
+                weight_shard_size=weight_shard_size,
             )
         else:
             if not self.use_presharded_weights:
                 if not is_bias and self.use_triton_kernels:
                     # do not transpose for bias
                     loaded_weight = loaded_weight.transpose(-2, -1)
                 loaded_weight = loaded_weight.narrow(
-                    shard_dim, shard_size * tp_rank, shard_size
+                    shard_dim, param_shard_size * tp_rank, param_shard_size
                 )
 
-            expert_data = expert_data.narrow(shard_dim, start, shard_size)
+            expert_data = expert_data.narrow(shard_dim, start, param_shard_size)
         loaded_weight = _maybe_copy_weight_view_before_h2d(loaded_weight)
         expert_data.copy_(loaded_weight)
 
@@ -699,11 +754,16 @@ def _load_w2(
         if is_bias:
             # this expert_data is a bias, not weight,
             # for w2_weight_bias in TP, it does not need to be sharded
-            shard_size = expert_data.shape[-1]
+            param_shard_size = expert_data.shape[-1]
         else:
             # this parameter is a weight matrix
             # for w2 in TP, it shards the input_features, i.e., shard_dim=2
-            shard_size = expert_data.shape[shard_dim]
+            param_shard_size = expert_data.shape[shard_dim]
+        weight_shard_size = param_shard_size
+        if self.use_padded_loading and not self.use_presharded_weights and not is_bias:
+            # FlashInfer TRTLLM may pad the destination shard to 128, while the
+            # checkpoint tensor is still unpadded.
+            weight_shard_size = loaded_weight.shape[shard_dim] // self.moe_tp_size
 
         if self.use_padded_loading:
             if _is_cpu and is_bias:
@@ -712,17 +772,18 @@ def _load_w2(
                 expert_data,
                 loaded_weight,
                 0,  # param_data_start
-                shard_size * tp_rank,
+                weight_shard_size * tp_rank,
                 shard_dim,
-                shard_size,
-                not self.use_presharded_weights,
+                param_shard_size,
+                narrow_weight=not self.use_presharded_weights,
+                weight_shard_size=weight_shard_size,
             )
         else:
             if not is_bias and not self.use_presharded_weights:
                 if self.use_triton_kernels:
                     loaded_weight = loaded_weight.transpose(-2, -1)
                 loaded_weight = loaded_weight.narrow(
-                    shard_dim, shard_size * tp_rank, shard_size
+                    shard_dim, param_shard_size * tp_rank, param_shard_size
                 )
 
         # w2, down_proj: Load into only logical weight of w2.

python/sglang/srt/layers/quantization/unquant.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import math
 from enum import Enum
 from typing import TYPE_CHECKING, List, Optional
 
@@ -59,6 +60,47 @@
 _is_npu = is_npu()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 
+_FLASHINFER_TRTLLM_BF16_LAYOUT_META_ATTR = "_sglang_flashinfer_trtllm_bf16_layout_meta"
+_FLASHINFER_TRTLLM_BF16_LAYOUT_CANONICAL = "canonical"
+_FLASHINFER_TRTLLM_BF16_LAYOUT_PACKED = "flashinfer_trtllm_bf16"
+
+
+def _get_flashinfer_trtllm_bf16_layout_meta(param: Parameter) -> dict:
+    return getattr(param, _FLASHINFER_TRTLLM_BF16_LAYOUT_META_ATTR, {})
+
+
+def _set_flashinfer_trtllm_bf16_layout_meta(
+    param: Parameter,
+    *,
+    canonical_shape: tuple[int, ...],
+    storage_layout: str,
+    packed_shape: tuple[int, ...] | None = None,
+) -> None:
+    meta = dict(_get_flashinfer_trtllm_bf16_layout_meta(param))
+    meta["canonical_shape"] = tuple(canonical_shape)
+    meta["storage_layout"] = storage_layout
+    if packed_shape is not None:
+        meta["packed_shape"] = tuple(packed_shape)
+    setattr(param, _FLASHINFER_TRTLLM_BF16_LAYOUT_META_ATTR, meta)
+
+
+def _zero_flashinfer_trtllm_bf16_padding(layer: torch.nn.Module) -> None:
+    padded_size = getattr(layer, "intermediate_size_per_partition", None)
+    unpadded_size = getattr(layer, "intermediate_size_per_partition_unpadded", None)
+    if padded_size is None or unpadded_size is None or unpadded_size >= padded_size:
+        return
+
+    w13_weight = layer.w13_weight.data
+    w2_weight = layer.w2_weight.data
+    if w13_weight.dim() == 3:
+        w13_weight[:, unpadded_size:padded_size, :].zero_()
+        if layer.moe_runner_config.is_gated:
+            w13_weight[:, padded_size + unpadded_size : 2 * padded_size, :].zero_()
+
+    if w2_weight.dim() == 3:
+        w2_weight[:, :, unpadded_size:padded_size].zero_()
+
+
 if _use_aiter:
     from aiter.ops.shuffle import shuffle_weight
     from aiter.tuned_gemm import tgemm
@@ -299,6 +341,12 @@ def create_weights(
         )
         layer.register_parameter("w13_weight", w13_weight)
         set_weight_attrs(w13_weight, extra_weight_attrs)
+        if self.use_flashinfer_trtllm_moe:
+            _set_flashinfer_trtllm_bf16_layout_meta(
+                w13_weight,
+                canonical_shape=tuple(w13_weight.data.shape),
+                storage_layout=_FLASHINFER_TRTLLM_BF16_LAYOUT_CANONICAL,
+            )
 
         if self.with_bias:
             w13_weight_bias = torch.nn.Parameter(
@@ -321,6 +369,12 @@ def create_weights(
         )
         layer.register_parameter("w2_weight", w2_weight)
         set_weight_attrs(w2_weight, extra_weight_attrs)
+        if self.use_flashinfer_trtllm_moe:
+            _set_flashinfer_trtllm_bf16_layout_meta(
+                w2_weight,
+                canonical_shape=tuple(w2_weight.data.shape),
+                storage_layout=_FLASHINFER_TRTLLM_BF16_LAYOUT_CANONICAL,
+            )
 
         if self.with_bias:
             w2_weight_bias = torch.nn.Parameter(
@@ -374,6 +428,31 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
 
         # Reorder rows of W1 for fused gated activation
         if self.use_flashinfer_trtllm_moe:
+            w13_meta = _get_flashinfer_trtllm_bf16_layout_meta(layer.w13_weight)
+            w2_meta = _get_flashinfer_trtllm_bf16_layout_meta(layer.w2_weight)
+            w13_layout = w13_meta.get("storage_layout")
+            w2_layout = w2_meta.get("storage_layout")
+            if (
+                w13_layout == _FLASHINFER_TRTLLM_BF16_LAYOUT_PACKED
+                and w2_layout == _FLASHINFER_TRTLLM_BF16_LAYOUT_PACKED
+            ):
+                return
+            if (
+                w13_layout != _FLASHINFER_TRTLLM_BF16_LAYOUT_CANONICAL
+                or w2_layout != _FLASHINFER_TRTLLM_BF16_LAYOUT_CANONICAL
+            ):
+                raise RuntimeError(
+                    "FlashInfer TRT-LLM BF16 MoE weights must have canonical "
+                    "layout metadata before process_weights_after_loading."
+                )
+
+            # Ensure padded slots are zero before packing. For example, hot
+            # update may restore params back to canonical shape with a reshape,
+            # not an inverse unpack, so padding can contain old packed values.
+            _zero_flashinfer_trtllm_bf16_padding(layer)
+
+            canonical_shape_w13 = tuple(layer.w13_weight.data.shape)
+            canonical_shape_w2 = tuple(layer.w2_weight.data.shape)
             # The cached indices are GPU tensors. Colocated weight offloading
             # can release their backing memory between reloads, so rebuild them
             # once per post-processing cycle.
@@ -442,6 +521,19 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             layer.w2_weight.data = layer.w2_weight.data.reshape(
                 layer.num_local_experts, *new_shape_w2
             )
+            _set_flashinfer_trtllm_bf16_layout_meta(
+                layer.w13_weight,
+                canonical_shape=canonical_shape_w13,
+                storage_layout=_FLASHINFER_TRTLLM_BF16_LAYOUT_PACKED,
+                packed_shape=tuple(layer.w13_weight.data.shape),
+            )
+            _set_flashinfer_trtllm_bf16_layout_meta(
+                layer.w2_weight,
+                canonical_shape=canonical_shape_w2,
+                storage_layout=_FLASHINFER_TRTLLM_BF16_LAYOUT_PACKED,
+                packed_shape=tuple(layer.w2_weight.data.shape),
+            )
+
         if _is_npu:
             # The kernels set the dispatcher output dtype themselves -- they are
             # the ones that know what their gmms expect. NPUUnquantMoEMethod
@@ -458,41 +550,59 @@ def maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load(
         param: torch.nn.Parameter,
         weight_name: str,
     ) -> None:
-        """Restore canonical BF16 MoE load shapes before hot weight copy.
+        """Prepare BF16 MoE params for canonical hot-update loading.
 
         The flashinfer TRT-LLM BF16 postprocess reshapes expert weights into
-        block layout. During weight update, checkpoint tensors are in
-        canonical layout and need a temporary shape restore for copy.
+        block layout. During weight update, checkpoint tensors are canonical
+        replacement values, so the packed storage can be discarded and reused as
+        a canonical load buffer before copy.
         """
-        if not get_moe_runner_backend().is_flashinfer_trtllm_routed():
+        if not self.use_flashinfer_trtllm_moe:
             return
 
         expected_shape = None
-        if weight_name.endswith(".experts.w13_weight"):
+        if param is getattr(layer, "w13_weight", None) or weight_name.endswith(
+            ".experts.w13_weight"
+        ):
             w13_rows = (
                 2 * layer.intermediate_size_per_partition
                 if layer.moe_runner_config.is_gated
                 else layer.intermediate_size_per_partition
             )
             expected_shape = (layer.num_local_experts, w13_rows, layer.hidden_size)
-        elif weight_name.endswith(".experts.w2_weight"):
+        elif param is getattr(layer, "w2_weight", None) or weight_name.endswith(
+            ".experts.w2_weight"
+        ):
             expected_shape = (
                 layer.num_local_experts,
                 layer.hidden_size,
                 layer.intermediate_size_per_partition,
             )
 
-        if expected_shape is None or tuple(param.data.shape) == expected_shape:
+        if expected_shape is None:
             return
 
-        expected_numel = expected_shape[0] * expected_shape[1] * expected_shape[2]
-        if param.data.numel() != expected_numel:
-            raise RuntimeError(
-                f"Cannot restore flashinfer TRT-LLM BF16 MoE weight shape for {weight_name}: "
-                f"current shape={tuple(param.data.shape)}, expected shape={expected_shape}."
-            )
+        meta = _get_flashinfer_trtllm_bf16_layout_meta(param)
+        expected_shape = tuple(meta.get("canonical_shape", expected_shape))
+        current_shape = tuple(param.data.shape)
+
+        if current_shape != expected_shape:
+            expected_numel = math.prod(expected_shape)
+            if param.data.numel() != expected_numel:
+                raise RuntimeError(
+                    f"Cannot restore flashinfer TRT-LLM BF16 MoE weight shape for {weight_name}: "
+                    f"current shape={current_shape}, expected shape={expected_shape}."
+                )
+
+            self._cache_permute_indices.clear()
+            param.data = param.data.reshape(expected_shape)
 
-        param.data = param.data.reshape(expected_shape)
+        _set_flashinfer_trtllm_bf16_layout_meta(
+            param,
+            canonical_shape=expected_shape,
+            storage_layout=_FLASHINFER_TRTLLM_BF16_LAYOUT_CANONICAL,
+            packed_shape=meta.get("packed_shape", current_shape),
+        )
 
     def _aiter_ck_moe_supported(self, layer) -> bool:
         # aiter CK fused-MoE requires intermediate_size_per_partition to be 128-aligned

python/sglang/srt/model_loader/weight_utils.py

@@ -1770,9 +1770,23 @@ def narrow_padded_param_and_loaded_weight(
     dim,
     shard_size,
     narrow_weight=True,
+    weight_shard_size=None,
 ):
-    actual_shard_size = get_actual_shard_size(
-        shard_size, weight_start, loaded_weight.size(dim)
+    """Return matching param/weight slices and zero padded param tail.
+
+    shard_size is the size of the destination param shard, including any
+    padding. weight_shard_size is the size of the incoming weight shard before
+    destination padding; it defaults to shard_size for existing callers.
+    """
+    if weight_shard_size is None:
+        weight_shard_size = shard_size
+
+    if not narrow_weight:
+        weight_start = 0
+
+    actual_shard_size = min(
+        shard_size,
+        get_actual_shard_size(weight_shard_size, weight_start, loaded_weight.size(dim)),
     )
 
     if narrow_weight: