[AMD] Fix weight checking for AITER-shuffled block FP8 weights (#34330)

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

@@ -66,6 +66,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
     normalize_e4m3fn_to_e4m3fnuz,
     requant_block_scale_ue8m0_for_deepgemm,
     resolve_mxfp8_dense_gemm_backend,
+    unshuffle_aiter_fp8_weight,
 )
 from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
 from sglang.srt.layers.quantization.marlin_utils_fp8 import prepare_fp8_layer_for_marlin
@@ -146,6 +147,13 @@ def _require_fp4_dtype():
     return fp4_dtype
 
 
+def unshuffle_fp8_weight(weight: torch.Tensor) -> torch.Tensor:
+    """Restore the logical layout of a backend-shuffled FP8 weight."""
+    if not _use_aiter:
+        raise RuntimeError("FP8 weight unshuffle requires AITER")
+    return unshuffle_aiter_fp8_weight(weight)
+
+
 if _use_aiter or _use_hip_int4:
     from aiter.ops.shuffle import (
         moe_shuffle_scale,
@@ -771,6 +779,7 @@ class Fp8LinearMethod(LinearMethodBase):
                 # it so a consumer that needs the row-major layout can assert
                 # instead of silently reading a permuted weight.
                 layer.aiter_bpreshuffled = True
+                layer.weight.is_shuffled = True
 
     def _process_mxfp8_linear_weight_scale(self, layer: Module) -> None:
         if not self.use_mxfp8:
@@ -1666,6 +1675,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                 layer.w2_weight.data = shuffle_weight(
                     layer.w2_weight.contiguous(), (16, 16)
                 )
+                layer.w13_weight.is_shuffled = True
+                layer.w2_weight.is_shuffled = True
             return
         elif self.use_mxfp8 and get_moe_a2a_backend().is_flashinfer_megamoe():
             from sglang.srt.layers.moe.flashinfer_megamoe import (
@@ -1711,6 +1722,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                 layer.w2_weight.data = shuffle_weight(
                     layer.w2_weight.contiguous(), (16, 16)
                 )
+                layer.w13_weight.is_shuffled = True
+                layer.w2_weight.is_shuffled = True
         elif _use_aiter:
             # Pre-shuffle weights
             t = shuffle_weight(layer.w13_weight, (16, 16))
@@ -1719,6 +1732,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             t = shuffle_weight(layer.w2_weight, (16, 16))
             layer.w2_weight.copy_(t)
             del t
+            layer.w13_weight.is_shuffled = True
+            layer.w2_weight.is_shuffled = True
         elif _is_cpu:
             assert _is_cpu_amx_available, (
                 "Fp8MoEMethod on CPU requires that CPU has AMX support"

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

@@ -137,6 +137,27 @@ def view_aiter_fused_rms_transposed_fp8_scale(scale: torch.Tensor) -> torch.Tens
     return torch.as_strided(scale, scale.shape, (1, scale.shape[0]))
 
 
+def unshuffle_aiter_fp8_weight(weight: torch.Tensor) -> torch.Tensor:
+    """Undo AITER ``shuffle_weight(..., layout=(16, 16))`` for FP8 weights."""
+    if weight.element_size() != 1:
+        raise ValueError("AITER FP8 unshuffle requires a one-byte element type")
+
+    shape = weight.shape
+    n, k = shape[-2:]
+    if n % 16 != 0 or k % 32 != 0:
+        raise ValueError(
+            "AITER (16, 16) FP8 layout requires N % 16 == 0 and K % 32 == 0, "
+            f"got shape {tuple(shape)}"
+        )
+
+    return (
+        weight.reshape(-1, n // 16, k // 32, 2, 16, 16)
+        .permute(0, 1, 4, 2, 3, 5)
+        .contiguous()
+        .reshape(shape)
+    )
+
+
 def materialize_bpreshuffle_fp8_scale_tuple(
     value: Tuple[torch.Tensor, ...],
 ) -> Tuple[torch.Tensor, ...]:

python/sglang/srt/utils/weight_checker.py

@@ -49,10 +49,13 @@ class CheckEntry(NamedTuple):
 class QuantizedWeight(NamedTuple):
     comparable_cls: type[ComparableWeight]
     scale_name: str
+    is_shuffled: bool = False
 
 
 _NON_PERSISTENT_BUFFER_PATTERNS = (
     "cos_sin_cache",
+    "cos_cache",
+    "sin_cache",
     "inv_freq",
     "freqs_cis",
     "expert_mask_gpu",
@@ -267,12 +270,14 @@ def _build_quantized_set(model) -> Dict[str, QuantizedWeight]:
         if comparable_cls is None:
             continue
         prefix = f"{module_name}." if module_name else ""
-        own = {name for name, _ in module.named_parameters(recurse=False)}
-        for name in own:
+        own = dict(module.named_parameters(recurse=False))
+        for name, parameter in own.items():
             scale = name.replace("weight", "weight_scale_inv")
             if name.endswith("weight") and scale in own:
                 quantized_set[prefix + name] = QuantizedWeight(
-                    comparable_cls, prefix + scale
+                    comparable_cls,
+                    prefix + scale,
+                    getattr(parameter, "is_shuffled", False),
                 )
     return quantized_set
 
@@ -293,7 +298,13 @@ def _build_check_entries(
             continue  # compared via its weight's comparable
         if name in quantized_set:
             qw = quantized_set[name]
-            yield CheckEntry(name, True, qw.comparable_cls(tensor, raw[qw.scale_name]))
+            yield CheckEntry(
+                name,
+                True,
+                qw.comparable_cls(
+                    tensor, raw[qw.scale_name], is_shuffled=qw.is_shuffled
+                ),
+            )
         else:
             should_compare = name not in skip_compare_names and (
                 not _is_non_persistent_buffer_name(name)

python/sglang/srt/utils/weight_checker_comparator.py

@@ -2,7 +2,11 @@ from typing import Iterable, NamedTuple, Optional, Tuple
 
 import torch
 
-from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod, Fp8MoEMethod
+from sglang.srt.layers.quantization.fp8 import (
+    Fp8LinearMethod,
+    Fp8MoEMethod,
+    unshuffle_fp8_weight,
+)
 from sglang.srt.layers.quantization.fp8_utils import (
     block_quant_dequant,
     inverse_transform_scale_ue8m0,
@@ -48,12 +52,21 @@ class ComparableWeight:
 class Fp8BlockComparable(ComparableWeight):
     """Deepseek-style FP8 quantization."""
 
-    def __init__(self, w_q: torch.Tensor, w_s: torch.Tensor):
+    def __init__(
+        self,
+        w_q: torch.Tensor,
+        w_s: torch.Tensor,
+        is_shuffled: bool = False,
+    ):
         self.w_q = w_q
         self.w_s = w_s
+        self.is_shuffled = is_shuffled
 
     def __repr__(self) -> str:
-        return f"fp8_block(shape={tuple(self.w_q.shape)} dtype={self.w_q.dtype})"
+        return (
+            f"fp8_block(shape={tuple(self.w_q.shape)} dtype={self.w_q.dtype} "
+            f"is_shuffled={self.is_shuffled})"
+        )
 
     @staticmethod
     def _normalize_scale(w_q: torch.Tensor, w_s: torch.Tensor) -> torch.Tensor:
@@ -90,6 +103,8 @@ class Fp8BlockComparable(ComparableWeight):
         s, block_size = self._scale_and_block_size()
         for q, s_chunk in self._iter_quant_chunks(self.w_q, s, block_size[0]):
             q, s_chunk = q.cuda(), s_chunk.cuda()
+            if self.is_shuffled:
+                q = unshuffle_fp8_weight(q)
             yield (
                 block_quant_dequant(q, s_chunk, block_size, dtype=torch.bfloat16),
                 block_quant_dequant(
@@ -99,7 +114,10 @@ class Fp8BlockComparable(ComparableWeight):
 
     def dequantize(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
         s, block_size = self._scale_and_block_size()
-        return block_quant_dequant(self.w_q, s, block_size, dtype=dtype)
+        w_q = self.w_q
+        if self.is_shuffled:
+            w_q = unshuffle_fp8_weight(w_q)
+        return block_quant_dequant(w_q, s, block_size, dtype=dtype)
 
 
 class RawComparable(ComparableWeight):