feat: add Nemotron 3.5 S5000 fused-MoE config

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,
@@ -316,6 +320,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."""
@@ -325,6 +331,11 @@ def benchmark_config(
     num_routed_experts = num_experts - num_fused_shared_experts
     assert num_routed_experts > 0
     x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
+    # The production runner commonly uses ``inplace=True``.  Keep that mode
+    # for representative timings, but restore the same input before every
+    # timed iteration so one benchmark run cannot feed its output back as the
+    # next iteration's input.
+    x_initial = x.clone()
 
     # Create random weights based on quantization type
     if use_int8_w8a16 or use_int8_w8a8:
@@ -338,7 +349,11 @@ def benchmark_config(
         w2 = torch.randint(
             -127,
             127,
-            (num_experts, hidden_size, shard_intermediate_size // 2),
+            (
+                num_experts,
+                hidden_size,
+                shard_intermediate_size // 2 if is_gated else shard_intermediate_size,
+            ),
             dtype=torch.int8,
             device=device,
         )
@@ -353,7 +368,11 @@ def benchmark_config(
         w2 = torch.randint(
             0,
             255,
-            (num_experts, hidden_size, shard_intermediate_size // 4),
+            (
+                num_experts,
+                hidden_size,
+                (shard_intermediate_size // 2 if is_gated else shard_intermediate_size) // 2,
+            ),
             dtype=torch.uint8,
             device=device,
         )
@@ -368,7 +387,7 @@ def benchmark_config(
         w2 = torch.randn(
             num_experts,
             hidden_size,
-            shard_intermediate_size // 2,
+            shard_intermediate_size // 2 if is_gated else shard_intermediate_size,
             dtype=init_dtype,
             device=device,
         )
@@ -396,7 +415,8 @@ 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
+        w2_k = shard_intermediate_size // 2 if is_gated else shard_intermediate_size
+        k_tiles_w2 = (w2_k + block_k - 1) // block_k
         w1_scale = torch.randn(
             (num_experts, n_tiles_w1, k_tiles_w1),
             dtype=torch.bfloat16,
@@ -423,7 +443,8 @@ 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
+            w2_k = shard_intermediate_size // 2 if is_gated else shard_intermediate_size
+            k_tiles_w2 = (w2_k + block_k - 1) // block_k
             w1_scale = torch.rand(
                 (num_experts, n_tiles_w1, k_tiles_w1),
                 dtype=torch.float32,
@@ -447,7 +468,25 @@ def benchmark_config(
     )
     topk_output = select_experts(x, input_gating, topk_config)
 
+    def reference_first_token() -> torch.Tensor:
+        """Compute one BF16 non-gated token with ordinary torch matmuls.
+
+        This is intentionally a small semantic gate for the Nemotron path;
+        timing still uses the Triton kernel for the complete batch.
+        """
+
+        token = x_initial[0].float()
+        result = torch.zeros(hidden_size, dtype=torch.float32, device=device)
+        for choice in range(topk):
+            expert = topk_output.topk_ids[0, choice]
+            gate_up = torch.matmul(token, w1[expert].float().transpose(0, 1))
+            activated = torch.relu(gate_up).square()
+            down = torch.matmul(activated, w2[expert].float().transpose(0, 1))
+            result.add_(down, alpha=float(topk_output.topk_weights[0, choice]))
+        return result
+
     def prepare(i: int):
+        x.copy_(x_initial)
         new_topk_output = select_experts(x, gating_output[i], topk_config)
         topk_output.topk_weights.copy_(new_topk_output.topk_weights)
         topk_output.topk_ids.copy_(new_topk_output.topk_ids)
@@ -461,7 +500,12 @@ def run():
             intermediate_size_per_partition=shard_intermediate_size,
             top_k=topk,
             num_fused_shared_experts=num_fused_shared_experts,
-            inplace=True,
+            # Keep the input immutable.  In-place output turns repeated
+            # warmups and graph capture into a feedback loop where each run
+            # consumes the previous run's result.
+            inplace=False,
+            activation=activation,
+            is_gated=is_gated,
         )
         with override_config(config):
             fused_moe(
@@ -483,22 +527,65 @@ def run():
             )
 
     # Warmup & JIT
-    run()
+    output = run()
     torch.cuda.synchronize()
 
+    if dtype == torch.bfloat16 and not is_gated:
+        if not torch.isfinite(output).all():
+            raise RuntimeError("BF16 non-gated MoE produced non-finite output")
+        reference = reference_first_token()
+        ref_scale = reference.abs().amax()
+        if not torch.isfinite(reference).all() or ref_scale <= 1e-12:
+            raise RuntimeError("BF16 non-gated reference output is degenerate")
+        out_token = output[0].float()
+        torch.testing.assert_close(out_token, reference, rtol=5e-2, atol=1e-6)
+        if out_token.abs().amax() <= ref_scale * 1e-2:
+            raise RuntimeError("BF16 non-gated MoE output is unexpectedly near zero")
+
     use_graph = _env_flag("TORCHADA_TUNE_USE_GRAPH") and hasattr(torch.cuda, "CUDAGraph")
     if use_graph:
         graph = torch.cuda.CUDAGraph()
         with torch.cuda.graph(graph):
-            for _ in range(10):
-                run()
+            graph_output = run()
         torch.cuda.synchronize()
-        for _ in range(5):
-            graph.replay()
+        # Verify that replay observes a changed routing tensor.  A cached
+        # alignment result can silently make every replay execute the first
+        # route; this gate catches that before timing is reported.
+        prepare(0)
+        graph.replay()
         torch.cuda.synchronize()
+        first_ids = topk_output.topk_ids.clone()
+        first_output = graph_output.detach().clone()
+        route_index = next(
+            (
+                i
+                for i in range(1, num_iters)
+                if not torch.equal(
+                    first_ids,
+                    select_experts(x_initial, gating_output[i], topk_config).topk_ids,
+                )
+            ),
+            None,
+        )
+        if route_index is None:
+            raise RuntimeError("graph route replay gate needs two distinct routing inputs")
+        prepare(route_index)
+        graph.replay()
+        torch.cuda.synchronize()
+        if torch.equal(first_ids, topk_output.topk_ids):
+            raise RuntimeError("graph route replay gate did not mutate top-k expert ids")
+        second_output = graph_output.detach().clone()
+        output_scale = max(
+            float(first_output.float().abs().amax()),
+            float(second_output.float().abs().amax()),
+            1e-12,
+        )
+        if not torch.isfinite(second_output).all():
+            raise RuntimeError("graph replay produced non-finite MoE output")
+        if (first_output.float() - second_output.float()).abs().amax() <= output_scale * 1e-4:
+            raise RuntimeError("graph route replay output did not change with routing inputs")
     else:
-        for _ in range(5):
-            run()
+        output = run()
         torch.cuda.synchronize()
 
     # Flush L2 cache
@@ -511,11 +598,15 @@ def run():
 
     for i in range(num_iters):
         prepare(i)
+        # Flush immediately before each timed sample so cache state from the
+        # previous sample cannot bias the comparison between configurations.
+        cache_flush.zero_()
+        torch.cuda.synchronize()
         start_events[i].record()
         if use_graph:
             graph.replay()
         else:
-            run()
+            output = run()
         end_events[i].record()
     torch.cuda.synchronize()
 
@@ -548,6 +639,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 +696,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 +752,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 +864,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 +906,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 +954,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 +1079,16 @@ 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",
+            "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,35 @@ 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 activations (e.g. SwiGLU) store gate and up projections side by
+    # side, while non-gated activations (e.g. NemotronH relu2_no_mul) have a
+    # single projection. Keep the returned size equal to the checkpoint w1
+    # dimension so callers can construct valid benchmark shapes.
+    multiplier = 2 if is_gated else 1
+    return multiplier * intermediate_size // moe_tp_size
+
+
+def infer_moe_activation(config) -> Tuple[str, bool]:
+    """Return the activation spelling and projection layout for a HF config.
+
+    NemotronH passes ``activation_without_mul(config.mlp_hidden_act)`` to its
+    FusedMoE layer. Its config says ``relu2`` but the expert checkpoint stores a
+    single (non-gated) projection, so represent that semantic explicitly here.
+    Other model families retain the historical gated-SiLU default.
+    """
+    raw = getattr(config, "mlp_hidden_act", None)
+    if raw is None:
+        raw = getattr(config, "hidden_act", "silu")
+    activation = str(raw).replace("torch.", "").lower()
+    architecture = str((getattr(config, "architectures", None) or [""])[0])
+    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 +169,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 +271,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 +285,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 +361,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 +373,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