add torch compile guard + tdm descriptor fix in routing.py

aiter/ops/triton/_gluon_kernels/gfx1250/moe/moe_op_gemm_a8w4.py

@@ -84,7 +84,7 @@ def unswizzle_mx_scale_gfx1250(
     return scale_buffer_slice
 
 
-@gluon.jit(launch_metadata=matmul_launch_metadata)
+@gluon.jit(launch_metadata=matmul_launch_metadata, do_not_specialize=["num_tokens"])
 def _moe_gemm_a8w4_decode(
     Y,
     stride_y_m,

aiter/ops/triton/_triton_kernels/moe/moe_routing/routing.py

@@ -48,7 +48,11 @@ def _routing_compute_indx(
     LOAD_SIZE: tl.constexpr = N_EXPTS_ACT_PAD * BLOCK_M
     local_offs = tl.arange(0, LOAD_SIZE)
     offs = pid_m * BLOCK_M * N_EXPTS_ACT + local_offs
-    if USE_TDM and EVEN_M and N_EXPTS_ACT == N_EXPTS_ACT_PAD:
+    # TDM tensor descriptors require >=16 bytes in the last dim. The expert-index
+    # load is int16 (2 bytes), so LOAD_SIZE must be >=8 elements; for tiny routing
+    # tiles (e.g. decode bs=1, where BLOCK_M=1 -> LOAD_SIZE=N_EXPTS_ACT_PAD) fall
+    # back to the functionally-identical plain-load branch below.
+    if USE_TDM and EVEN_M and N_EXPTS_ACT == N_EXPTS_ACT_PAD and LOAD_SIZE >= 8:
         expt_desc = tl.make_tensor_descriptor(
             base=ExptIndx + pid_m * BLOCK_M * N_EXPTS_ACT,
             shape=(1, LOAD_SIZE),
@@ -122,7 +126,11 @@ def _routing_compute_indx_fused(
     LOAD_SIZE: tl.constexpr = N_EXPTS_ACT_PAD * BLOCK_M
     local_offs = tl.arange(0, LOAD_SIZE)
     offs = local_offs
-    if USE_TDM and EVEN_M and N_EXPTS_ACT == N_EXPTS_ACT_PAD:
+    # TDM tensor descriptors require >=16 bytes in the last dim. The expert-index
+    # load is int16 (2 bytes), so LOAD_SIZE must be >=8 elements; for tiny routing
+    # tiles (e.g. decode bs=1, where BLOCK_M=1 -> LOAD_SIZE=N_EXPTS_ACT_PAD) fall
+    # back to the functionally-identical plain-load branch below.
+    if USE_TDM and EVEN_M and N_EXPTS_ACT == N_EXPTS_ACT_PAD and LOAD_SIZE >= 8:
         expt_desc = tl.make_tensor_descriptor(
             base=ExptIndx,
             shape=(1, LOAD_SIZE),

aiter/ops/triton/normalization/fused_rmsnorm_add.py

@@ -1,22 +1,28 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
+from typing import Tuple
+
 import torch
 import triton
 from aiter.ops.triton.utils._triton.arch_info import get_arch
+from aiter.jit.utils.torch_guard import torch_compile_guard
+from aiter.ops.triton._gluon_kernels.gfx1250.norm.fused_rmsnorm_add import (
+    _gluon_fused_rms_kernel,
+)
+from aiter.ops.triton._triton_kernels.normalization.fused_rmsnorm_add import (
+    _triton_fused_rms_kernel,
+)
 
 
-def fused_rmsnorm_add(x, weight, epsilon, res1=None):
-    """RMSNorm over the last dim of a 2D tensor, with an optional residual add.
-
-    x:      (M, N) tensor (bf16/fp16). Made contiguous if it isn't already.
-    weight: (N,) tensor.
-    res1:   optional (M, N) residual; when given, computes x += res1 first and
-            returns (out, out_res1) where out_res1 is the pre-norm sum.
-    Returns out (M, N) if res1 is None, else (out, out_res1).
+def _fused_rmsnorm_add_core(x, weight, epsilon, res1):
+    """Shared RMSNorm (+ optional residual add) launcher.
 
-    Dispatches to a gfx1250 Gluon kernel when running on gfx1250, otherwise
-    falls back to a portable Triton kernel.
+    Returns ``(out, out_res1)`` where ``out_res1`` is ``None`` when ``res1`` is
+    ``None``. The two guarded entry points below wrap this with a *fixed* return
+    type each, so they can register as torch custom ops (a single op cannot have
+    a Tensor-or-tuple return schema). Dispatches to the gfx1250 Gluon kernel on
+    gfx1250, otherwise a portable Triton kernel.
     """
     assert x.dim() == 2, "fused_rmsnorm_add expects a 2D tensor"
     M, N = x.shape
@@ -52,9 +58,6 @@ def fused_rmsnorm_add(x, weight, epsilon, res1=None):
         out_res1_stride_m = out_res1.stride(0)
 
     if get_arch() == "gfx1250":
-        from aiter.ops.triton._gluon_kernels.gfx1250.norm.fused_rmsnorm_add import (
-            _gluon_fused_rms_kernel,
-        )
 
         BLOCK_SIZE_M = 1
         grid = (triton.cdiv(M, BLOCK_SIZE_M),)
@@ -76,9 +79,6 @@ def fused_rmsnorm_add(x, weight, epsilon, res1=None):
             FIRST_INPUT_RES=(res1 is not None),
         )
     else:
-        from aiter.ops.triton._triton_kernels.normalization.fused_rmsnorm_add import (
-            _triton_fused_rms_kernel,
-        )
 
         grid = (M,)
         _triton_fused_rms_kernel[grid](
@@ -98,6 +98,53 @@ def fused_rmsnorm_add(x, weight, epsilon, res1=None):
             FIRST_INPUT_RES=(res1 is not None),
         )
 
-    if res1 is not None:
-        return out1, out_res1
+    return out1, out_res1
+
+
+def _fused_rmsnorm_fake(
+    x: torch.Tensor, weight: torch.Tensor, epsilon: float
+) -> torch.Tensor:
+    return torch.empty_like(x)
+
+
+@torch_compile_guard(gen_fake=_fused_rmsnorm_fake)
+def _fused_rmsnorm(
+    x: torch.Tensor, weight: torch.Tensor, epsilon: float
+) -> torch.Tensor:
+    out1, _ = _fused_rmsnorm_add_core(x, weight, epsilon, None)
     return out1
+
+
+def _fused_rmsnorm_add_fake(
+    x: torch.Tensor, weight: torch.Tensor, epsilon: float, res1: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    return torch.empty_like(x), torch.empty_like(x)
+
+
+@torch_compile_guard(gen_fake=_fused_rmsnorm_add_fake)
+def _fused_rmsnorm_add(
+    x: torch.Tensor, weight: torch.Tensor, epsilon: float, res1: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    out1, out_res1 = _fused_rmsnorm_add_core(x, weight, epsilon, res1)
+    return out1, out_res1
+
+
+def fused_rmsnorm_add(x, weight, epsilon, res1=None):
+    """RMSNorm over the last dim of a 2D tensor, with an optional residual add.
+
+    x:      (M, N) tensor (bf16/fp16). Made contiguous if it isn't already.
+    weight: (N,) tensor.
+    res1:   optional (M, N) residual; when given, computes x += res1 first and
+            returns (out, out_res1) where out_res1 is the pre-norm sum.
+    Returns out (M, N) if res1 is None, else (out, out_res1).
+
+    Dispatches to a gfx1250 Gluon kernel when running on gfx1250, otherwise
+    falls back to a portable Triton kernel. The actual compute is registered as
+    two torch custom ops (``_fused_rmsnorm`` / ``_fused_rmsnorm_add``) so the
+    path is safe under torch.compile / CUDAGraph capture; this thin dispatcher
+    selects the right one based on whether a residual was supplied (a static
+    branch at trace time).
+    """
+    if res1 is None:
+        return _fused_rmsnorm(x, weight, epsilon)
+    return _fused_rmsnorm_add(x, weight, epsilon, res1)