feat: support Nemotron non-gated MoE tuning on Triton 3.6

src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/E=128,N=1856,device_name=MTT_S5000.json

@@ -0,0 +1,18 @@
+{
+    "1": {
+        "BLOCK_SIZE_M": 32,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 4,
+        "num_stages": 1
+    },
+    "8": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 32,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 1
+    }
+}

src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/headdim=64,dstate=128,device_name=MTT_S5000,cache_dtype=float32.json

@@ -0,0 +1,11 @@
+{
+    "triton_version": "3.6.0",
+    "64": {
+        "BLOCK_SIZE_M": 8,
+        "num_warps": 8
+    },
+    "512": {
+        "BLOCK_SIZE_M": 8,
+        "num_warps": 4
+    }
+}

src/torchada/triton/autotune/fused_moe/tune_moe.py

@@ -140,6 +140,8 @@ class ModelEntry:
     shard_intermediate_size: int = 0
     topk: int = 0
     num_fused_shared_experts: int = 0
+    activation: str = "silu"
+    is_gated: bool = True
     dtype: torch.dtype = torch.float16
     block_shape: Optional[Tuple[int, int]] = None
 
@@ -167,6 +169,8 @@ def unique_key(self) -> Tuple:
             self.shard_intermediate_size,
             self.topk,
             self.num_fused_shared_experts,
+            self.activation,
+            self.is_gated,
             str(self.dtype),
             self.use_fp8,
             self.use_int8,
@@ -188,7 +192,7 @@ def validate_and_log_entries(entries: List[ModelEntry]) -> None:
     for i, e in enumerate(entries):
         logger.info(
             "[%d] model=%s tp=%d ep=%d experts=%d hidden=%d "
-            "intermediate=%d topk=%d shared=%d dtype=%s block=%s",
+            "intermediate=%d topk=%d shared=%d activation=%s gated=%s dtype=%s block=%s",
             i,
             e.path,
             e.tp_size,
@@ -198,6 +202,8 @@ def validate_and_log_entries(entries: List[ModelEntry]) -> None:
             e.shard_intermediate_size,
             e.topk,
             e.num_fused_shared_experts,
+            e.activation,
+            e.is_gated,
             e.dtype_str,
             e.block_shape,
         )
@@ -316,6 +322,8 @@ def benchmark_config(
     per_channel_quant: bool,
     block_shape: List[int] = None,
     num_fused_shared_experts: int = 0,
+    activation: str = "silu",
+    is_gated: bool = True,
     num_iters: int = 100,
 ) -> float:
     """Run the fused MoE kernel and return latency in microseconds."""
