[k3] enable dual stream for shared expert and some fusions

.github/benchmark/models_accuracy.json

@@ -249,7 +249,7 @@
   {
     "model_name": "Kimi-K3",
     "model_path": "moonshotai/Kimi-K3",
-    "extraArgs": "--kv_cache_dtype fp8 -tp 8 --trust-remote-code --gpu-memory-utilization 0.93 --block-size 128 --no-enable_prefix_caching",
+    "extraArgs": "--kv_cache_dtype fp8 -tp 8 --trust-remote-code --gpu-memory-utilization 0.93 --block-size 128 --no-enable_prefix_caching --online_quant_config '{\"global_quant_config\": \"ptpc_fp8\", \"exclude_layer\": [\"lm_head\", \"model.embed_tokens\", \"*self_attn.[qkv]_conv1d*\", \"*block_sparse_moe.experts*\", \"*block_sparse_moe.routed_expert_*\", \"*vision_tower*\", \"*mm_projector*\"]}'",
     "env_vars": "",
     "runner": "linux-atom-do-mi350x-8",
     "test_level": "pr",

atom/config.py

@@ -314,8 +314,15 @@ def __init__(
             "mxfp4",
             "mxfp8",
             "quark",
+            "compressed-tensors",
         ]:
             self.online_quant = True
+            if self.quant_method == "compressed-tensors":
+                logger.warning(
+                    "Online quant with compressed-tensors is not fully supported. "
+                    "Be careful about the online quant config setting when launching "
+                    "the server."
+                )
             online_parser = get_quant_parser("online_quant")
             online_parsed_quant_config = online_parser.parse(online_quant_config)
             self.online_global_spec = online_parsed_quant_config.global_spec

atom/model_ops/kimi_k3/__init__.py

@@ -3,7 +3,10 @@
 
 """Fused model operations for Kimi-K3."""
 
-from atom.model_ops.kimi_k3.activations import rmsnorm_gated, situ_and_mul
+from atom.model_ops.kimi_k3.activations import (
+    rmsnorm_gated,
+    situ_and_mul,
+)
 from atom.model_ops.kimi_k3.attention_residual import apply_attn_res
 from atom.model_ops.kimi_k3.kda_state import gather_kda_initial_state
 

atom/model_ops/kimi_k3/activations.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import torch
+from aiter import QuantType, dtypes, get_hip_quant
 
 try:
     import triton
@@ -81,6 +82,53 @@ def _rmsnorm_gated_kernel(
             y_ptr + row * stride_ym + cols, y.to(y_ptr.dtype.element_ty), mask=mask
         )
 
+    @triton.jit
+    def _rmsnorm_gated_fp8_per_token_kernel(
+        x_ptr,
+        w_ptr,
+        g_ptr,
+        y_ptr,
+        s_ptr,
+        H,
+        eps,
+        fp8_max,
+        stride_xm,
+        stride_xh,
+        stride_g_outer,
+        stride_g_head,
+        stride_ym,
+        HEADS: tl.constexpr,
+        HEADS_POW2: tl.constexpr,
+        BLOCK: tl.constexpr,
+    ):
+        tok = tl.program_id(0)
+        head_ids = tl.arange(0, HEADS_POW2)
+        cols = tl.arange(0, BLOCK)
+        mask = (head_ids[:, None] < HEADS) & (cols[None, :] < H)  # [HEADS_POW2, BLOCK]
+        # Padding heads (head_ids >= HEADS) are masked out on every load/store,
+        # but their raw offset (head_ids * stride) can still address past the end
+        # of the buffer -- forming an out-of-bounds pointer is UB on ROCm/triton
+        # and faults when the allocation abuts an unmapped page. Clamp the head
+        # index used for addressing to a valid row; the mask (other=0.0) still
+        # discards the value, so numerics are unchanged.
+        h_safe = tl.where(head_ids < HEADS, head_ids, 0)
+        x_off = tok * stride_xm + h_safe[:, None] * stride_xh + cols[None, :]
+        x = tl.load(x_ptr + x_off, mask=mask, other=0.0).to(tl.float32)
+        var = tl.sum(x * x, axis=1) / H  # [HEADS]
+        rstd = 1.0 / tl.sqrt(var + eps)  # [HEADS]
+        w = tl.load(w_ptr + cols, mask=cols < H, other=0.0).to(tl.float32)  # [BLOCK]
+        g_off = tok * stride_g_outer + h_safe[:, None] * stride_g_head + cols[None, :]
+        gate = tl.load(g_ptr + g_off, mask=mask, other=0.0).to(tl.float32)
+        normed = (x * rstd[:, None] * w[None, :]) * tl.sigmoid(gate)  # [HEADS, BLOCK]
+        amax = tl.max(tl.abs(normed))  # scalar per token
+        scale = amax / fp8_max
+        inv = tl.where(scale > 0.0, 1.0 / scale, 0.0)
+        q = normed * inv
+        q = tl.minimum(tl.maximum(q, -fp8_max), fp8_max)
+        y_off = tok * stride_ym + h_safe[:, None] * H + cols[None, :]
+        tl.store(y_ptr + y_off, q.to(y_ptr.dtype.element_ty), mask=mask)
+        tl.store(s_ptr + tok, scale)
+
 
 def situ_and_mul(
     x: torch.Tensor, beta: float, linear_beta: float | None
@@ -115,15 +163,36 @@ def situ_and_mul(
 
 
 def rmsnorm_gated(
-    x: torch.Tensor, weight: torch.Tensor, gate: torch.Tensor, eps: float
-) -> torch.Tensor:
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    gate: torch.Tensor,
+    eps: float,
+    quant_type: QuantType | None = None,
+    quant_dtype: torch.dtype | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """rmsnorm(x) over last dim * weight * sigmoid(gate).
 
+    When ``(quant_type, quant_dtype)`` is the per-token FP8 scheme, the normed
+    output is also quantized and the function returns
+    ``(fp8 [t, heads*H], scale [t, 1])`` ready for the consuming GEMM's
+    ``x_scale=`` path. With no quant (``None``/``QuantType.No``) it returns the
+    bf16 tensor shaped like ``x``. Per-token FP8 is the only fused scheme today
+    (the consuming o_proj's a8w8 scheme); any other requested scheme asserts.
+
     ``gate`` may be strided (e.g. a column slice of a fused GEMM output): the
     kernel reads it via (outer, head) strides so no contiguous copy is needed.
     ``x`` is normed row-wise and is made contiguous (cheap; the caller's ``out``
     already is). Supports a 2D ``[M, H]`` or 3D ``[outer, heads, H]`` gate.
     """
+    if quant_type == QuantType.per_Token and quant_dtype == dtypes.fp8:
+        return _rmsnorm_gated_per_token_quant(x, weight, gate, eps, quant_dtype)
+    # Only the no-quant (bf16) path remains. Any other requested scheme is
+    # unsupported here -- fail loud rather than silently feed bf16 activations to
+    # a GEMM that expects quantized input.
+    assert quant_type in (None, QuantType.No), (
+        "rmsnorm_gated only fuses per-token FP8 quant; got "
+        f"quant_type={quant_type}, quant_dtype={quant_dtype}"
+    )
     h = x.shape[-1]
     x2 = x.reshape(-1, h)
     m = x2.shape[0]
@@ -156,6 +225,59 @@ def rmsnorm_gated(
     return y.reshape_as(x)
 
 
+def _rmsnorm_gated_per_token_quant(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    gate: torch.Tensor,
+    eps: float,
+    quant_dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-head sigmoid-gated RMSNorm fused to per-token quant.
+
+    Same math as ``rmsnorm_gated`` (rmsnorm(x) over the last dim * weight *
+    sigmoid(gate)), but instead of a bf16 output it emits a ``quant_dtype`` tensor
+    plus one per-token scale (amax / dtype_max) over the flattened
+    ``heads * head_dim`` row, ready for o_proj's per-token a8w8 GEMM. ``gate`` may
+    be strided (a column slice of the fused in_proj output).
+
+    Returns ``(out [t, heads*head_dim], scale [t, 1] float32)``.
+    """
+    assert x.ndim == 3, f"expected [t, heads, head_dim], got {tuple(x.shape)}"
+    t, heads, H = x.shape
+    fp8_max = float(torch.finfo(quant_dtype).max)
+    if not _HAS_TRITON or t == 0 or H > 8192:
+        normed = _rmsnorm_gated_torch(x, weight, gate, eps).reshape(t, heads * H)
+        return get_hip_quant(QuantType.per_Token)(normed, quant_dtype=quant_dtype)
+    x = x.contiguous()
+    out = torch.empty((t, heads * H), dtype=quant_dtype, device=x.device)
+    scale = torch.empty((t, 1), dtype=torch.float32, device=x.device)
+    if gate.ndim == 3:
+        stride_g_outer, stride_g_head = gate.stride(0), gate.stride(1)
+    else:
+        # 2D [t, heads*H]: one logical head per row; head term drops out.
+        stride_g_outer, stride_g_head = gate.stride(0), 0
+    BLOCK = triton.next_power_of_2(H)
+    _rmsnorm_gated_fp8_per_token_kernel[(t,)](
+        x,
+        weight,
+        gate,
+        out,
+        scale,
+        H,
+        float(eps),
+        fp8_max,
+        x.stride(0),
+        x.stride(1),
+        stride_g_outer,
+        stride_g_head,
+        out.stride(0),
+        HEADS=heads,
+        HEADS_POW2=triton.next_power_of_2(heads),
+        BLOCK=BLOCK,
+    )
+    return out, scale
+
+
 # --------------------------------------------------------------------------- #
 # torch references (also the fallback when triton is unavailable)
 # --------------------------------------------------------------------------- #

atom/model_ops/linear.py

@@ -373,6 +373,24 @@ def gemm_a8w8_per_token_impl(
     )
 
 
+def _can_use_a8w8_preshuffle(output_size: int, input_size: int) -> bool:
+    """Whether an a8w8 weight can use the AITER bpreshuffle GEMM as-is.
+
+    ``shuffle_weight(..., layout=(16, 16))`` packs the output (N) dim in 16-row
+    tiles and the input (K) dim in ``BK = IK * 2 = 32``-col tiles (it asserts
+    ``x.shape[-1] % 32 == 0``). So N must be 16-aligned and K must be 32-aligned.
+    """
+    return output_size % 16 == 0 and input_size % 32 == 0
+
+
+def _a8w8_preshuffle_output_padding(output_size: int) -> int:
+    """Rows needed to pad an a8w8 weight's output dim (N) up to the GEMM's N-tile
+    (128). Returns 0 when already tile-aligned. Padding N to 128 also makes it
+    16-aligned, so the tuned preshuffle GEMM can run instead of falling back."""
+    remainder = output_size % 128
+    return 0 if remainder == 0 else 128 - remainder
+
+
 class LinearBase(nn.Module):
     def __init__(
         self,
@@ -487,6 +505,7 @@ def __init__(
             self.weight_scale.weight_loader = self.weight_loader
         self.need_normalize_e4m3fn_to_e4m3fnuz = params_dtype == torch.float8_e4m3fnuz
         self.quant_func = get_hip_quant(self.quant_type)
+        self.is_output_padded = False
 
     @staticmethod
     def weight_loader_process(
@@ -688,6 +707,8 @@ def online_quantize_weight(self):
         }
 
     def process_weights_after_loading(self):
+        if self.weight.numel() == 0:
+            return
         # Re-quantize before process_weights if online quantization is enabled
         if self.quant_config is not None and self.quant_config.online_quant:
             self.online_quantize_weight()
@@ -773,16 +794,52 @@ def process_weights_after_loading(self):
                     self, "needs_preshuffled_weight", False
                 ):
                     need_shuffle = True
-            if need_shuffle:
-                if self.weight.dim() == 2:
-                    shuffle_weights(self.weight)
+            if need_shuffle and self.weight.dim() == 2:
+                self.is_output_padded = self._maybe_pad_a8w8_preshuffle_output()
+                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 (
             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)
 
+    def _maybe_pad_a8w8_preshuffle_output(self) -> bool:
+        if not (
+            self.quant_type == QuantType.per_Token and self.params_dtype == dtypes.fp8
+        ):
+            return False
+        if self.weight.dim() != 2:
+            return False
+        output_size, input_size = self.weight.shape
+        padding_size = _a8w8_preshuffle_output_padding(output_size)
+        if not _can_use_a8w8_preshuffle(output_size + padding_size, input_size):
+            # Padding the output (N) cannot make this weight preshuffle-able, i.e.
+            # the input dim K is not 32-aligned. Fail loudly here rather than let
+            # shuffle_weights hit its cryptic `x.shape[-1] % 32 == 0` assertion.
+            raise RuntimeError(
+                f"{self.prefix}: a8w8 bpreshuffle GEMM requires K % 32 == 0, got "
+                f"K={input_size}. Align K or run this layer via the triton a8w8 "
+                f"path (ATOM_USE_TRITON_GEMM=1)."
+            )
+        if padding_size == 0:
+            return False
+        self._output_size_before_padding = output_size
+        self.weight.data = torch.nn.functional.pad(
+            self.weight.data, (0, 0, 0, padding_size)
+        )
+        ws = self.weight_scale.data
+        self.weight_scale.data = torch.cat(
+            [ws, ws.new_ones((padding_size, *ws.shape[1:]))], dim=0
+        )
+        # Bias is also per-output-channel
+        if self.bias is not None:
+            b = self.bias.data
+            self.bias.data = torch.cat(
+                [b, b.new_zeros((padding_size, *b.shape[1:]))], dim=0
+            )
+        return True
+
     # linear mark trace shape/dtype helper
     def get_trace_prefix(
         self,
@@ -925,6 +982,9 @@ def forward(
                 )
                 if self.bias is not None:
                     y += self.bias
+        if self.is_output_padded:
+            # Drop the padded output rows
+            y = y[..., : self._output_size_before_padding]
         if self.tp_dim == 1 and self.tp_size > 1 and self.reduce_results:
             y = tensor_model_parallel_all_reduce(y)
         return y

atom/model_ops/module_dispatch_ops.py

@@ -19,18 +19,19 @@
 the methods listed in each op's docstring.
 
 Currently registered:
-  - torch.ops.aiter.maybe_dual_stream_forward — V2/V3.2/V4 MoE
+  - torch.ops.aiter.maybe_dual_stream_forward — V2/V3.2/V4/K3 MoE
   - torch.ops.aiter.indexer_score_topk       — V4 sparse indexer
 """
 
 import torch
 
-from atom.config import get_current_atom_config
+from atom.config import CUDAGraphMode, get_current_atom_config
 from atom.utils import envs
 from atom.utils.custom_register import direct_register_custom_op
+from atom.utils.forward_context import get_current_cudagraph_runtime_mode
 
 # ---------------------------------------------------------------------------
-# Dual-stream MoE dispatch (V2 / V3.2 / V4)
+# Dual-stream MoE dispatch (V2 / V3.2 / V4 / K3)
 # ---------------------------------------------------------------------------
 #
 # Caller contract (the MoE module looked up by `layer_name`):
@@ -54,13 +55,12 @@ def maybe_dual_stream_forward(
     # Under TBO the two micro-batches already overlap on separate threads
     from atom.utils.tbo.ubatching import tbo_active
 
-    # PIECEWISE cudagraph only: dual_stream_moe_forward forks work onto
-    # `alt_stream` and does a caching-allocator alloc there; under PIECEWISE
-    # per-piece capture close the dual stream
-    compilation_config = get_current_atom_config().compilation_config
-    cudagraph_mode = getattr(compilation_config, "cudagraph_mode", None)
+    # Graph ownership belongs to the active frontend.  Only a concrete
+    # PIECEWISE runtime decision is unsafe here: per-piece capture closes over
+    # the main stream while this forward forks work onto `alt_stream`.  Eager
+    # NONE and whole-model FULL capture both support the fork/join topology.
     is_piecewise_cudagraph = (
-        cudagraph_mode is not None and cudagraph_mode.requires_piecewise_compilation()
+        get_current_cudagraph_runtime_mode() == CUDAGraphMode.PIECEWISE
     )
 
     if (