[AMD] Perf Kimi-K3 MoE optimization (#33838)

Co-authored-by: wunhuang <wunhuang@amd.com>

python/sglang/srt/layers/moe/topk.py

@@ -16,7 +16,7 @@ from __future__ import annotations
 
 import logging
 import math
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import IntEnum, auto
 from typing import (
     TYPE_CHECKING,
@@ -238,6 +238,32 @@ class TopKConfig:
     # Draft-side MoE blocks set this False so they never write the target's
     # process-global routed-experts capture buffer.
     allow_routed_experts_capture: bool = True
+    _correction_bias_dtype_cache: Optional[torch.Tensor] = field(
+        default=None, init=False, repr=False, compare=False
+    )
+    _correction_bias_cache_key: Optional[Tuple] = field(
+        default=None, init=False, repr=False, compare=False
+    )
+
+    def correction_bias_for_dtype(self, dtype: torch.dtype) -> Optional[torch.Tensor]:
+        """Return correction bias in ``dtype``, reusing a per-TopK lazy copy."""
+        correction_bias = self.correction_bias
+        if correction_bias is None or correction_bias.dtype == dtype:
+            return correction_bias
+
+        # Weight loaders update parameters in place. Including the version in
+        # the key prevents an early access from retaining pre-load contents.
+        cache_key = (
+            correction_bias.data_ptr(),
+            correction_bias._version,
+            correction_bias.device,
+            correction_bias.dtype,
+            dtype,
+        )
+        if self._correction_bias_cache_key != cache_key:
+            self._correction_bias_dtype_cache = correction_bias.to(dtype=dtype)
+            self._correction_bias_cache_key = cache_key
+        return self._correction_bias_dtype_cache
 
 
 # -------------------------------- TopKOutput ---------------------------------------
@@ -2306,6 +2332,9 @@ def select_experts(
         info=expert_location_dispatch_info,
     )
 
+    if _use_aiter and use_grouped_topk and correction_bias is not None:
+        correction_bias = topk_config.correction_bias_for_dtype(router_logits.dtype)
+
     # DeepSeek V2/V3/R1 series models use grouped_top_k
     # remove num_fused_shared_experts from grouped_topk/biased_grouped_topk
     num_routed_topk = top_k - num_fused_shared_experts

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

@@ -156,6 +156,13 @@ _flashinfer_mxfp4_permute_indices_device_cache: dict[
 ] = {}
 
 
+def _aiter_situ_uses_gu_interleaved_weights() -> bool:
+    """Match AITER's SiTU activation-mode precedence when choosing weight layout."""
+    a8w4 = get_bool_env_var("AITER_SITUV2_A8W4", "false")
+    a4w4 = get_bool_env_var("AITER_SITUV2_A4W4", "false")
+    return a8w4 or not a4w4
+
+
 def _get_flashinfer_mxfp4_device_permute_indices(
     x: torch.Tensor,
     epilogue_tile_m: int,
@@ -951,12 +958,17 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                         .view(-1, n)
                     )
 
-            k3_situ_a8w4 = (
-                os.environ.get("AITER_SITUV2_A8W4", "0") == "1"
-                and getattr(layer.moe_runner_config, "activation", None) == "situ"
-            )
-            use_aiter_gu_interleave = k3_situ_a8w4 or (
-                envs.SGLANG_USE_AITER_MOE_GU_ITLV.get() and gate_up_interleaved
+            # AITER selects the activation dtype at runtime. A8W4 takes precedence
+            # and, together with A16W4, uses the preshuffled GU-interleaved layout.
+            # A4W4 uses the generic separated layout instead; feeding it the
+            # A16/A8 layout makes real-checkpoint MoE outputs nearly orthogonal.
+            k3_situ = getattr(layer.moe_runner_config, "activation", None) == "situ"
+            use_aiter_gu_interleave = (
+                k3_situ and _aiter_situ_uses_gu_interleaved_weights()
+            ) or (
+                not k3_situ
+                and envs.SGLANG_USE_AITER_MOE_GU_ITLV.get()
+                and gate_up_interleaved
             )
             if use_aiter_gu_interleave:
                 layer.w13_weight.data = shuffle_weight_a16w4(layer.w13_weight, 16, True)
@@ -1733,9 +1745,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                             expanded_idx_to_permuted_idx=expanded_idx,
                             top_k=packed_topk.shape[1],
                         )
-                    else:
-                        result = result[0]
-                    return StandardCombineInput(hidden_states=result)
+                        return StandardCombineInput(hidden_states=result)
+                    # The finalized kernel writes to its explicit output
+                    # argument. Do not propagate the FFI return tensor: some
+                    # SiTU runner versions return a distinct wrapper/allocation
+                    # even though symm_output contains the published result.
+                    # Returning the destination makes the pointer contract
+                    # explicit for K3's zero-copy latent buffer.
+                    return StandardCombineInput(hidden_states=symm_output)
 
                 # Bypassed topk: route from logits inside the op.
                 correction_bias = topk_output.topk_config.correction_bias