From 26a85c335bbfc390e029c7b88b4aff93306aea47 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Sun, 13 Sep 2026 11:00:36 +0800 Subject: [PATCH 01/22] fix: recognize Triton 3.6 MoE configs --- src/torchada/triton/runtime/fused_moe/config.py | 8 +++++++- tests/test_moe_config_dir.py | 8 ++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/src/torchada/triton/runtime/fused_moe/config.py b/src/torchada/triton/runtime/fused_moe/config.py index c34c3d6..69b97e5 100644 --- a/src/torchada/triton/runtime/fused_moe/config.py +++ b/src/torchada/triton/runtime/fused_moe/config.py @@ -71,7 +71,13 @@ def get_moe_configs( be picked and the associated configuration chosen to invoke the kernel. """ # Supported Triton versions, should be sorted from the newest to the oldest - supported_triton_versions = ["3.4.0", "3.3.1", "3.2.0", "3.1.0"] + supported_triton_versions = [ + "3.6.0", + "3.4.0", + "3.3.1", + "3.2.0", + "3.1.0", + ] # First look up if an optimized configuration is available in the configs # directory diff --git a/tests/test_moe_config_dir.py b/tests/test_moe_config_dir.py index 535d4e0..f63e4e2 100644 --- a/tests/test_moe_config_dir.py +++ b/tests/test_moe_config_dir.py @@ -55,6 +55,14 @@ def test_exact_triton_version_dir_wins(tmp_path, monkeypatch): assert resolved == str(tmp_path / "configs" / "triton_3_1_0") +def test_triton36_exact_version_dir_wins(tmp_path, monkeypatch): + for ver in ("triton_3_2_0", "triton_3_6_0"): + (tmp_path / "configs" / ver).mkdir(parents=True) + monkeypatch.setattr(fused_moe, "_installed_triton_version", lambda: "3.6.0") + resolved = fused_moe._vllm_tuned_config_dir(str(tmp_path)) + assert resolved == str(tmp_path / "configs" / "triton_3_6_0") + + @pytest.mark.parametrize( "raw", ["3.2.0", "3.2.0.post1", "3.2.0rc1", "3.2.0+git9d8d5e91", "3.2.0.dev20260601"], From 100b5f29f6724d7d493a30cf4b017266f6d7b7c4 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 16:48:20 +0800 Subject: [PATCH 02/22] feat: model Nemotron non-gated MoE tuning semantics --- .../triton/autotune/fused_moe/tune_moe.py | 41 +++++++++++++++-- .../triton/autotune/fused_moe/utils.py | 39 ++++++++++++++-- tests/test_nemotron_moe_config.py | 46 +++++++++++++++++++ 3 files changed, 118 insertions(+), 8 deletions(-) create mode 100644 tests/test_nemotron_moe_config.py diff --git a/src/torchada/triton/autotune/fused_moe/tune_moe.py b/src/torchada/triton/autotune/fused_moe/tune_moe.py index 702c429..6f7cdb8 100755 --- a/src/torchada/triton/autotune/fused_moe/tune_moe.py +++ b/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,9 +320,21 @@ 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.""" + """Run the fused MoE kernel and return latency in microseconds. + + The current torchada Triton sequence only implements the gated activation + layout. Keep non-gated model entries visible in metadata, but fail closed + instead of tuning a shape with an incorrect projection width. + """ + if not is_gated: + raise NotImplementedError( + f"MoE activation {activation!r} is non-gated; relu2_no_mul benchmark " + "requires an activation-capable Triton sequence" + ) device = "cuda" torch.set_default_device(device) init_dtype = torch.float16 if use_fp8_w8a8 else dtype @@ -338,7 +354,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, ) @@ -462,6 +482,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 +570,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 +627,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 +683,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 +795,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 +837,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 +885,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: diff --git a/src/torchada/triton/autotune/fused_moe/utils.py b/src/torchada/triton/autotune/fused_moe/utils.py index 469b362..e49003b 100644 --- a/src/torchada/triton/autotune/fused_moe/utils.py +++ b/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 diff --git a/tests/test_nemotron_moe_config.py b/tests/test_nemotron_moe_config.py new file mode 100644 index 0000000..be33e34 --- /dev/null +++ b/tests/test_nemotron_moe_config.py @@ -0,0 +1,46 @@ +from types import SimpleNamespace + +import torch + +from torchada.triton.autotune.fused_moe.utils import ( + calculate_shard_intermediate_size, + get_config_filename, + infer_moe_activation, +) + + +def test_nemotron_relu2_uses_non_gated_layout(): + config = SimpleNamespace( + architectures=["NemotronHForCausalLM"], + mlp_hidden_act="relu2", + ) + assert infer_moe_activation(config) == ("relu2_no_mul", False) + assert calculate_shard_intermediate_size(1856, tp_size=1, is_gated=False) == 1856 + + +def test_other_models_keep_gated_layout(): + config = SimpleNamespace(architectures=["Qwen3MoeForCausalLM"], hidden_act="silu") + assert infer_moe_activation(config) == ("silu", True) + assert calculate_shard_intermediate_size(512, tp_size=2) == 512 + + +def test_nemotron_config_filename_uses_w2_width(): + gated = get_config_filename( + 128, 3712, 2688, 6, torch.bfloat16, False, False, False, False, False, None + ) + nongated = get_config_filename( + 128, + 1856, + 2688, + 6, + torch.bfloat16, + False, + False, + False, + False, + False, + None, + is_gated=False, + ) + assert "E=128,N=1856" in gated + assert "E=128,N=1856" in nongated From f90fe33ffa7286a99abae0e0988028bb04accaa0 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:00:51 +0800 Subject: [PATCH 03/22] fix: implement non-gated relu2 MoE activation in tuner --- .../triton/runtime/fused_moe/fused_moe.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/src/torchada/triton/runtime/fused_moe/fused_moe.py b/src/torchada/triton/runtime/fused_moe/fused_moe.py index 44728b2..e288f45 100644 --- a/src/torchada/triton/runtime/fused_moe/fused_moe.py +++ b/src/torchada/triton/runtime/fused_moe/fused_moe.py @@ -272,11 +272,20 @@ def _fused_moe_kernel_sequence( topk_ids, ) - intermediate_cache2 = torch.empty( - (total_tokens, N // 2), - device=hidden_states.device, - dtype=hidden_states.dtype, - ) + if is_gated: + intermediate_cache2 = torch.empty( + (total_tokens, N // 2), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + gate, up = intermediate_cache1.chunk(2, dim=-1) + intermediate_cache2.copy_(torch.nn.functional.silu(gate) * up) + else: + intermediate_cache2 = torch.empty_like(intermediate_cache1) + # Nemotron-H's relu2 path is non-gated: relu(x)^2, with no second + # projection half to multiply. Keep this explicit so the tuner cannot + # benchmark an uninitialized intermediate buffer. + intermediate_cache2.copy_(torch.relu(intermediate_cache1).square()) intermediate_cache3 = torch.empty( (num_tokens, topk, w2.shape[1]), From 8e4089de8c62add73bbbabeb48795d6cc769eef2 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:09:10 +0800 Subject: [PATCH 04/22] feat: enable BF16 Nemotron MoE tuning --- .../triton/autotune/fused_moe/tune_moe.py | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/src/torchada/triton/autotune/fused_moe/tune_moe.py b/src/torchada/triton/autotune/fused_moe/tune_moe.py index 6f7cdb8..ea698bc 100755 --- a/src/torchada/triton/autotune/fused_moe/tune_moe.py +++ b/src/torchada/triton/autotune/fused_moe/tune_moe.py @@ -324,17 +324,7 @@ def benchmark_config( is_gated: bool = True, num_iters: int = 100, ) -> float: - """Run the fused MoE kernel and return latency in microseconds. - - The current torchada Triton sequence only implements the gated activation - layout. Keep non-gated model entries visible in metadata, but fail closed - instead of tuning a shape with an incorrect projection width. - """ - if not is_gated: - raise NotImplementedError( - f"MoE activation {activation!r} is non-gated; relu2_no_mul benchmark " - "requires an activation-capable Triton sequence" - ) + """Run the fused MoE kernel and return latency in microseconds.""" device = "cuda" torch.set_default_device(device) init_dtype = torch.float16 if use_fp8_w8a8 else dtype @@ -1010,7 +1000,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.", ) From d78fdb7a0943a689b80ddd25b3221ff55f577e7a Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:15:14 +0800 Subject: [PATCH 05/22] fix: use full Nemotron non-gated MoE projection width --- src/torchada/triton/autotune/fused_moe/tune_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/torchada/triton/autotune/fused_moe/tune_moe.py b/src/torchada/triton/autotune/fused_moe/tune_moe.py index ea698bc..f414b99 100755 --- a/src/torchada/triton/autotune/fused_moe/tune_moe.py +++ b/src/torchada/triton/autotune/fused_moe/tune_moe.py @@ -378,7 +378,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, ) From 58672a86e04e80c576c03345ab59de685e3022fc Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:26:51 +0800 Subject: [PATCH 06/22] fix: make Nemotron MoE tuner return combined outputs --- .../triton/runtime/fused_moe/fused_moe.py | 70 +++++++++++-------- 1 file changed, 42 insertions(+), 28 deletions(-) diff --git a/src/torchada/triton/runtime/fused_moe/fused_moe.py b/src/torchada/triton/runtime/fused_moe/fused_moe.py index e288f45..8899126 100644 --- a/src/torchada/triton/runtime/fused_moe/fused_moe.py +++ b/src/torchada/triton/runtime/fused_moe/fused_moe.py @@ -336,6 +336,19 @@ def _fused_moe_kernel_sequence( router_topk=topk, ) + if no_combine: + return intermediate_cache3 + + if _use_intermediate: + combined = intermediate_cache3 + if apply_router_weight_on_input: + combined = combined * topk_weights.to(combined.dtype).unsqueeze(-1) + combined = combined.float().sum(dim=1).to(out_hidden_states.dtype) + out_hidden_states.copy_(combined) + if routed_scaling_factor is not None: + out_hidden_states.mul_(routed_scaling_factor) + return out_hidden_states + def fused_experts_impl( hidden_states: torch.Tensor, @@ -506,32 +519,33 @@ def fused_moe( moe_runner_config.num_experts is None or moe_runner_config.num_experts != moe_runner_config.num_local_experts ) - fused_experts_impl( - hidden_states, - w1, - w2, - topk_weights, - topk_ids, - b1, - b2, - True, - moe_runner_config.activation, - moe_runner_config.is_gated, - moe_runner_config.apply_router_weight_on_input, - use_fp8_w8a8, - use_int8_w8a8, - use_int8_w8a16, - use_int4_w4a16, - per_channel_quant, - w1_scale, - w2_scale, - w1_zp, - w2_zp, - a1_scale, - a2_scale, - block_shape, - moe_runner_config.routed_scaling_factor, - moe_runner_config.gemm1_alpha, - moe_runner_config.gemm1_clamp_limit, - filter_expert, + return fused_experts_impl( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + b1=b1, + b2=b2, + inplace=True, + activation=moe_runner_config.activation, + is_gated=moe_runner_config.is_gated, + apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + use_int8_w8a16=use_int8_w8a16, + use_int4_w4a16=use_int4_w4a16, + per_channel_quant=per_channel_quant, + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_zp=w1_zp, + w2_zp=w2_zp, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + no_combine=False, + routed_scaling_factor=moe_runner_config.routed_scaling_factor, + gemm1_alpha=moe_runner_config.gemm1_alpha, + gemm1_limit=moe_runner_config.gemm1_clamp_limit, + filter_expert=filter_expert, ) From ec67ada8918d8a9b8906bfa14592a8e39b23f6bc Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:29:40 +0800 Subject: [PATCH 07/22] test: gate Nemotron MoE pipeline semantics --- tests/test_moe_pipeline.py | 159 ++++++++++++++++++++++++++++++ tests/test_nemotron_moe_config.py | 48 +++++++++ 2 files changed, 207 insertions(+) create mode 100644 tests/test_moe_pipeline.py diff --git a/tests/test_moe_pipeline.py b/tests/test_moe_pipeline.py new file mode 100644 index 0000000..a96f3ea --- /dev/null +++ b/tests/test_moe_pipeline.py @@ -0,0 +1,159 @@ +"""CPU contract tests for the out-of-tree fused MoE implementation. + +The Triton launch is replaced with a tiny reference GEMM. These tests cover +the Python pipeline and argument plumbing without requiring a CUDA/MUSA device. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from torchada.triton.runtime.fused_moe import fused_moe as moe + + +def _fake_gemm(a, w, bias, out, *args, **kwargs): + """Reference implementation of invoke_fused_moe_kernel for CPU tests.""" + ids, sorted_ids, expert_ids = args[4], args[5], args[6] + topk = int(args[9]) + call_index = getattr(_fake_gemm, "calls", 0) + _fake_gemm.calls = call_index + 1 + # The first launch receives [tokens, hidden], the second receives routed + # [tokens * topk, intermediate]. sorted_ids is identity in this test. + nrows = a.shape[0] + out_rows = out.reshape(-1, out.shape[-1]) + + def write(row, value): + out_rows[row].copy_(value) + + if call_index == 0: + expert = ids[:, 0] + for row, e in enumerate(expert.tolist()): + write(row, a[row].to(out.dtype) @ w[e].to(out.dtype).T) + if bias is not None: + out_rows[row].add_(bias[e].to(out.dtype)) + else: + for row in range(nrows): + token = row // topk + e = int(ids[token, row % topk]) + write(row, a[row].to(out.dtype) @ w[e].to(out.dtype).T) + if bias is not None: + out_rows[row].add_(bias[e].to(out.dtype)) + + +def _args(*, activation="relu2_no_mul", is_gated=False, no_combine=False, inplace=False): + tokens, hidden, inter, experts, topk = 2, 3, 4, 2, 1 + x = torch.tensor([[1.0, -2.0, 0.5], [-0.5, 2.0, 1.0]]) + w1 = ( + torch.arange(experts * inter * hidden, dtype=torch.float32).reshape(experts, inter, hidden) + / 10 + ) + w2_width = inter // 2 if is_gated else inter + w2 = ( + torch.arange(experts * hidden * w2_width, dtype=torch.float32).reshape( + experts, hidden, w2_width + ) + / 10 + ) + weights = torch.ones(tokens, topk) + ids = torch.tensor([[0], [1]], dtype=torch.long) + ident = torch.arange(tokens, dtype=torch.int32) + config = {"BLOCK_SIZE_M": 1} + return dict( + hidden_states=x, + w1=w1, + w2=w2, + topk_weights=weights, + topk_ids=ids, + sorted_token_ids=ident, + expert_ids=torch.zeros(tokens, dtype=torch.int32), + num_tokens_post_padded=torch.tensor([tokens], dtype=torch.int32), + config=config, + down_config=config, + down_moe_use_tma=False, + b1=None, + b2=None, + use_fp8_w8a8=False, + use_int8_w8a8=False, + use_int8_w8a16=False, + use_int4_w4a16=False, + per_channel_quant=False, + w1_scale=None, + w2_scale=None, + w1_zp=None, + w2_zp=None, + a1_scale=None, + a2_scale=None, + block_shape=None, + activation=activation, + is_gated=is_gated, + no_combine=no_combine, + inplace=inplace, + apply_router_weight_on_input=False, + routed_scaling_factor=None, + gemm1_alpha=None, + gemm1_limit=None, + filter_expert=False, + ) + + +@pytest.mark.parametrize("activation,is_gated", [("relu2_no_mul", False), ("silu", True)]) +@pytest.mark.parametrize("no_combine,inplace", [(False, False), (False, True), (True, False)]) +def test_fused_moe_pipeline_activation_and_output_modes( + monkeypatch, activation, is_gated, no_combine, inplace +): + monkeypatch.setattr(moe, "invoke_fused_moe_kernel", _fake_gemm) + _fake_gemm.calls = 0 + kwargs = _args(activation=activation, is_gated=is_gated, no_combine=no_combine, inplace=inplace) + original = kwargs["hidden_states"].clone() + result = moe._fused_moe_kernel_sequence(**kwargs) + assert isinstance(result, torch.Tensor) + assert result.shape == ((2, 1, 3) if no_combine else (2, 3)) + if inplace: + assert result.data_ptr() == kwargs["hidden_states"].data_ptr() + # Ensure both activation branches actually contribute finite values. + assert torch.isfinite(result).all() + if inplace: + assert not torch.equal(result, original) + + +def test_fused_moe_rejects_unknown_activation(monkeypatch): + monkeypatch.setattr(moe, "invoke_fused_moe_kernel", _fake_gemm) + _fake_gemm.calls = 0 + with pytest.raises(ValueError, match="activation"): + moe._fused_moe_kernel_sequence(**_args(activation="gelu", is_gated=False)) + + +def test_fused_moe_wrapper_forwards_runner_options_by_keyword(monkeypatch): + captured = {} + + def fake_impl(*args, **kwargs): + captured["args"] = args + captured.update(kwargs) + return torch.tensor([7.0]) + + monkeypatch.setattr(moe, "fused_experts_impl", fake_impl) + topk = (torch.ones(1, 1), torch.zeros(1, 1, dtype=torch.long), None) + runner = SimpleNamespace( + num_experts=2, + num_local_experts=2, + activation="relu2_no_mul", + is_gated=False, + apply_router_weight_on_input=True, + routed_scaling_factor=1.5, + gemm1_alpha=0.25, + gemm1_clamp_limit=2.0, + no_combine=True, + inplace=False, + ) + x = torch.ones(1, 3) + out = moe.fused_moe(x, torch.ones(2, 4, 3), torch.ones(2, 3, 4), topk, runner) + assert out.item() == 7.0 + assert captured["activation"] == "relu2_no_mul" + assert captured["is_gated"] is False + assert captured["no_combine"] is True + assert captured["inplace"] is False + assert captured["apply_router_weight_on_input"] is True + assert captured["routed_scaling_factor"] == 1.5 + assert captured["gemm1_alpha"] == 0.25 + assert captured["gemm1_limit"] == 2.0 diff --git a/tests/test_nemotron_moe_config.py b/tests/test_nemotron_moe_config.py index be33e34..f93cb96 100644 --- a/tests/test_nemotron_moe_config.py +++ b/tests/test_nemotron_moe_config.py @@ -44,3 +44,51 @@ def test_nemotron_config_filename_uses_w2_width(): ) assert "E=128,N=1856" in gated assert "E=128,N=1856" in nongated + + +def test_nemotron_benchmark_uses_single_projection_width(monkeypatch): + """The BF16 Nemotron benchmark must allocate w2 with N, rather than N/2.""" + from torchada.triton.autotune.fused_moe import tune_moe + + seen = [] + + def fake_randn(*shape, **kwargs): + seen.append(tuple(shape)) + # Stop before the benchmark starts routing/timing. This keeps the + # shape assertion CPU-only and avoids allocating the full checkpoint + # dimensions. + if len(seen) == 3: + raise RuntimeError("stop after expert weights") + return torch.empty((1,), dtype=kwargs.get("dtype", torch.float32)) + + monkeypatch.setattr(tune_moe.torch, "set_default_device", lambda *_: None) + monkeypatch.setattr(tune_moe.torch, "randn", fake_randn) + config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 1, + "num_stages": 1, + } + try: + tune_moe.benchmark_config( + config, + 1, + 128, + 1856, + 2688, + 6, + torch.bfloat16, + False, + False, + False, + False, + False, + activation="relu2_no_mul", + is_gated=False, + ) + except RuntimeError as exc: + assert str(exc) == "stop after expert weights" + assert seen[1] == (128, 1856, 2688) + assert seen[2] == (128, 2688, 1856) From 1737c4fe4b1d0029361e51fcdc2b854b493aa650 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:37:49 +0800 Subject: [PATCH 08/22] test: cover topk MoE combine semantics --- tests/test_moe_pipeline.py | 41 +++++++++++++++++++++++++++++--------- 1 file changed, 32 insertions(+), 9 deletions(-) diff --git a/tests/test_moe_pipeline.py b/tests/test_moe_pipeline.py index a96f3ea..088758d 100644 --- a/tests/test_moe_pipeline.py +++ b/tests/test_moe_pipeline.py @@ -15,7 +15,9 @@ def _fake_gemm(a, w, bias, out, *args, **kwargs): """Reference implementation of invoke_fused_moe_kernel for CPU tests.""" ids, sorted_ids, expert_ids = args[4], args[5], args[6] - topk = int(args[9]) + # The down projection launch passes top_k=1 because inputs are already + # expanded; routing metadata still carries the original top-k width. + topk = ids.shape[1] call_index = getattr(_fake_gemm, "calls", 0) _fake_gemm.calls = call_index + 1 # The first launch receives [tokens, hidden], the second receives routed @@ -27,9 +29,10 @@ def write(row, value): out_rows[row].copy_(value) if call_index == 0: - expert = ids[:, 0] - for row, e in enumerate(expert.tolist()): - write(row, a[row].to(out.dtype) @ w[e].to(out.dtype).T) + for row in range(nrows): + token = row // topk + e = int(ids[token, row % topk]) + write(row, a[token].to(out.dtype) @ w[e].to(out.dtype).T) if bias is not None: out_rows[row].add_(bias[e].to(out.dtype)) else: @@ -42,7 +45,7 @@ def write(row, value): def _args(*, activation="relu2_no_mul", is_gated=False, no_combine=False, inplace=False): - tokens, hidden, inter, experts, topk = 2, 3, 4, 2, 1 + tokens, hidden, inter, experts, topk = 2, 3, 4, 2, 2 x = torch.tensor([[1.0, -2.0, 0.5], [-0.5, 2.0, 1.0]]) w1 = ( torch.arange(experts * inter * hidden, dtype=torch.float32).reshape(experts, inter, hidden) @@ -55,9 +58,9 @@ def _args(*, activation="relu2_no_mul", is_gated=False, no_combine=False, inplac ) / 10 ) - weights = torch.ones(tokens, topk) - ids = torch.tensor([[0], [1]], dtype=torch.long) - ident = torch.arange(tokens, dtype=torch.int32) + weights = torch.tensor([[0.25, 0.75], [0.4, 0.6]]) + ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + ident = torch.arange(tokens * topk, dtype=torch.int32) config = {"BLOCK_SIZE_M": 1} return dict( hidden_states=x, @@ -108,11 +111,31 @@ def test_fused_moe_pipeline_activation_and_output_modes( original = kwargs["hidden_states"].clone() result = moe._fused_moe_kernel_sequence(**kwargs) assert isinstance(result, torch.Tensor) - assert result.shape == ((2, 1, 3) if no_combine else (2, 3)) + assert result.shape == ((2, 2, 3) if no_combine else (2, 3)) if inplace: assert result.data_ptr() == kwargs["hidden_states"].data_ptr() # Ensure both activation branches actually contribute finite values. assert torch.isfinite(result).all() + # Check the complete GEMM1 -> activation -> GEMM2 pipeline against a + # compact PyTorch reference (the fake launcher only replaces the kernels). + expected_rows = [] + for token in range(kwargs["hidden_states"].shape[0]): + token_outputs = [] + for choice, expert in enumerate(kwargs["topk_ids"][token].tolist()): + gate_up = original[token] @ kwargs["w1"][expert].T + if is_gated: + width = gate_up.shape[-1] // 2 + activated = torch.nn.functional.silu(gate_up[:width]) * gate_up[width:] + else: + activated = torch.relu(gate_up).square() + token_outputs.append(activated @ kwargs["w2"][expert].T) + expected_rows.append(torch.stack(token_outputs)) + expected = torch.stack(expected_rows) + if no_combine: + assert torch.allclose(result, expected, atol=1e-5, rtol=1e-5) + else: + weighted = (expected * kwargs["topk_weights"].unsqueeze(-1)).sum(dim=1) + assert torch.allclose(result, weighted, atol=1e-5, rtol=1e-5) if inplace: assert not torch.equal(result, original) From ddc57579d6ec2cca5a1272e9994dacf4cb073a7b Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:38:24 +0800 Subject: [PATCH 09/22] fix: complete MoE activation and reduction semantics --- .../triton/runtime/fused_moe/fused_moe.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/src/torchada/triton/runtime/fused_moe/fused_moe.py b/src/torchada/triton/runtime/fused_moe/fused_moe.py index 8899126..70c511f 100644 --- a/src/torchada/triton/runtime/fused_moe/fused_moe.py +++ b/src/torchada/triton/runtime/fused_moe/fused_moe.py @@ -206,6 +206,10 @@ def _fused_moe_kernel_sequence( num_tokens = hidden_states.shape[0] E, N, _ = w1.shape topk = topk_ids.shape[1] + if (is_gated and activation != "silu") or (not is_gated and activation != "relu2_no_mul"): + raise ValueError( + f"Unsupported MoE activation/layout: activation={activation!r}, " f"is_gated={is_gated}" + ) compute_type = tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16 padded_tokens = ( @@ -309,7 +313,11 @@ def _fused_moe_kernel_sequence( ( out_slice if use_fused_moe_sum_all_reduce - else (intermediate_cache3 if _use_intermediate else out_hidden_states.unsqueeze(0)) + else ( + intermediate_cache3 + if no_combine or _use_intermediate + else out_hidden_states.unsqueeze(0) + ) ), a2_scale, w2_scale, @@ -380,6 +388,10 @@ def fused_experts_impl( gemm1_limit: Optional[float] = None, filter_expert: bool = True, ): + if (is_gated and activation != "silu") or (not is_gated and activation != "relu2_no_mul"): + raise ValueError( + f"Unsupported MoE activation/layout: activation={activation!r}, " f"is_gated={is_gated}" + ) padded_size = 128 if not (use_fp8_w8a8 or use_int8_w8a8) or block_shape is not None: padded_size = 0 @@ -527,7 +539,7 @@ def fused_moe( topk_ids=topk_ids, b1=b1, b2=b2, - inplace=True, + inplace=getattr(moe_runner_config, "inplace", False), activation=moe_runner_config.activation, is_gated=moe_runner_config.is_gated, apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input, @@ -543,7 +555,7 @@ def fused_moe( a1_scale=a1_scale, a2_scale=a2_scale, block_shape=block_shape, - no_combine=False, + no_combine=getattr(moe_runner_config, "no_combine", False), routed_scaling_factor=moe_runner_config.routed_scaling_factor, gemm1_alpha=moe_runner_config.gemm1_alpha, gemm1_limit=moe_runner_config.gemm1_clamp_limit, From 5dd124fa7c3b69f486ec72712c48c9e14f0e1c5b Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:40:19 +0800 Subject: [PATCH 10/22] test: cover all routed expert rows --- tests/test_moe_pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_moe_pipeline.py b/tests/test_moe_pipeline.py index 088758d..5ae0b17 100644 --- a/tests/test_moe_pipeline.py +++ b/tests/test_moe_pipeline.py @@ -22,8 +22,8 @@ def _fake_gemm(a, w, bias, out, *args, **kwargs): _fake_gemm.calls = call_index + 1 # The first launch receives [tokens, hidden], the second receives routed # [tokens * topk, intermediate]. sorted_ids is identity in this test. - nrows = a.shape[0] out_rows = out.reshape(-1, out.shape[-1]) + nrows = out_rows.shape[0] def write(row, value): out_rows[row].copy_(value) From 8182e68d6d98dfe50b8e967ea9b0d94a3002f7a5 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:41:18 +0800 Subject: [PATCH 11/22] test: model routed weight in fake down kernel --- tests/test_moe_pipeline.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_moe_pipeline.py b/tests/test_moe_pipeline.py index 5ae0b17..1d4e918 100644 --- a/tests/test_moe_pipeline.py +++ b/tests/test_moe_pipeline.py @@ -38,10 +38,15 @@ def write(row, value): else: for row in range(nrows): token = row // topk - e = int(ids[token, row % topk]) + choice = row % topk + e = int(ids[token, choice]) write(row, a[row].to(out.dtype) @ w[e].to(out.dtype).T) if bias is not None: out_rows[row].add_(bias[e].to(out.dtype)) + # The down launch receives MUL_ROUTED_WEIGHT=True when the + # caller asks the kernel to apply routing weights on output. + if args[8]: + out_rows[row].mul_(args[3][token, choice]) def _args(*, activation="relu2_no_mul", is_gated=False, no_combine=False, inplace=False): From 754ca7fa826ee8794745b501eb9b949b3dd7c6d3 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:42:46 +0800 Subject: [PATCH 12/22] test: cover python-side routed weighting --- tests/test_moe_pipeline.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_moe_pipeline.py b/tests/test_moe_pipeline.py index 1d4e918..d4c6645 100644 --- a/tests/test_moe_pipeline.py +++ b/tests/test_moe_pipeline.py @@ -97,7 +97,8 @@ def _args(*, activation="relu2_no_mul", is_gated=False, no_combine=False, inplac is_gated=is_gated, no_combine=no_combine, inplace=inplace, - apply_router_weight_on_input=False, + # Exercise Python-side routing-weight application for combined output. + apply_router_weight_on_input=True, routed_scaling_factor=None, gemm1_alpha=None, gemm1_limit=None, From 8528e3cc74c3c2de6c770a4d6c62a790e56633d8 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:52:47 +0800 Subject: [PATCH 13/22] fix: provide standalone MoE alignment fallback --- .../triton/runtime/fused_moe/fused_moe.py | 32 ++++++++++++++++--- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/src/torchada/triton/runtime/fused_moe/fused_moe.py b/src/torchada/triton/runtime/fused_moe/fused_moe.py index 70c511f..c008536 100644 --- a/src/torchada/triton/runtime/fused_moe/fused_moe.py +++ b/src/torchada/triton/runtime/fused_moe/fused_moe.py @@ -93,10 +93,34 @@ def moe_align_block_size( ) return sorted_ids, expert_ids, num_tokens_post_pad - except ImportError: - raise ImportError( - "No implementation of moe_align_block_size found. " "Please install sgl_kernel or vllm" - ) + except (ImportError, AttributeError): + # Standalone torchada tuning images may contain vLLM Python sources + # without the optional _moe_C extension. Keep the reference path + # usable for correctness/tuning; production vLLM uses the native op. + flat_ids = topk_ids.reshape(-1) + routes = [] + block_experts = [] + for expert in range(num_experts): + positions = torch.nonzero(flat_ids == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + pad = (-positions.numel()) % block_size + if pad: + positions = torch.cat( + [positions, torch.full_like(positions[:1], flat_ids.numel()).expand(pad)] + ) + routes.append(positions) + block_experts.extend([expert] * (positions.numel() // block_size)) + if routes: + sorted_routes = torch.cat(routes) + blocks = torch.tensor(block_experts, dtype=torch.int32, device=topk_ids.device) + else: + sorted_routes = torch.empty(0, dtype=torch.int32, device=topk_ids.device) + blocks = torch.empty(0, dtype=torch.int32, device=topk_ids.device) + sorted_ids[: sorted_routes.numel()].copy_(sorted_routes) + expert_ids[: blocks.numel()].copy_(blocks) + num_tokens_post_pad[0] = sorted_routes.numel() + return sorted_ids, expert_ids, num_tokens_post_pad def _prepare_fused_moe_run( From 44754027427299bdcd643c39eb5930c5b630ac70 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 19:05:41 +0800 Subject: [PATCH 14/22] fix: normalize torchada fused MoE split config --- src/torchada/triton/kernels/moe/kernel.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/torchada/triton/kernels/moe/kernel.py b/src/torchada/triton/kernels/moe/kernel.py index 8c0a762..960797c 100644 --- a/src/torchada/triton/kernels/moe/kernel.py +++ b/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 From f0eb59e9f887a3cd146cf2c8f01f37929459097f Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 19:09:27 +0800 Subject: [PATCH 15/22] fix: cache MoE routing metadata across graph capture --- .../triton/runtime/fused_moe/fused_moe.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/torchada/triton/runtime/fused_moe/fused_moe.py b/src/torchada/triton/runtime/fused_moe/fused_moe.py index c008536..8547530 100644 --- a/src/torchada/triton/runtime/fused_moe/fused_moe.py +++ b/src/torchada/triton/runtime/fused_moe/fused_moe.py @@ -11,6 +11,8 @@ try_get_optimal_moe_config, ) +_ALIGNMENT_CACHE: dict[int, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} + try: _support_tensor_descriptor = True except: @@ -47,6 +49,11 @@ def moe_align_block_size( ensuring divisibility by block_size. """ + cache_key = id(topk_ids) + is_capturing = bool(getattr(torch.cuda, "is_current_stream_capturing", lambda: False)()) + if is_capturing and cache_key in _ALIGNMENT_CACHE: + return _ALIGNMENT_CACHE[cache_key] + if topk_ids.numel() < num_experts + 1: max_num_tokens_padded = topk_ids.numel() * block_size else: @@ -73,7 +80,9 @@ def moe_align_block_size( cumsum_buffer, True, ) - return sorted_ids, expert_ids, num_tokens_post_pad + result = sorted_ids, expert_ids, num_tokens_post_pad + _ALIGNMENT_CACHE[cache_key] = result + return result except ImportError: pass @@ -91,7 +100,9 @@ def moe_align_block_size( num_tokens_post_pad, None, ) - return sorted_ids, expert_ids, num_tokens_post_pad + result = sorted_ids, expert_ids, num_tokens_post_pad + _ALIGNMENT_CACHE[cache_key] = result + return result except (ImportError, AttributeError): # Standalone torchada tuning images may contain vLLM Python sources @@ -120,7 +131,9 @@ def moe_align_block_size( sorted_ids[: sorted_routes.numel()].copy_(sorted_routes) expert_ids[: blocks.numel()].copy_(blocks) num_tokens_post_pad[0] = sorted_routes.numel() - return sorted_ids, expert_ids, num_tokens_post_pad + result = sorted_ids, expert_ids, num_tokens_post_pad + _ALIGNMENT_CACHE[cache_key] = result + return result def _prepare_fused_moe_run( From aa5ac5e0457bb0eac2b9c5e5a3999a4e61093075 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 19:40:48 +0800 Subject: [PATCH 16/22] feat: add Nemotron E128 N1856 S5000 MoE config --- .../E=128,N=1856,device_name=MTT_S5000.json | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=128,N=1856,device_name=MTT_S5000.json diff --git a/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=128,N=1856,device_name=MTT_S5000.json b/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=128,N=1856,device_name=MTT_S5000.json new file mode 100644 index 0000000..351b2c3 --- /dev/null +++ b/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=128,N=1856,device_name=MTT_S5000.json @@ -0,0 +1,98 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 1 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 1 + }, + "3": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 1 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 1 + }, + "5": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 1 + }, + "6": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 1 + }, + "7": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 1 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 1 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 1 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 1 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 1 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 1 + } +} From 758ea5aa7cfd7f75be78871ad420c6f843b7e725 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 21:05:29 +0800 Subject: [PATCH 17/22] fix: keep MoE routing graph safe and tune stable inputs --- .../triton/autotune/fused_moe/tune_moe.py | 18 ++++++-- .../triton/runtime/fused_moe/fused_moe.py | 44 ++++++++++++------- tests/test_moe_pipeline.py | 28 ++++++++++++ 3 files changed, 70 insertions(+), 20 deletions(-) diff --git a/src/torchada/triton/autotune/fused_moe/tune_moe.py b/src/torchada/triton/autotune/fused_moe/tune_moe.py index f414b99..9d677ca 100755 --- a/src/torchada/triton/autotune/fused_moe/tune_moe.py +++ b/src/torchada/triton/autotune/fused_moe/tune_moe.py @@ -331,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: @@ -363,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, ) @@ -406,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, @@ -433,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, @@ -458,6 +469,7 @@ def benchmark_config( topk_output = select_experts(x, input_gating, topk_config) 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) diff --git a/src/torchada/triton/runtime/fused_moe/fused_moe.py b/src/torchada/triton/runtime/fused_moe/fused_moe.py index 8547530..dcdbb29 100644 --- a/src/torchada/triton/runtime/fused_moe/fused_moe.py +++ b/src/torchada/triton/runtime/fused_moe/fused_moe.py @@ -11,8 +11,6 @@ try_get_optimal_moe_config, ) -_ALIGNMENT_CACHE: dict[int, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} - try: _support_tensor_descriptor = True except: @@ -23,6 +21,24 @@ def support_tensor_descriptor(): return _support_tensor_descriptor +def _is_current_stream_capturing() -> bool: + """Return whether the active accelerator stream is graph-capturing. + + The CUDA compatibility namespace can exist on CPU-only test hosts while + the underlying capture query is unavailable. Treat that case as an + ordinary eager call; on an actual capture stream the backend query is + expected to return a boolean. + """ + + query = getattr(torch.cuda, "is_current_stream_capturing", None) + if query is None: + return False + try: + return bool(query()) + except (AttributeError, RuntimeError): + return False + + @functools.lru_cache() def _down_moe_use_tma(): return support_tensor_descriptor() @@ -49,11 +65,6 @@ def moe_align_block_size( ensuring divisibility by block_size. """ - cache_key = id(topk_ids) - is_capturing = bool(getattr(torch.cuda, "is_current_stream_capturing", lambda: False)()) - if is_capturing and cache_key in _ALIGNMENT_CACHE: - return _ALIGNMENT_CACHE[cache_key] - if topk_ids.numel() < num_experts + 1: max_num_tokens_padded = topk_ids.numel() * block_size else: @@ -80,11 +91,9 @@ def moe_align_block_size( cumsum_buffer, True, ) - result = sorted_ids, expert_ids, num_tokens_post_pad - _ALIGNMENT_CACHE[cache_key] = result - return result + return sorted_ids, expert_ids, num_tokens_post_pad - except ImportError: + except (ImportError, AttributeError): pass # Try to import from vllm._custom_ops @@ -100,11 +109,14 @@ def moe_align_block_size( num_tokens_post_pad, None, ) - result = sorted_ids, expert_ids, num_tokens_post_pad - _ALIGNMENT_CACHE[cache_key] = result - return result + return sorted_ids, expert_ids, num_tokens_post_pad except (ImportError, AttributeError): + if _is_current_stream_capturing(): + raise RuntimeError( + "MoE graph capture requires the native moe_align_block_size " + "operator; the Python fallback is not graph-safe" + ) from None # Standalone torchada tuning images may contain vLLM Python sources # without the optional _moe_C extension. Keep the reference path # usable for correctness/tuning; production vLLM uses the native op. @@ -131,9 +143,7 @@ def moe_align_block_size( sorted_ids[: sorted_routes.numel()].copy_(sorted_routes) expert_ids[: blocks.numel()].copy_(blocks) num_tokens_post_pad[0] = sorted_routes.numel() - result = sorted_ids, expert_ids, num_tokens_post_pad - _ALIGNMENT_CACHE[cache_key] = result - return result + return sorted_ids, expert_ids, num_tokens_post_pad def _prepare_fused_moe_run( diff --git a/tests/test_moe_pipeline.py b/tests/test_moe_pipeline.py index d4c6645..48f63e7 100644 --- a/tests/test_moe_pipeline.py +++ b/tests/test_moe_pipeline.py @@ -4,6 +4,7 @@ the Python pipeline and argument plumbing without requiring a CUDA/MUSA device. """ +import sys from types import SimpleNamespace import pytest @@ -12,6 +13,33 @@ from torchada.triton.runtime.fused_moe import fused_moe as moe +def _force_python_alignment(monkeypatch): + """Make alignment tests independent of optional native extensions.""" + monkeypatch.setitem(sys.modules, "sgl_kernel", None) + monkeypatch.setitem(sys.modules, "vllm._custom_ops", None) + + +def test_moe_alignment_recomputes_mutated_routes(monkeypatch): + _force_python_alignment(monkeypatch) + monkeypatch.setattr(moe.torch.cuda, "is_current_stream_capturing", lambda: False) + topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + + first = moe.moe_align_block_size(topk_ids, block_size=2, num_experts=2) + topk_ids.copy_(torch.tensor([[1, 1], [1, 1]], dtype=torch.long)) + second = moe.moe_align_block_size(topk_ids, block_size=2, num_experts=2) + + assert not torch.equal(first[0], second[0]) + assert second[2].item() == 4 + assert second[1][:2].tolist() == [1, 1] + + +def test_moe_alignment_rejects_python_fallback_during_capture(monkeypatch): + _force_python_alignment(monkeypatch) + monkeypatch.setattr(moe.torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError, match="native moe_align_block_size"): + moe.moe_align_block_size(torch.tensor([[0]], dtype=torch.long), 1, 1) + + def _fake_gemm(a, w, bias, out, *args, **kwargs): """Reference implementation of invoke_fused_moe_kernel for CPU tests.""" ids, sorted_ids, expert_ids = args[4], args[5], args[6] From cef6de3f4b3342ef9296917dbec5f98aa362c9a1 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 21:23:30 +0800 Subject: [PATCH 18/22] fix: make MoE tuner graph and reference gates meaningful --- .../E=128,N=1856,device_name=MTT_S5000.json | 98 ------------------- .../triton/autotune/fused_moe/tune_moe.py | 85 ++++++++++++++-- 2 files changed, 76 insertions(+), 107 deletions(-) delete mode 100644 src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=128,N=1856,device_name=MTT_S5000.json diff --git a/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=128,N=1856,device_name=MTT_S5000.json b/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=128,N=1856,device_name=MTT_S5000.json deleted file mode 100644 index 351b2c3..0000000 --- a/src/torchada/triton/autotune/fused_moe/configs/triton_3_2_0/E=128,N=1856,device_name=MTT_S5000.json +++ /dev/null @@ -1,98 +0,0 @@ -{ - "1": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1 - }, - "2": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1 - }, - "3": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1 - }, - "4": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1 - }, - "5": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1 - }, - "6": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1 - }, - "7": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1 - }, - "8": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1 - }, - "16": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1 - }, - "32": { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1 - }, - "64": { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1 - }, - "128": { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1 - } -} diff --git a/src/torchada/triton/autotune/fused_moe/tune_moe.py b/src/torchada/triton/autotune/fused_moe/tune_moe.py index 9d677ca..610617f 100755 --- a/src/torchada/triton/autotune/fused_moe/tune_moe.py +++ b/src/torchada/triton/autotune/fused_moe/tune_moe.py @@ -468,6 +468,23 @@ 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) @@ -483,7 +500,10 @@ 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, ) @@ -507,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 @@ -535,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() From 8152a5800c5627d36a18196b841ba416059b0366 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Sun, 13 Sep 2026 11:04:59 +0800 Subject: [PATCH 19/22] feat: tune Nemotron non-gated MoE shapes --- .../triton/autotune/fused_moe/tune_moe.py | 122 ++++-------------- .../triton/autotune/fused_moe/utils.py | 26 ++-- .../triton/runtime/fused_moe/config.py | 8 +- tests/test_nemotron_moe_config.py | 76 ++++++++--- 4 files changed, 97 insertions(+), 135 deletions(-) diff --git a/src/torchada/triton/autotune/fused_moe/tune_moe.py b/src/torchada/triton/autotune/fused_moe/tune_moe.py index 610617f..d110b73 100755 --- a/src/torchada/triton/autotune/fused_moe/tune_moe.py +++ b/src/torchada/triton/autotune/fused_moe/tune_moe.py @@ -192,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, @@ -202,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, ) @@ -330,12 +332,11 @@ 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) - # 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: @@ -349,11 +350,7 @@ def benchmark_config( w2 = torch.randint( -127, 127, - ( - num_experts, - hidden_size, - shard_intermediate_size // 2 if is_gated else shard_intermediate_size, - ), + (num_experts, hidden_size, w2_input_size), dtype=torch.int8, device=device, ) @@ -368,11 +365,7 @@ def benchmark_config( w2 = torch.randint( 0, 255, - ( - num_experts, - hidden_size, - (shard_intermediate_size // 2 if is_gated else shard_intermediate_size) // 2, - ), + (num_experts, hidden_size, w2_input_size // 2), dtype=torch.uint8, device=device, ) @@ -387,7 +380,7 @@ def benchmark_config( w2 = torch.randn( num_experts, hidden_size, - shard_intermediate_size // 2 if is_gated else shard_intermediate_size, + w2_input_size, dtype=init_dtype, device=device, ) @@ -404,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, ) @@ -415,8 +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 - w2_k = shard_intermediate_size // 2 if is_gated else shard_intermediate_size - k_tiles_w2 = (w2_k + 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, @@ -443,8 +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 - w2_k = shard_intermediate_size // 2 if is_gated else shard_intermediate_size - k_tiles_w2 = (w2_k + 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, @@ -468,25 +459,7 @@ 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) @@ -500,10 +473,7 @@ def run(): intermediate_size_per_partition=shard_intermediate_size, top_k=topk, num_fused_shared_experts=num_fused_shared_experts, - # 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, + inplace=True, activation=activation, is_gated=is_gated, ) @@ -527,65 +497,22 @@ def run(): ) # Warmup & JIT - output = run() + 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): - graph_output = run() - torch.cuda.synchronize() - # 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() + for _ in range(10): + run() 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() + for _ in range(5): + 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: - output = run() + for _ in range(5): + run() torch.cuda.synchronize() # Flush L2 cache @@ -598,15 +525,11 @@ 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: - output = run() + run() end_events[i].record() torch.cuda.synchronize() @@ -1084,6 +1007,7 @@ def main(args: argparse.Namespace) -> None: "bf16", "bfloat16", "fp16", + "float16", "fp8_w8a8", "int8_w8a16", "int8_w8a8", diff --git a/src/torchada/triton/autotune/fused_moe/utils.py b/src/torchada/triton/autotune/fused_moe/utils.py index e49003b..16a9ab4 100644 --- a/src/torchada/triton/autotune/fused_moe/utils.py +++ b/src/torchada/triton/autotune/fused_moe/utils.py @@ -24,32 +24,36 @@ class BenchmarkConfig(TypedDict): def calculate_shard_intermediate_size( - intermediate_size: int, tp_size: int, ep_size: int = 1, is_gated: bool = True + 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 - # 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. + # 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]: - """Return the activation spelling and projection layout for a HF config. + """Infer the activation name and projection layout from 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. + ``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() - architecture = str((getattr(config, "architectures", None) or [""])[0]) + 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") diff --git a/src/torchada/triton/runtime/fused_moe/config.py b/src/torchada/triton/runtime/fused_moe/config.py index 69b97e5..19da803 100644 --- a/src/torchada/triton/runtime/fused_moe/config.py +++ b/src/torchada/triton/runtime/fused_moe/config.py @@ -71,13 +71,7 @@ def get_moe_configs( be picked and the associated configuration chosen to invoke the kernel. """ # Supported Triton versions, should be sorted from the newest to the oldest - supported_triton_versions = [ - "3.6.0", - "3.4.0", - "3.3.1", - "3.2.0", - "3.1.0", - ] + supported_triton_versions = ["3.6.0", "3.4.0", "3.3.1", "3.2.0", "3.1.0"] # First look up if an optimized configuration is available in the configs # directory diff --git a/tests/test_nemotron_moe_config.py b/tests/test_nemotron_moe_config.py index f93cb96..323fa16 100644 --- a/tests/test_nemotron_moe_config.py +++ b/tests/test_nemotron_moe_config.py @@ -1,10 +1,12 @@ from types import SimpleNamespace +import pytest import torch from torchada.triton.autotune.fused_moe.utils import ( calculate_shard_intermediate_size, get_config_filename, + get_model_config, infer_moe_activation, ) @@ -14,21 +16,59 @@ def test_nemotron_relu2_uses_non_gated_layout(): architectures=["NemotronHForCausalLM"], mlp_hidden_act="relu2", ) + assert infer_moe_activation(config) == ("relu2_no_mul", False) assert calculate_shard_intermediate_size(1856, tp_size=1, is_gated=False) == 1856 -def test_other_models_keep_gated_layout(): +def test_other_models_keep_historical_gated_layout(): config = SimpleNamespace(architectures=["Qwen3MoeForCausalLM"], hidden_act="silu") + assert infer_moe_activation(config) == ("silu", True) assert calculate_shard_intermediate_size(512, tp_size=2) == 512 -def test_nemotron_config_filename_uses_w2_width(): +def test_nemotron_model_config_exposes_projection_metadata(monkeypatch): + config = SimpleNamespace( + architectures=["NemotronHForCausalLM"], + mlp_hidden_act="relu2", + hidden_size=6144, + moe_latent_size=2688, + n_routed_experts=128, + num_experts_per_tok=6, + moe_intermediate_size=1856, + torch_dtype=torch.bfloat16, + ) + monkeypatch.setattr( + "torchada.triton.autotune.fused_moe.utils._load_model_config", + lambda _: config, + ) + + params = get_model_config("nemotron", tp_size=1) + + assert params["num_experts"] == 128 + assert params["topk"] == 6 + assert params["hidden_size"] == 2688 + assert params["shard_intermediate_size"] == 1856 + assert params["activation"] == "relu2_no_mul" + assert params["is_gated"] is False + + +def test_config_filename_uses_second_gemm_width(): gated = get_config_filename( - 128, 3712, 2688, 6, torch.bfloat16, False, False, False, False, False, None + 128, + 3712, + 2688, + 6, + torch.bfloat16, + False, + False, + False, + False, + False, + None, ) - nongated = get_config_filename( + non_gated = get_config_filename( 128, 1856, 2688, @@ -42,21 +82,21 @@ def test_nemotron_config_filename_uses_w2_width(): None, is_gated=False, ) + assert "E=128,N=1856" in gated - assert "E=128,N=1856" in nongated + assert "E=128,N=1856" in non_gated -def test_nemotron_benchmark_uses_single_projection_width(monkeypatch): - """The BF16 Nemotron benchmark must allocate w2 with N, rather than N/2.""" +@pytest.mark.parametrize("is_gated,expected_w2", [(True, 1856), (False, 3712)]) +def test_benchmark_allocates_matching_second_projection(monkeypatch, is_gated, expected_w2): + """w2's K dimension follows the actual post-activation width.""" from torchada.triton.autotune.fused_moe import tune_moe seen = [] def fake_randn(*shape, **kwargs): seen.append(tuple(shape)) - # Stop before the benchmark starts routing/timing. This keeps the - # shape assertion CPU-only and avoids allocating the full checkpoint - # dimensions. + # Stop after x, w1 and w2 allocation; this test only checks shapes. if len(seen) == 3: raise RuntimeError("stop after expert weights") return torch.empty((1,), dtype=kwargs.get("dtype", torch.float32)) @@ -71,12 +111,13 @@ def fake_randn(*shape, **kwargs): "num_warps": 1, "num_stages": 1, } - try: + + with pytest.raises(RuntimeError, match="stop after expert weights"): tune_moe.benchmark_config( config, 1, 128, - 1856, + 3712 if is_gated else 1856, 2688, 6, torch.bfloat16, @@ -85,10 +126,9 @@ def fake_randn(*shape, **kwargs): False, False, False, - activation="relu2_no_mul", - is_gated=False, + activation="silu" if is_gated else "relu2_no_mul", + is_gated=is_gated, ) - except RuntimeError as exc: - assert str(exc) == "stop after expert weights" - assert seen[1] == (128, 1856, 2688) - assert seen[2] == (128, 2688, 1856) + + assert seen[1] == (128, 3712 if is_gated else 1856, 2688) + assert seen[2] == (128, 2688, expected_w2) From 9c785826917b9d9c13d3df99e07d72fbe5d2b77a Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Sun, 13 Sep 2026 11:08:09 +0800 Subject: [PATCH 20/22] test: correct non-gated Nemotron projection width --- tests/test_nemotron_moe_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_nemotron_moe_config.py b/tests/test_nemotron_moe_config.py index 323fa16..ae4da4d 100644 --- a/tests/test_nemotron_moe_config.py +++ b/tests/test_nemotron_moe_config.py @@ -87,7 +87,7 @@ def test_config_filename_uses_second_gemm_width(): assert "E=128,N=1856" in non_gated -@pytest.mark.parametrize("is_gated,expected_w2", [(True, 1856), (False, 3712)]) +@pytest.mark.parametrize("is_gated,expected_w2", [(True, 1856), (False, 1856)]) def test_benchmark_allocates_matching_second_projection(monkeypatch, is_gated, expected_w2): """w2's K dimension follows the actual post-activation width.""" from torchada.triton.autotune.fused_moe import tune_moe From 143fb4a2bfe34b796cc8827bb2d822c0a10e595d Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Sun, 13 Sep 2026 14:52:32 +0800 Subject: [PATCH 21/22] fix: keep runtime resolver version scope unchanged --- src/torchada/triton/runtime/fused_moe/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/torchada/triton/runtime/fused_moe/config.py b/src/torchada/triton/runtime/fused_moe/config.py index 19da803..c34c3d6 100644 --- a/src/torchada/triton/runtime/fused_moe/config.py +++ b/src/torchada/triton/runtime/fused_moe/config.py @@ -71,7 +71,7 @@ def get_moe_configs( be picked and the associated configuration chosen to invoke the kernel. """ # Supported Triton versions, should be sorted from the newest to the oldest - supported_triton_versions = ["3.6.0", "3.4.0", "3.3.1", "3.2.0", "3.1.0"] + supported_triton_versions = ["3.4.0", "3.3.1", "3.2.0", "3.1.0"] # First look up if an optimized configuration is available in the configs # directory From c37ddfbcd36388fd8e146e819877840f52d1dad1 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Sun, 13 Sep 2026 16:13:20 +0800 Subject: [PATCH 22/22] perf: add Nemotron Triton 3.6 S5000 configs --- .../E=128,N=1856,device_name=MTT_S5000.json | 18 ++++++++++++++++++ ...ice_name=MTT_S5000,cache_dtype=float32.json | 11 +++++++++++ 2 files changed, 29 insertions(+) create mode 100644 src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/E=128,N=1856,device_name=MTT_S5000.json create mode 100644 src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/headdim=64,dstate=128,device_name=MTT_S5000,cache_dtype=float32.json diff --git a/src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/E=128,N=1856,device_name=MTT_S5000.json b/src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/E=128,N=1856,device_name=MTT_S5000.json new file mode 100644 index 0000000..d6edacc --- /dev/null +++ b/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 + } +} diff --git a/src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/headdim=64,dstate=128,device_name=MTT_S5000,cache_dtype=float32.json b/src/torchada/triton/autotune/fused_moe/configs/triton_3_6_0/headdim=64,dstate=128,device_name=MTT_S5000,cache_dtype=float32.json new file mode 100644 index 0000000..892a267 --- /dev/null +++ b/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 + } +} \ No newline at end of file