From 1bcb56a8b001f860789c9f425a0d6655433b6022 Mon Sep 17 00:00:00 2001 From: Amelia Moore Date: Mon, 6 Jul 2026 11:37:52 +0000 Subject: [PATCH 1/8] fixes to update aiter main + add shuffle --- atom/model_ops/fused_moe_triton.py | 9 ++++- atom/model_ops/linear.py | 65 +++++++++++++++++++++++++----- atom/models/deepseek_v2.py | 36 +++++++++++------ 3 files changed, 85 insertions(+), 25 deletions(-) diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index 507ba70b97..b01101d40e 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,7 +36,13 @@ from aiter.ops.triton.moe.moe_op_gemm_a4w4 import ( moe_gemm_a4w4, mxfp4_quant, - swizzle_scales as swizzle_scales_cdna4, + ) + # aiter #3900 unified the per-module swizzle_scales into utils/shuffle.py's + # arch-aware shuffle_scale_moe (a8w4/a8w8/a16w4/a4w4 family). Alias to keep + # the a8w4 / cdna4 call sites below unchanged. + from aiter.ops.triton.utils.shuffle import ( + shuffle_scale_moe as swizzle_scales_a8w4, + shuffle_scale_moe as swizzle_scales_cdna4, ) from aiter.ops.triton.moe.quant_moe import downcast_to_static_fp8 diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index d2539d957a..a217edc6de 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -22,6 +22,8 @@ from aiter.jit.utils.torch_guard import torch_compile_guard from aiter.tuned_gemm import tgemm from aiter.utility import fp4_utils +import aiter.ops.triton.utils._triton.arch_info as arch_info +from aiter.ops.triton.utils.shuffle import shuffle_scale_gemm, 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 ( @@ -219,7 +221,7 @@ def gemm_a4w4_quant( 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 + weight_scale.shape[0] // MXFP4_QUANT_BLOCK_SIZE, -1 # NOTE: this may only work for gfx950 ), y=y, ) @@ -750,11 +752,13 @@ def process_weights_after_loading(self): # weight; only the AITER bpreshuffle fallback needs the shuffle. and not (use_triton_gemm() and gemm_a8w8_triton is not None) ) or ( + # gemma4w4 weight shuffled here (-> shuffle_weights below) self.quant_type == QuantType.per_1x32 and (not is_fp4_blockscale or not use_fp4_non_shuffle_triton_gemm()) ) # per_1x128 only needs shuffle when using the preshuffle GEMM path if not need_shuffle and self.quant_type == QuantType.per_1x128: + # gemma8w8 blockscale weight shuffled here (-> shuffle_weights below) need_shuffle = envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE # Modules whose fused forward calls a *preshuffle* blockscale GEMM # directly (e.g. DeepSeek fused qkv_a_proj) need the 16x16-shuffled @@ -767,10 +771,51 @@ def process_weights_after_loading(self): need_shuffle = True if need_shuffle: if self.weight.dim() == 2: - shuffle_weights(self.weight) + if ( + arch_info.get_arch() == "gfx1250" + and use_triton_gemm() + and self.params_dtype == dtypes.fp4x2 + and not use_fp4_non_shuffle_triton_gemm() + ): + # gemma4w4 on gfx1250: WMMA preshuffle layout + from aiter.ops.triton.utils.shuffle import ( + _shuffle_weight_gfx1250, + ) + + self.weight.data = _shuffle_weight_gfx1250(self.weight.data) + elif ( + arch_info.get_arch() == "gfx1250" + and use_triton_gemm() + and self.quant_type == QuantType.per_1x128 + ): + # gemma8w8 blockscale on gfx1250: WMMA preshuffle layout + from aiter.ops.shuffle import preshuffle_fp8_weights_gfx1250 + + self.weight.data = preshuffle_fp8_weights_gfx1250( + self.weight.data + ) + 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 - if self.quant_type == QuantType.per_1x32 and ( + _fp4_triton_preshuffle = ( + self.quant_type == QuantType.per_1x32 + and self.params_dtype == dtypes.fp4x2 + and use_triton_gemm() + and not use_fp4_non_shuffle_triton_gemm() + and gemm_afp4wfp4_preshuffle is not None + ) + if _fp4_triton_preshuffle: + # Shuffle the weight scale into the gemm_afp4wfp4_preshuffle kernel layout + # at load. shuffle_scale_gemm_e8m0 is built on shuffle_scale_gemm's gfx950 + # tile (preshuffle_factor=32, scale_kwidth=8) but returns the UN-collapsed + # (M_pad, N_pad) shape -- byte-identical to the original fp4_utils.e8m0_shuffle + # that the forward (gemm_a4w4_quant, ~line 223) still re-collapses via + # `.view(shape[0] // 32, -1)`. (Raw shuffle_scale_gemm returns the already + # collapsed (M_pad//32, N*32), which the forward would divide by 32 again + # -> wrong-strided scale -> OOB read in the preshuffle kernel -> GPU fault.) + self.weight_scale.data = shuffle_scale_gemm_e8m0(self.weight_scale.data) + elif 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) @@ -822,23 +867,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..e0ff0217e6 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, From 453e9abf4209120c0c98afeb46cce377ebc4ecca Mon Sep 17 00:00:00 2001 From: Amelia Moore Date: Tue, 7 Jul 2026 12:44:09 +0000 Subject: [PATCH 2/8] scale shuffle adjustments for arch --- atom/model_ops/linear.py | 91 +++++++++++++++++++--------------------- 1 file changed, 44 insertions(+), 47 deletions(-) diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index a217edc6de..8fe6bd3526 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -198,30 +198,51 @@ def gemm_a4w4_quant( dtype=otype, device=x.device, ) + # Scale preshuffle factor per arch: 32 on gfx950, 16 on gfx1250 (WMMA + # tiles both M and N in 16-lane groups). Applies to BOTH the activation + # (M-axis) and weight (N-axis) e8m0 scales. + _arch = arch_info.get_arch() + _scale_pf = 16 if _arch == "gfx1250" else MXFP4_QUANT_BLOCK_SIZE + 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), - ) + if _arch == "gfx1250" and m >= MXFP4_QUANT_BLOCK_SIZE: + # get_hip_quant's built-in scale shuffle is gfx950-pinned + # (per_1x32_mx_quant_hip pads to (256, 8) -> the 32/8 tile, + # no arch branch). On gfx1250 quantize WITHOUT that shuffle and + # apply the arch-aware e8m0 tile (16/4) in-line, mirroring the + # weight-scale path so the activation scale matches the gfx1250 + # GEMM instead of arriving in the gfx950 layout. + x, x_scale = quant_func( + x, quant_dtype=params_dtype, shuffle=False + ) + x_scale = shuffle_scale_gemm_e8m0(x_scale.view(torch.uint8), arch=_arch) + else: + x, x_scale = quant_func( + x, + quant_dtype=params_dtype, + shuffle=(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 + x_scale.shape[0] // _scale_pf, -1 ) else: x_scale = x_scale[:m, ...].view(torch.uint8) + # Re-collapse the un-collapsed (rows_pad, kgroups_pad) weight scale + # produced by shuffle_scale_gemm_e8m0 into the kernel layout, by the same + # arch preshuffle factor used for the activation scale above. 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 # NOTE: this may only work for gfx950 + weight_scale.shape[0] // _scale_pf, -1 ), y=y, ) @@ -771,54 +792,30 @@ def process_weights_after_loading(self): need_shuffle = True if need_shuffle: if self.weight.dim() == 2: - if ( - arch_info.get_arch() == "gfx1250" - and use_triton_gemm() - and self.params_dtype == dtypes.fp4x2 - and not use_fp4_non_shuffle_triton_gemm() - ): - # gemma4w4 on gfx1250: WMMA preshuffle layout + if use_triton_gemm(): + # gfx1250 WMMA preshuffle layout for both gemma4w4 (FP4) and + # gemma8w8 blockscale, via the arch-aware triton shuffle_weight + # (dispatches to _shuffle_weight_gfx1250 on gfx1250). from aiter.ops.triton.utils.shuffle import ( - _shuffle_weight_gfx1250, + shuffle_weight as _triton_shuffle_weight, ) - self.weight.data = _shuffle_weight_gfx1250(self.weight.data) - elif ( - arch_info.get_arch() == "gfx1250" - and use_triton_gemm() - and self.quant_type == QuantType.per_1x128 - ): - # gemma8w8 blockscale on gfx1250: WMMA preshuffle layout - from aiter.ops.shuffle import preshuffle_fp8_weights_gfx1250 - - self.weight.data = preshuffle_fp8_weights_gfx1250( - self.weight.data - ) + 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 - _fp4_triton_preshuffle = ( - self.quant_type == QuantType.per_1x32 - and self.params_dtype == dtypes.fp4x2 - and use_triton_gemm() - and not use_fp4_non_shuffle_triton_gemm() - and gemm_afp4wfp4_preshuffle is not None - ) - if _fp4_triton_preshuffle: - # Shuffle the weight scale into the gemm_afp4wfp4_preshuffle kernel layout - # at load. shuffle_scale_gemm_e8m0 is built on shuffle_scale_gemm's gfx950 - # tile (preshuffle_factor=32, scale_kwidth=8) but returns the UN-collapsed - # (M_pad, N_pad) shape -- byte-identical to the original fp4_utils.e8m0_shuffle - # that the forward (gemm_a4w4_quant, ~line 223) still re-collapses via - # `.view(shape[0] // 32, -1)`. (Raw shuffle_scale_gemm returns the already - # collapsed (M_pad//32, N*32), which the forward would divide by 32 again - # -> wrong-strided scale -> OOB read in the preshuffle kernel -> GPU fault.) - self.weight_scale.data = shuffle_scale_gemm_e8m0(self.weight_scale.data) - elif self.quant_type == QuantType.per_1x32 and ( + # Shuffle the weight scale once at load. One call covers every per_1x32 case + # that needs a shuffled scale: the FP4 triton-preshuffle path (config C) and + # the non-triton / non-fp4x2 fallbacks. The only per_1x32 case skipped is FP4 + + # non-shuffle triton (config D), which passes raw scales. shuffle_scale_gemm_e8m0 + # is arch-aware (gfx950 32/8 tile == fp4_utils.e8m0_shuffle byte-for-byte, proven; + # gfx1250 16/4 tile) and returns the UN-collapsed (M_pad, N_pad) view that the + # forward (gemm_a4w4_quant) re-collapses by the arch preshuffle_factor. + 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( From 5e44c1071f007fca7cb0c0e9683f208f8b5c1269 Mon Sep 17 00:00:00 2001 From: Amelia Moore Date: Tue, 7 Jul 2026 15:36:11 +0000 Subject: [PATCH 3/8] remove arch dependency --- atom/model_ops/linear.py | 60 ++++++++++++---------------------------- 1 file changed, 17 insertions(+), 43 deletions(-) diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index 8fe6bd3526..fadf0031be 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -21,9 +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 -import aiter.ops.triton.utils._triton.arch_info as arch_info -from aiter.ops.triton.utils.shuffle import shuffle_scale_gemm, shuffle_scale_gemm_e8m0 +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 ( @@ -198,52 +200,26 @@ def gemm_a4w4_quant( dtype=otype, device=x.device, ) - # Scale preshuffle factor per arch: 32 on gfx950, 16 on gfx1250 (WMMA - # tiles both M and N in 16-lane groups). Applies to BOTH the activation - # (M-axis) and weight (N-axis) e8m0 scales. - _arch = arch_info.get_arch() - _scale_pf = 16 if _arch == "gfx1250" else MXFP4_QUANT_BLOCK_SIZE - + # Arch-specific MXFP4 preshuffle operand prep (scale preshuffle factor, + # gfx1250-vs-gfx950 activation quant + e8m0 tile, scale re-collapse) lives + # in aiter.ops.triton.utils.shuffle so load-time shuffle and forward GEMM + # stay in one place. if x_scale is None: - quant_func = get_hip_quant(QuantType.per_1x32) - if _arch == "gfx1250" and m >= MXFP4_QUANT_BLOCK_SIZE: - # get_hip_quant's built-in scale shuffle is gfx950-pinned - # (per_1x32_mx_quant_hip pads to (256, 8) -> the 32/8 tile, - # no arch branch). On gfx1250 quantize WITHOUT that shuffle and - # apply the arch-aware e8m0 tile (16/4) in-line, mirroring the - # weight-scale path so the activation scale matches the gfx1250 - # GEMM instead of arriving in the gfx950 layout. - x, x_scale = quant_func( - x, quant_dtype=params_dtype, shuffle=False - ) - x_scale = shuffle_scale_gemm_e8m0(x_scale.view(torch.uint8), arch=_arch) - else: - x, x_scale = quant_func( - x, - quant_dtype=params_dtype, - shuffle=(m >= MXFP4_QUANT_BLOCK_SIZE), - ) + x, x_scale = quant_mxfp4_act_preshuffle(x, params_dtype, m) 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] // _scale_pf, -1 + x_scale = collapse_mxfp4_gemm_scale( + x_scale.view(torch.float8_e8m0fnu), rows_valid=m ) - else: - x_scale = x_scale[:m, ...].view(torch.uint8) + x = x.view(torch.float4_e2m1fn_x2) - # Re-collapse the un-collapsed (rows_pad, kgroups_pad) weight scale - # produced by shuffle_scale_gemm_e8m0 into the kernel layout, by the same - # arch preshuffle factor used for the activation scale above. + # collapse_mxfp4_gemm_scale re-collapses the un-collapsed (rows_pad, + # kgroups_pad) weight scale produced by shuffle_scale_gemm_e8m0 into the + # kernel layout, by the same arch preshuffle factor as the activation scale. 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] // _scale_pf, -1 - ), + collapse_mxfp4_gemm_scale(weight_scale), y=y, ) # Default AITER path: quantize/shuffle into the layout expected by gemm_a4w4 @@ -773,13 +749,11 @@ def process_weights_after_loading(self): # weight; only the AITER bpreshuffle fallback needs the shuffle. and not (use_triton_gemm() and gemm_a8w8_triton is not None) ) or ( - # gemma4w4 weight shuffled here (-> shuffle_weights below) self.quant_type == QuantType.per_1x32 and (not is_fp4_blockscale or not use_fp4_non_shuffle_triton_gemm()) ) # per_1x128 only needs shuffle when using the preshuffle GEMM path if not need_shuffle and self.quant_type == QuantType.per_1x128: - # gemma8w8 blockscale weight shuffled here (-> shuffle_weights below) need_shuffle = envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE # Modules whose fused forward calls a *preshuffle* blockscale GEMM # directly (e.g. DeepSeek fused qkv_a_proj) need the 16x16-shuffled From 919262d2f839d2147a2c3e3cabe5ca4e68f97d20 Mon Sep 17 00:00:00 2001 From: Amelia Moore Date: Tue, 7 Jul 2026 15:43:15 +0000 Subject: [PATCH 4/8] cleanup --- atom/model_ops/fused_moe_triton.py | 3 --- atom/model_ops/linear.py | 19 ++----------------- 2 files changed, 2 insertions(+), 20 deletions(-) diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index b01101d40e..e7deac7c2b 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -37,9 +37,6 @@ moe_gemm_a4w4, mxfp4_quant, ) - # aiter #3900 unified the per-module swizzle_scales into utils/shuffle.py's - # arch-aware shuffle_scale_moe (a8w4/a8w8/a16w4/a4w4 family). Alias to keep - # the a8w4 / cdna4 call sites below unchanged. from aiter.ops.triton.utils.shuffle import ( shuffle_scale_moe as swizzle_scales_a8w4, shuffle_scale_moe as swizzle_scales_cdna4, diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index fadf0031be..747da415d8 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -200,10 +200,6 @@ def gemm_a4w4_quant( dtype=otype, device=x.device, ) - # Arch-specific MXFP4 preshuffle operand prep (scale preshuffle factor, - # gfx1250-vs-gfx950 activation quant + e8m0 tile, scale re-collapse) lives - # in aiter.ops.triton.utils.shuffle so load-time shuffle and forward GEMM - # stay in one place. if x_scale is None: x, x_scale = quant_mxfp4_act_preshuffle(x, params_dtype, m) else: @@ -212,9 +208,6 @@ def gemm_a4w4_quant( ) x = x.view(torch.float4_e2m1fn_x2) - # collapse_mxfp4_gemm_scale re-collapses the un-collapsed (rows_pad, - # kgroups_pad) weight scale produced by shuffle_scale_gemm_e8m0 into the - # kernel layout, by the same arch preshuffle factor as the activation scale. y = gemm_afp4wfp4_preshuffle( x.view(torch.uint8), weight.view(torch.uint8).view(weight.shape[0] // 16, -1), @@ -767,9 +760,7 @@ def process_weights_after_loading(self): if need_shuffle: if self.weight.dim() == 2: if use_triton_gemm(): - # gfx1250 WMMA preshuffle layout for both gemma4w4 (FP4) and - # gemma8w8 blockscale, via the arch-aware triton shuffle_weight - # (dispatches to _shuffle_weight_gfx1250 on gfx1250). + # arch-aware weight shuffle from aiter.ops.triton.utils.shuffle import ( shuffle_weight as _triton_shuffle_weight, ) @@ -779,13 +770,7 @@ def process_weights_after_loading(self): else: shuffle_weights(self.weight) # self.weight_scale.data = fp4_utils.e8m0_shuffle(self.weight_scale.data) - # Shuffle the weight scale once at load. One call covers every per_1x32 case - # that needs a shuffled scale: the FP4 triton-preshuffle path (config C) and - # the non-triton / non-fp4x2 fallbacks. The only per_1x32 case skipped is FP4 + - # non-shuffle triton (config D), which passes raw scales. shuffle_scale_gemm_e8m0 - # is arch-aware (gfx950 32/8 tile == fp4_utils.e8m0_shuffle byte-for-byte, proven; - # gfx1250 16/4 tile) and returns the UN-collapsed (M_pad, N_pad) view that the - # forward (gemm_a4w4_quant) re-collapses by the arch preshuffle_factor. + # shuffle the weight once if self.quant_type == QuantType.per_1x32 and ( self.params_dtype != dtypes.fp4x2 or not use_fp4_non_shuffle_triton_gemm() ): From 5eb675bf34ccb6413642c1145bb5c1cf6ff47971 Mon Sep 17 00:00:00 2001 From: Amelia Moore Date: Tue, 7 Jul 2026 15:54:42 +0000 Subject: [PATCH 5/8] black --- atom/models/deepseek_v2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index e0ff0217e6..0b416f4092 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -1774,7 +1774,7 @@ def __init__( source_quant_dtype=source_quant_dtype, prefix=f"{prefix}.fused_qkv_a_proj", ) - + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) self.q_b_proj = ColumnParallelLinear( q_lora_rank, From a7d0edc4aa23a972cc9bdd2f65a9eeaa64ea27ab Mon Sep 17 00:00:00 2001 From: Amelia Moore Date: Tue, 7 Jul 2026 18:24:15 +0000 Subject: [PATCH 6/8] decouple MXFP4_QUANT_BLOCK_SIZE + ruff + black --- atom/model_ops/linear.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index 747da415d8..3b0193b5a9 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -201,10 +201,14 @@ def gemm_a4w4_quant( device=x.device, ) if x_scale is None: - x, x_scale = quant_mxfp4_act_preshuffle(x, params_dtype, m) + # quant with no shuffle, then shuffle + x, x_scale = quant_mxfp4_act_preshuffle( + x, params_dtype, m, MXFP4_QUANT_BLOCK_SIZE + ) else: + # collapse in arch-aware layout x_scale = collapse_mxfp4_gemm_scale( - x_scale.view(torch.float8_e8m0fnu), rows_valid=m + x_scale.view(torch.float8_e8m0fnu), MXFP4_QUANT_BLOCK_SIZE, rows_valid=m ) x = x.view(torch.float4_e2m1fn_x2) @@ -212,7 +216,7 @@ def gemm_a4w4_quant( x.view(torch.uint8), weight.view(torch.uint8).view(weight.shape[0] // 16, -1), x_scale, - collapse_mxfp4_gemm_scale(weight_scale), + collapse_mxfp4_gemm_scale(weight_scale, MXFP4_QUANT_BLOCK_SIZE), y=y, ) # Default AITER path: quantize/shuffle into the layout expected by gemm_a4w4 From dbf71ec3465d2268595288bdc5be886db08d8692 Mon Sep 17 00:00:00 2001 From: Amelia Moore Date: Wed, 8 Jul 2026 20:28:12 +0000 Subject: [PATCH 7/8] import fix --- atom/model_ops/fused_moe_triton.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index e7deac7c2b..834869e952 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -37,10 +37,7 @@ moe_gemm_a4w4, mxfp4_quant, ) - from aiter.ops.triton.utils.shuffle import ( - shuffle_scale_moe as swizzle_scales_a8w4, - shuffle_scale_moe 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 @@ -53,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 From 0131a4696f5699c9ad75632ffb448a53eaa78967 Mon Sep 17 00:00:00 2001 From: Amelia Moore Date: Wed, 8 Jul 2026 20:48:57 +0000 Subject: [PATCH 8/8] comments --- atom/model_ops/linear.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index 3b0193b5a9..e2430a6af8 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -201,7 +201,7 @@ def gemm_a4w4_quant( device=x.device, ) if x_scale is None: - # quant with no shuffle, then shuffle + # quant with no shuffle + manual shuffle + collapse x, x_scale = quant_mxfp4_act_preshuffle( x, params_dtype, m, MXFP4_QUANT_BLOCK_SIZE ) @@ -774,7 +774,7 @@ def process_weights_after_loading(self): else: shuffle_weights(self.weight) # self.weight_scale.data = fp4_utils.e8m0_shuffle(self.weight_scale.data) - # shuffle the weight once + # 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() ):