enable ar+norm+quant fusion (#1494)

* enable ar+norm+quant fusion

* fix fusion with online quant

Author: gbyu-amd

atom/model_ops/layernorm.py

@@ -219,6 +219,7 @@ def __init__(
         x_pad_to_multiple: int = 0,
         fused_allreduce: bool = False,
         fused_quant: bool = False,
+        fused_quant_emit_bf16: bool = False,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
     ) -> None:
@@ -229,6 +230,12 @@ def __init__(
         self.x_pad_to_multiple = x_pad_to_multiple
         self.fused_allreduce = fused_allreduce
         self.use_fused_quant = fused_quant
+        # When the fused AllReduce+RMSNorm+quant path is active AND a downstream
+        # consumer also needs the unquantized (bf16) normed activation (e.g. the
+        # GLM-5.2 DSA indexer, whose wk/weights_proj GEMMs run in BF16), the
+        # kernel additionally emits that bf16 mirror. forward() then returns
+        # ((fp8, scale, bf16), residual) instead of ((fp8, scale), residual).
+        self.fused_quant_emit_bf16 = fused_quant_emit_bf16
         self.tp_size = get_tensor_model_parallel_world_size()
         self.quant_config = quant_config
         self.prefix = prefix
@@ -335,6 +342,44 @@ def forward(
             assert (
                 residual is not None
             ), "fused_allreduce_rmsnorm requires residual input!"
+            if self.use_fused_quant and self.quant_type.value in (
+                _QV_PER_1X128,
+                _QV_PER_TOKEN,
+            ):
+                # Combined AllReduce + RMSNorm + FP8 quant: the downstream GEMM
+                # (e.g. fused_qkv_a_proj) consumes the (fp8, scale) output
+                # directly, dropping a standalone per-token/per-group quant kernel
+                # from the hot path. `_aiter_transpose_scale` (resolved at init)
+                # selects the scale layout for that GEMM. The fused kernel does
+                # not support non-contiguous input.
+                if self.fused_quant_emit_bf16:
+                    # Also emit the pre-quant bf16 normed activation for a
+                    # co-consumer that runs in BF16 (GLM-5.2 DSA indexer).
+                    x, residual, x_scale, x_bf16 = (
+                        tensor_model_parallel_fused_allreduce_rmsnorm_quant(
+                            x.contiguous(),
+                            residual,
+                            self.weight,
+                            self.eps,
+                            quant_type=self.quant_type,
+                            group_size=128,
+                            transpose_scale=self._aiter_transpose_scale,
+                            emit_bf16=True,
+                        )
+                    )
+                    return (x, x_scale, x_bf16), residual
+                x, residual, x_scale = (
+                    tensor_model_parallel_fused_allreduce_rmsnorm_quant(
+                        x.contiguous(),
+                        residual,
+                        self.weight,
+                        self.eps,
+                        quant_type=self.quant_type,
+                        group_size=128,
+                        transpose_scale=self._aiter_transpose_scale,
+                    )
+                )
+                return (x, x_scale), residual
             # tensor_model_parallel_fused_allreduce_rmsnorm does not support non-contiguous input
             x, residual = tensor_model_parallel_fused_allreduce_rmsnorm(
                 x.contiguous(),

atom/models/deepseek_v2.py

@@ -81,6 +81,7 @@
     make_layers,
     maybe_prefix,
 )
+from atom.quant_spec import should_skip_online_quant
 from atom.utils import envs
 from atom.utils.custom_register import direct_register_custom_op
 
@@ -1949,13 +1950,23 @@ def __init__(
         # qkv_a_proj + reduce + RMSNorm + quant fusion remains gated by
         # use_triton_gemm() in forward(), because that path depends on Triton GEMM.
         self.prefix = prefix
-        self.quant_dtype = layer_quant_dtype
-        self.qknorm_quant_type = layer_quant_type_value
+        # Online-aware scheme for the fused q/kv norm+quant feeding q_b_proj:
+        # use the online target (applied after __init__), else the static config.
+        # layer_quant_dtype/type stay static for the fp4x2 weight-loading branch.
+        eff_dtype, eff_type = layer_quant_dtype, layer_quant_type
+        if quant_config is not None and quant_config.online_quant:
+            online_cfg = quant_config.get_layer_quant_config(
+                f"{prefix}.{q_a_proj_name}", use_online_quant=True
+            )
+            if not should_skip_online_quant(eff_type, eff_dtype, online_cfg):
+                eff_dtype, eff_type = online_cfg.quant_dtype, online_cfg.quant_type
+        self.quant_dtype = eff_dtype
+        self.qknorm_quant_type = None if eff_type is None else eff_type.value
         self.fuse_qknorm_quant = False
         # always fuse qknorm
         self.fuse_qknorm = ENABLE_DS_QKNORM_FUSION
         if quant_config is not None and ENABLE_DS_QKNORM_QUANT_FUSION:
-            if layer_quant_dtype in (dtypes.fp8, dtypes.fp4x2):
+            if eff_dtype in (dtypes.fp8, dtypes.fp4x2):
                 self.fuse_qknorm_quant = True
 
     def forward(
@@ -1964,8 +1975,16 @@ def forward(
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
         hidden_states_scale = None
+        # When input_layernorm fused AR+RMSNorm+quant, hidden_states is a tuple.
+        # A 3-tuple (fp8, scale, bf16) additionally carries the unquantized bf16
+        # normed activation for the v32 indexer (see RMSNorm.fused_quant_emit_bf16);
+        # a 2-tuple (fp8, scale) is the plain fused-quant output.
+        indexer_hidden = None
         if isinstance(hidden_states, tuple):
-            hidden_states, hidden_states_scale = hidden_states
+            if len(hidden_states) == 3:
+                hidden_states, hidden_states_scale, indexer_hidden = hidden_states
+            else:
+                hidden_states, hidden_states_scale = hidden_states
 
         if self.q_lora_rank is not None:
             if self.fuse_qknorm_quant and use_triton_gemm():
@@ -2042,8 +2061,11 @@ def forward(
             kv_c_normed = self.kv_a_layernorm(kv_c)
             hidden_states_or_q_c_scale = None
         if self.is_v32 and self.indexer is not None and not self.skip_topk:
+            # The indexer's wk/weights_proj GEMMs run in BF16. When input_layernorm
+            # fused the quant it emits a bf16 mirror (indexer_hidden); otherwise
+            # hidden_states is already the bf16 normed activation.
             self.indexer(
-                hidden_states,
+                indexer_hidden if indexer_hidden is not None else hidden_states,
                 hidden_states_or_q_c,
                 hidden_states_or_q_c_scale,
                 positions,
@@ -2158,12 +2180,57 @@ def __init__(
                 reduce_results=not self.fuse_ar_input_norm,
                 prefix=f"{prefix}.mlp",
             )
+        # Fuse activation quant into AR+RMSNorm when fused_qkv_a_proj is
+        # per-1x128/per-token FP8, so the GEMM consumes the (fp8, scale) directly.
+        # Use the online-quant target (get_layer_quant_config(..., online)), not
+        # the static params_dtype: online quant flips the qkv proj to fp8 only in
+        # process_weights_after_loading, which runs after this __init__.
+        qkv_proj = getattr(self.self_attn, "fused_qkv_a_proj", None)
+        eff_quant_type = None if qkv_proj is None else qkv_proj.quant_type
+        eff_quant_dtype = None if qkv_proj is None else qkv_proj.params_dtype
+        if (
+            qkv_proj is not None
+            and quant_config is not None
+            and quant_config.online_quant
+        ):
+            online_cfg = quant_config.get_layer_quant_config(
+                qkv_proj.prefix, use_online_quant=True
+            )
+            if not should_skip_online_quant(
+                eff_quant_type, eff_quant_dtype, online_cfg
+            ):
+                eff_quant_type = online_cfg.quant_type
+                eff_quant_dtype = online_cfg.quant_dtype
+        input_norm_fused_quant = (
+            qkv_proj is not None
+            and eff_quant_dtype == dtypes.fp8
+            and eff_quant_type.value
+            in (QuantType.per_1x128.value, QuantType.per_Token.value)
+        )
+        fused_allreduce = (
+            self.fuse_ar_input_norm and self.layer_idx > 0 and not is_mtp_block
+        )
+        # GLM-5.2 (v32) DSA: the indexer's wk/weights_proj GEMMs run in BF16 and
+        # consume the same normed activation, so the fused kernel must also emit
+        # the pre-quant bf16 mirror alongside the fp8 for fused_qkv_a_proj.
+        emit_bf16_for_indexer = (
+            getattr(self.self_attn, "is_v32", False)
+            and getattr(self.self_attn, "indexer", None) is not None
+        )
         self.input_layernorm = RMSNorm(
             config.hidden_size,
             eps=config.rms_norm_eps,
-            fused_allreduce=self.fuse_ar_input_norm
-            and self.layer_idx > 0
-            and not is_mtp_block,
+            fused_allreduce=fused_allreduce,
+            fused_quant=fused_allreduce and input_norm_fused_quant,
+            fused_quant_emit_bf16=(
+                fused_allreduce and input_norm_fused_quant and emit_bf16_for_indexer
+            ),
+            quant_config=quant_config,
+            prefix=(
+                f"{prefix}.self_attn.fused_qkv_a_proj"
+                if qkv_proj is not None
+                else f"{prefix}.self_attn.q_a_proj"
+            ),
         )
         self.post_attention_layernorm = RMSNorm(
             config.hidden_size,