Enable GFX12 Preshuffle Weights

atom/model_ops/fused_moe_triton.py

@@ -29,16 +29,15 @@
     from aiter.ops.triton.moe.moe_routing.routing import routing
     from aiter.ops.triton.moe.moe_op_gemm_a8w4 import (
         moe_gemm_a8w4,
-        swizzle_scales as swizzle_scales_a8w4,
     )
     from aiter.ops.triton.moe.moe_op_gemm_a16w4 import (
         moe_gemm_a16w4,
     )
     from aiter.ops.triton.moe.moe_op_gemm_a4w4 import (
         moe_gemm_a4w4,
         mxfp4_quant,
-        swizzle_scales as swizzle_scales_cdna4,
     )
+    from aiter.ops.triton.utils.shuffle import shuffle_scale_moe
     from aiter.ops.triton.moe.quant_moe import downcast_to_static_fp8
 
 from atom.model_ops.moe import MoEActivationQuant
@@ -51,10 +50,10 @@ def _swizzle_scales_for_kernel(scale, act_quant: MoEActivationQuant):
     BF16/FP4 (a16w4/a4w4): CDNA4 swizzle on gfx942/gfx950, no swizzle elsewhere.
     """
     if act_quant == MoEActivationQuant.FP8:
-        return swizzle_scales_a8w4(scale)
+        return shuffle_scale_moe(scale)
     # TODO: move arch dispatch into aiter's a4w4/a16w4 swizzle_scales (like a8w4)
     if get_arch() in ("gfx942", "gfx950"):
-        return swizzle_scales_cdna4(scale), "CDNA4_SCALE"
+        return shuffle_scale_moe(scale), "CDNA4_SCALE"
     return scale, None
 
 

atom/model_ops/linear.py

@@ -21,7 +21,11 @@
 from aiter.dist.parallel_state import get_tp_group
 from aiter.jit.utils.torch_guard import torch_compile_guard
 from aiter.tuned_gemm import tgemm
-from aiter.utility import fp4_utils
+from aiter.ops.triton.utils.shuffle import (
+    collapse_mxfp4_gemm_scale,
+    quant_mxfp4_act_preshuffle,
+    shuffle_scale_gemm_e8m0,
+)
 from atom.config import QuantizationConfig, get_current_atom_config
 from atom.quant_spec import LayerQuantConfig, should_skip_online_quant
 from atom.model_ops.utils import (
@@ -197,30 +201,22 @@ def gemm_a4w4_quant(
             device=x.device,
         )
         if x_scale is None:
-            quant_func = get_hip_quant(QuantType.per_1x32)
-            x, x_scale = quant_func(
-                x,
-                quant_dtype=params_dtype,
-                shuffle=(m >= MXFP4_QUANT_BLOCK_SIZE),
+            # quant with no shuffle + manual shuffle + collapse
+            x, x_scale = quant_mxfp4_act_preshuffle(
+                x, params_dtype, m, MXFP4_QUANT_BLOCK_SIZE
             )
         else:
-            x_scale = x_scale.view(torch.float8_e8m0fnu)
-            x = x.view(torch.float4_e2m1fn_x2)
-
-        if m >= MXFP4_QUANT_BLOCK_SIZE:
-            x_scale = x_scale.view(torch.uint8).view(
-                x_scale.shape[0] // MXFP4_QUANT_BLOCK_SIZE, -1
+            # collapse in arch-aware layout
+            x_scale = collapse_mxfp4_gemm_scale(
+                x_scale.view(torch.float8_e8m0fnu), MXFP4_QUANT_BLOCK_SIZE, rows_valid=m
             )
-        else:
-            x_scale = x_scale[:m, ...].view(torch.uint8)
+            x = x.view(torch.float4_e2m1fn_x2)
 
         y = gemm_afp4wfp4_preshuffle(
             x.view(torch.uint8),
             weight.view(torch.uint8).view(weight.shape[0] // 16, -1),
             x_scale,
-            weight_scale.view(torch.uint8).view(
-                weight_scale.shape[0] // MXFP4_QUANT_BLOCK_SIZE, -1
-            ),
+            collapse_mxfp4_gemm_scale(weight_scale, MXFP4_QUANT_BLOCK_SIZE),
             y=y,
         )
     # Default AITER path: quantize/shuffle into the layout expected by gemm_a4w4
@@ -767,13 +763,22 @@ def process_weights_after_loading(self):
                     need_shuffle = True
             if need_shuffle:
                 if self.weight.dim() == 2:
-                    shuffle_weights(self.weight)
+                    if use_triton_gemm():
+                        # arch-aware weight shuffle
+                        from aiter.ops.triton.utils.shuffle import (
+                            shuffle_weight as _triton_shuffle_weight,
+                        )
+
+                        self.weight.data = _triton_shuffle_weight(self.weight.data)
+                        self.weight.is_shuffled = True
+                    else:
+                        shuffle_weights(self.weight)
                 # self.weight_scale.data = fp4_utils.e8m0_shuffle(self.weight_scale.data)
-        # shuffle weight scale once so no reshuffling for every gemm
+        # shuffle the weight scale once
         if self.quant_type == QuantType.per_1x32 and (
             self.params_dtype != dtypes.fp4x2 or not use_fp4_non_shuffle_triton_gemm()
         ):
-            self.weight_scale.data = fp4_utils.e8m0_shuffle(self.weight_scale.data)
+            self.weight_scale.data = shuffle_scale_gemm_e8m0(self.weight_scale.data)
 
     @mark_trace
     def forward(
@@ -822,23 +827,23 @@ def forward(
                         scale_b=self.weight_scale,
                     )
             elif self.quant_type.value == QuantType.per_Token.value:
-                if self.params_dtype == dtypes.i8:
-                    y = gemm_a8w8(
+                if use_triton_gemm() and gemm_a8w8_triton is not None:
+                    # Triton a8w8 per-token-per-channel GEMM (unshuffled weight).
+                    y = gemm_a8w8_per_token_impl(
                         x,
                         self.weight,
                         x_scale,
                         self.weight_scale,
-                        self.bias,
+                        bias=self.bias,
                         dtype=otype,
                     )
-                elif use_triton_gemm() and gemm_a8w8_triton is not None:
-                    # Triton a8w8 per-token-per-channel GEMM (unshuffled weight).
-                    y = gemm_a8w8_per_token_impl(
+                elif self.params_dtype == dtypes.i8:
+                    y = gemm_a8w8(
                         x,
                         self.weight,
                         x_scale,
                         self.weight_scale,
-                        bias=self.bias,
+                        self.bias,
                         dtype=otype,
                     )
                 else:

atom/models/deepseek_v2.py

@@ -733,12 +733,6 @@ def _fuse_qkv_a_proj_reduce_rmsnorm_quant_fp8(
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     M = hidden_states_quant.shape[0]
 
-    # NOTE: this fused path always calls aiter's *preshuffle* blockscale GEMMs,
-    # which require a 16x16-shuffled weight. fused_qkv_a_proj is flagged with
-    # needs_preshuffled_weight=True so the loader shuffles it once even under the
-    # non-preshuffle path (ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0) -- see
-    # LinearBase.process_weights_after_loading.
-
     if hidden_states_quant_scale is None:
         if M <= 32:
             qkv_lora = gemm_a16w8_blockscale_preshuffle(
@@ -1735,8 +1729,28 @@ def __init__(
                     quant_config = None
                     base_quant_config = None
             else:
-                source_quant_dtype = torch.bfloat16
-                base_quant_config = None
+                # Mirror the non-triton path above: quark pre-quantized MXFP4
+                # checkpoints (weight_format=real_quantized) store the attention
+                # projections as packed FP4 + e8m0 on disk, so load them directly
+                # (source_quant_dtype=None, keep quant_config via base_quant_config).
+                # The bf16 branch is the online-quant path (load a BF16 checkpoint,
+                # quantize to FP4 at load); using it for a static MXFP4 checkpoint
+                # allocates a bf16 weight param the packed FP4 tensor cannot load
+                # into -> uninitialized garbage -> NaN logits.
+                q_a_proj_quant_config = quant_config.get_layer_quant_config(
+                    f"{prefix}.{q_a_proj_name}"
+                )
+                is_quark_static_mxfp4 = (
+                    q_a_proj_quant_config.quant_method == "quark"
+                    and layer_quant_type == QuantType.per_1x32
+                )
+                if is_quark_static_mxfp4:
+                    source_quant_dtype = None
+                    base_quant_config = quant_config
+                else:
+                    source_quant_dtype = torch.bfloat16
+                    base_quant_config = None
+
         else:
             source_quant_dtype = None
             # Check exclude patterns (e.g. W4A8 checkpoints exclude attention)
@@ -1760,11 +1774,7 @@ def __init__(
                 source_quant_dtype=source_quant_dtype,
                 prefix=f"{prefix}.fused_qkv_a_proj",
             )
-            # The fused qkv_a_proj forward calls *preshuffle* blockscale GEMMs, so
-            # its weight must be 16x16-shuffled even when the global non-preshuffle
-            # path (ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0) is selected. The loader
-            # honors this flag in LinearBase.process_weights_after_loading.
-            self.fused_qkv_a_proj.needs_preshuffled_weight = True
+
             self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
             self.q_b_proj = ColumnParallelLinear(
                 q_lora_rank,