diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index 507ba70b97..834869e952 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -29,7 +29,6 @@ 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, @@ -37,8 +36,8 @@ 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 diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index d2539d957a..e2430a6af8 100644 --- a/atom/model_ops/linear.py +++ b/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: diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index bd72e39cb2..0b416f4092 100644 --- a/atom/models/deepseek_v2.py +++ b/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,