From b2fee5821a3ffbc676972b51986674a2cddc809d Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 16:48:20 +0800 Subject: [PATCH 01/17] 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 75f1a886be98f4c6a6e8ef6980aba0d31ccb4774 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:00:51 +0800 Subject: [PATCH 02/17] 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 ad6b5e296825be5170cad531d2ba80b7f1134380 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:09:10 +0800 Subject: [PATCH 03/17] 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 d5b7c8387834a8a09f06e7bfb6792b4824fc4ad6 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:15:14 +0800 Subject: [PATCH 04/17] 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 b0e9c0a181ca2c88f16e95bc5da7bff0bccbf032 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:26:51 +0800 Subject: [PATCH 05/17] 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 4d062cd3c913905b15d19941bd309f8dc8ce526b Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:29:40 +0800 Subject: [PATCH 06/17] 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 68c7221291c1aa942b55455ef8ba1f70300ac7a4 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:37:49 +0800 Subject: [PATCH 07/17] 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 ff61745170af7b10f8296680364e0524e42556f1 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:38:24 +0800 Subject: [PATCH 08/17] 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 bf63ff1df907d12290ec5603ac0f79b8cb098406 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:40:19 +0800 Subject: [PATCH 09/17] 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 275f7849ffb209cdc939d24257a116571cbc76f3 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:41:18 +0800 Subject: [PATCH 10/17] 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 602d7f2286260c4d554cb1a86ff510e68464eb88 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:42:46 +0800 Subject: [PATCH 11/17] 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 deb82948a2c0c59d66a3312229fade76956c509d Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 18:52:47 +0800 Subject: [PATCH 12/17] 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 f8226975073dbd2f4ca5b70f8183a2a90201719a Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 19:05:41 +0800 Subject: [PATCH 13/17] 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 d7fe76259f6a54d43349e436101177d18914f3f8 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 19:09:27 +0800 Subject: [PATCH 14/17] 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 7b31ef542f757180fe96bbe7a5ad9af025c8f531 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 19:40:48 +0800 Subject: [PATCH 15/17] 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 e7a33a4d6477cf1991a96d8545474cf7fff662c9 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 21:05:29 +0800 Subject: [PATCH 16/17] 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 7be2178bdf1d78760133b4793aacea679ab7fe9c Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Wed, 9 Sep 2026 21:23:30 +0800 Subject: [PATCH 17/17] 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()