@@ -324,6 +332,10 @@ def benchmark_config(
     init_dtype = torch.float16 if use_fp8_w8a8 else dtype
     num_routed_experts = num_experts - num_fused_shared_experts
     assert num_routed_experts > 0
+    # ``shard_intermediate_size`` is the first GEMM output width (w1.shape[1]).
+    # Gated models use half of it as the second GEMM input width because w1
+    # contains gate and up projections; non-gated models use the full width.
+    w2_input_size = shard_intermediate_size // 2 if is_gated else shard_intermediate_size
     x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
 
     # Create random weights based on quantization type
@@ -338,7 +350,7 @@ def benchmark_config(
         w2 = torch.randint(
             -127,
             127,
-            (num_experts, hidden_size, shard_intermediate_size // 2),
+            (num_experts, hidden_size, w2_input_size),
             dtype=torch.int8,
             device=device,
         )
@@ -353,7 +365,7 @@ def benchmark_config(
         w2 = torch.randint(
             0,
             255,
-            (num_experts, hidden_size, shard_intermediate_size // 4),
+            (num_experts, hidden_size, w2_input_size // 2),
             dtype=torch.uint8,
             device=device,
         )
@@ -368,7 +380,7 @@ def benchmark_config(
         w2 = torch.randn(
             num_experts,
             hidden_size,
-            shard_intermediate_size // 2,
+            w2_input_size,
             dtype=init_dtype,
             device=device,
         )
@@ -385,7 +397,7 @@ def benchmark_config(
     w1_scale = w2_scale = a1_scale = a2_scale = None
     if use_int8_w8a16:
         w1_scale = torch.randn(
-            (num_experts, 2 * shard_intermediate_size),
+            (num_experts, shard_intermediate_size),
             dtype=torch.float32,
             device=device,
         )
@@ -396,7 +408,7 @@ def benchmark_config(
         n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n
         n_tiles_w2 = (hidden_size + block_n - 1) // block_n
         k_tiles_w1 = (hidden_size + block_k - 1) // block_k
-        k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k
+        k_tiles_w2 = (w2_input_size + block_k - 1) // block_k
         w1_scale = torch.randn(
             (num_experts, n_tiles_w1, k_tiles_w1),
             dtype=torch.bfloat16,
@@ -423,7 +435,7 @@ def benchmark_config(
             n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n
             n_tiles_w2 = (hidden_size + block_n - 1) // block_n
             k_tiles_w1 = (hidden_size + block_k - 1) // block_k
-            k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k
+            k_tiles_w2 = (w2_input_size + block_k - 1) // block_k
             w1_scale = torch.rand(
                 (num_experts, n_tiles_w1, k_tiles_w1),
                 dtype=torch.float32,
@@ -462,6 +474,8 @@ def run():
             top_k=topk,
             num_fused_shared_experts=num_fused_shared_experts,
             inplace=True,
+            activation=activation,
+            is_gated=is_gated,
         )
         with override_config(config):
             fused_moe(
@@ -548,6 +562,8 @@ def build_model_entries(args: argparse.Namespace) -> List[ModelEntry]:
             entry.shard_intermediate_size = params["shard_intermediate_size"]
             entry.topk = params["topk"]
             entry.num_fused_shared_experts = params.get("num_fused_shared_experts", 0)
+            entry.activation = params.get("activation", "silu")
+            entry.is_gated = params.get("is_gated", True)
             entry.dtype_str = _resolve_dtype_str(entry.dtype_str, params)
             entry.dtype = _resolve_torch_dtype(entry.dtype_str, params)
             entry.block_shape = tuple(params["block_shape"]) if params["block_shape"] else None
@@ -603,6 +619,8 @@ def build_model_entries(args: argparse.Namespace) -> List[ModelEntry]:
                     entry.shard_intermediate_size = params["shard_intermediate_size"]
                     entry.topk = params["topk"]
                     entry.num_fused_shared_experts = params.get("num_fused_shared_experts", 0)
+                    entry.activation = params.get("activation", "silu")
+                    entry.is_gated = params.get("is_gated", True)
                     entry.dtype_str = _resolve_dtype_str(entry.dtype_str, params)
                     entry.dtype = _resolve_torch_dtype(entry.dtype_str, params)
                     entry.block_shape = (
@@ -657,6 +675,8 @@ def _tune_worker(
                     entry.per_channel_quant,
                     list(entry.block_shape) if entry.block_shape else None,
                     entry.num_fused_shared_experts,
+                    activation=entry.activation,
+                    is_gated=entry.is_gated,
                     num_iters=10,
                 )
             except (triton.runtime.autotuner.OutOfResources, RuntimeError, AssertionError):
@@ -767,6 +787,7 @@ def run_tuning(entries: List[ModelEntry], batch_sizes: List[int], args: argparse
             entry.use_int4,
             entry.per_channel_quant,
             entry.block_shape,
+            is_gated=entry.is_gated,
         )
         sorted_batches = sorted(bs_to_config.keys())
         best_configs = {bs: sort_config(bs_to_config[bs]) for bs in sorted_batches}
@@ -808,7 +829,11 @@ def _benchmark_worker(
             )
             block_n = entry.block_shape[0] if entry.block_shape else 0
             block_k = entry.block_shape[1] if entry.block_shape else 0
-            N = entry.shard_intermediate_size // 2
+            N = (
+                entry.shard_intermediate_size // 2
+                if entry.is_gated
+                else entry.shard_intermediate_size
+            )
             if entry.use_int4:
                 N = N // 2
             op_config = get_moe_configs(
@@ -852,6 +877,8 @@ def _benchmark_worker(
                 entry.per_channel_quant,
                 list(entry.block_shape) if entry.block_shape else None,
                 entry.num_fused_shared_experts,
+                activation=entry.activation,
+                is_gated=entry.is_gated,
             )
             result_queue.put((entry, batch_size, kernel_time, None))
         except Exception as e:
@@ -975,7 +1002,17 @@ def main(args: argparse.Namespace) -> None:
     parser.add_argument(
         "--dtype",
         type=str,
-        choices=["auto", "fp8_w8a8", "int8_w8a16", "int8_w8a8", "int4_w4a16"],
+        choices=[
+            "auto",
+            "bf16",
+            "bfloat16",
+            "fp16",
+            "float16",
+            "fp8_w8a8",
+            "int8_w8a16",
+            "int8_w8a8",
+            "int4_w4a16",
+        ],
         default="auto",
         help="Quantization dtype.",
     )

src/torchada/triton/autotune/fused_moe/utils.py

@@ -1,6 +1,6 @@
 import json
 import os
-from typing import Dict, List, TypedDict
+from typing import Dict, List, Tuple, TypedDict
 
 import torch
 from transformers import AutoConfig
@@ -24,12 +24,39 @@ class BenchmarkConfig(TypedDict):
 
 
 def calculate_shard_intermediate_size(
-    intermediate_size: int, tp_size: int, ep_size: int = 1
+    intermediate_size: int,
+    tp_size: int,
+    ep_size: int = 1,
+    is_gated: bool = True,
 ) -> int:
     assert tp_size % ep_size == 0
     moe_tp_size = tp_size // ep_size
     assert intermediate_size % moe_tp_size == 0
-    return 2 * intermediate_size // moe_tp_size
+    # Gated projections (for example SwiGLU) store gate and up next to each
+    # other in w1.  Non-gated projections such as Nemotron-H's relu2 use one
+    # projection, so the w1 output width is the model's intermediate size.
+    multiplier = 2 if is_gated else 1
+    return multiplier * intermediate_size // moe_tp_size
+
+
+def infer_moe_activation(config) -> Tuple[str, bool]:
+    """Infer the activation name and projection layout from a HF config.
+
+    ``NemotronHForCausalLM`` passes ``activation_without_mul`` to its fused
+    MoE layer.  The HF config advertises ``relu2`` while the checkpoint has a
+    single projection (``relu2_no_mul``).  Keep all other architectures on
+    the historical gated-SiLU default unless their activation already carries
+    the explicit ``_no_mul`` suffix.
+    """
+    raw = getattr(config, "mlp_hidden_act", None)
+    if raw is None:
+        raw = getattr(config, "hidden_act", "silu")
+    activation = str(raw).replace("torch.", "").lower()
+    architectures = getattr(config, "architectures", None) or []
+    architecture = str(architectures[0]) if architectures else type(config).__name__
+    if architecture == "NemotronHForCausalLM" and not activation.endswith("_no_mul"):
+        activation = f"{activation}_no_mul"
+    return activation, not activation.endswith("_no_mul")
 
 
 def get_num_shared_experts(config, disable_shared_experts_fusion: bool) -> int:
@@ -146,6 +173,7 @@ def get_model_config(
     config = _load_model_config(model_name)
 
     architecture = config.architectures[0]
+    activation, is_gated = infer_moe_activation(config)
     quant_dtype_str = infer_quant_dtype_str(config)
     block_shape = None
     if hasattr(config, "quantization_config") and "weight_block_size" in config.quantization_config:
@@ -247,7 +275,9 @@ def get_model_config(
         topk = config.num_experts_per_tok
         intermediate_size = config.intermediate_size
 
-    shard_intermediate_size = calculate_shard_intermediate_size(intermediate_size, tp_size, ep_size)
+    shard_intermediate_size = calculate_shard_intermediate_size(
+        intermediate_size, tp_size, ep_size, is_gated=is_gated
+    )
 
     return {
         "num_experts": E,
@@ -259,6 +289,8 @@ def get_model_config(
         "architecture": architecture,
         "num_fused_shared_experts": num_fused_shared_experts,
         "quant_dtype_str": quant_dtype_str,
+        "activation": activation,
+        "is_gated": is_gated,
     }
 
 
@@ -333,6 +365,7 @@ def get_config_filename(
     use_int4_w4a16: bool,
     per_channel_quant: bool,
     block_shape: List[int],
+    is_gated: bool = True,
 ) -> str:
     dtype_str = get_config_dtype_str(
         dtype,
@@ -344,7 +377,7 @@ def get_config_filename(
 
     # NOTE(woosuk): The current naming convention uses w2.shape[2], which
     # is the intermediate size after silu_and_mul.
-    N = shard_intermediate_size // 2
+    N = shard_intermediate_size // 2 if is_gated else shard_intermediate_size
     if use_int4_w4a16:
         N = N // 2
 

src/torchada/triton/kernels/moe/kernel.py

@@ -668,6 +668,10 @@ def invoke_fused_moe_kernel(
     fuse_add_to_output: bool = False,
     add_output_mask: Optional[torch.Tensor] = None,
 ) -> None:
+    config = dict(config)
+    split_k = config.pop("SPLIT_K", 1)
+    if split_k != 1:
+        raise ValueError("The torchada fused MoE kernel only supports SPLIT_K=1")
     assert topk_weights.stride(1) == 1
     assert sorted_token_ids.stride(0) == 1