diff --git a/src/torchada/triton/autotune/fused_moe/tune_moe.py b/src/torchada/triton/autotune/fused_moe/tune_moe.py index 702c429..610617f 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,6 +320,8 @@ def benchmark_config( per_channel_quant: bool, block_shape: List[int] = None, num_fused_shared_experts: int = 0, + activation: str = "silu", + is_gated: bool = True, num_iters: int = 100, ) -> float: """Run the fused MoE kernel and return latency in microseconds.""" @@ -325,6 +331,11 @@ def benchmark_config( num_routed_experts = num_experts - num_fused_shared_experts assert num_routed_experts > 0 x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + # The production runner commonly uses ``inplace=True``. Keep that mode + # for representative timings, but restore the same input before every + # timed iteration so one benchmark run cannot feed its output back as the + # next iteration's input. + x_initial = x.clone() # Create random weights based on quantization type if use_int8_w8a16 or use_int8_w8a8: @@ -338,7 +349,11 @@ def benchmark_config( w2 = torch.randint( -127, 127, - (num_experts, hidden_size, shard_intermediate_size // 2), + ( + num_experts, + hidden_size, + shard_intermediate_size // 2 if is_gated else shard_intermediate_size, + ), dtype=torch.int8, device=device, ) @@ -353,7 +368,11 @@ def benchmark_config( w2 = torch.randint( 0, 255, - (num_experts, hidden_size, shard_intermediate_size // 4), + ( + num_experts, + hidden_size, + (shard_intermediate_size // 2 if is_gated else shard_intermediate_size) // 2, + ), dtype=torch.uint8, device=device, ) @@ -368,7 +387,7 @@ def benchmark_config( w2 = torch.randn( num_experts, hidden_size, - shard_intermediate_size // 2, + shard_intermediate_size // 2 if is_gated else shard_intermediate_size, dtype=init_dtype, device=device, ) @@ -396,7 +415,8 @@ def benchmark_config( n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n n_tiles_w2 = (hidden_size + block_n - 1) // block_n k_tiles_w1 = (hidden_size + block_k - 1) // block_k - k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k + w2_k = shard_intermediate_size // 2 if is_gated else shard_intermediate_size + k_tiles_w2 = (w2_k + block_k - 1) // block_k w1_scale = torch.randn( (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.bfloat16, @@ -423,7 +443,8 @@ def benchmark_config( n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n n_tiles_w2 = (hidden_size + block_n - 1) // block_n k_tiles_w1 = (hidden_size + block_k - 1) // block_k - k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k + w2_k = shard_intermediate_size // 2 if is_gated else shard_intermediate_size + k_tiles_w2 = (w2_k + block_k - 1) // block_k w1_scale = torch.rand( (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.float32, @@ -447,7 +468,25 @@ def benchmark_config( ) topk_output = select_experts(x, input_gating, topk_config) + def reference_first_token() -> torch.Tensor: + """Compute one BF16 non-gated token with ordinary torch matmuls. + + This is intentionally a small semantic gate for the Nemotron path; + timing still uses the Triton kernel for the complete batch. + """ + + token = x_initial[0].float() + result = torch.zeros(hidden_size, dtype=torch.float32, device=device) + for choice in range(topk): + expert = topk_output.topk_ids[0, choice] + gate_up = torch.matmul(token, w1[expert].float().transpose(0, 1)) + activated = torch.relu(gate_up).square() + down = torch.matmul(activated, w2[expert].float().transpose(0, 1)) + result.add_(down, alpha=float(topk_output.topk_weights[0, choice])) + return result + def prepare(i: int): + x.copy_(x_initial) new_topk_output = select_experts(x, gating_output[i], topk_config) topk_output.topk_weights.copy_(new_topk_output.topk_weights) topk_output.topk_ids.copy_(new_topk_output.topk_ids) @@ -461,7 +500,12 @@ def run(): intermediate_size_per_partition=shard_intermediate_size, top_k=topk, num_fused_shared_experts=num_fused_shared_experts, - inplace=True, + # Keep the input immutable. In-place output turns repeated + # warmups and graph capture into a feedback loop where each run + # consumes the previous run's result. + inplace=False, + activation=activation, + is_gated=is_gated, ) with override_config(config): fused_moe( @@ -483,22 +527,65 @@ def run(): ) # Warmup & JIT - run() + output = run() torch.cuda.synchronize() + if dtype == torch.bfloat16 and not is_gated: + if not torch.isfinite(output).all(): + raise RuntimeError("BF16 non-gated MoE produced non-finite output") + reference = reference_first_token() + ref_scale = reference.abs().amax() + if not torch.isfinite(reference).all() or ref_scale <= 1e-12: + raise RuntimeError("BF16 non-gated reference output is degenerate") + out_token = output[0].float() + torch.testing.assert_close(out_token, reference, rtol=5e-2, atol=1e-6) + if out_token.abs().amax() <= ref_scale * 1e-2: + raise RuntimeError("BF16 non-gated MoE output is unexpectedly near zero") + use_graph = _env_flag("TORCHADA_TUNE_USE_GRAPH") and hasattr(torch.cuda, "CUDAGraph") if use_graph: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - for _ in range(10): - run() + graph_output = run() torch.cuda.synchronize() - for _ in range(5): - graph.replay() + # Verify that replay observes a changed routing tensor. A cached + # alignment result can silently make every replay execute the first + # route; this gate catches that before timing is reported. + prepare(0) + graph.replay() torch.cuda.synchronize() + first_ids = topk_output.topk_ids.clone() + first_output = graph_output.detach().clone() + route_index = next( + ( + i + for i in range(1, num_iters) + if not torch.equal( + first_ids, + select_experts(x_initial, gating_output[i], topk_config).topk_ids, + ) + ), + None, + ) + if route_index is None: + raise RuntimeError("graph route replay gate needs two distinct routing inputs") + prepare(route_index) + graph.replay() + torch.cuda.synchronize() + if torch.equal(first_ids, topk_output.topk_ids): + raise RuntimeError("graph route replay gate did not mutate top-k expert ids") + second_output = graph_output.detach().clone() + output_scale = max( + float(first_output.float().abs().amax()), + float(second_output.float().abs().amax()), + 1e-12, + ) + if not torch.isfinite(second_output).all(): + raise RuntimeError("graph replay produced non-finite MoE output") + if (first_output.float() - second_output.float()).abs().amax() <= output_scale * 1e-4: + raise RuntimeError("graph route replay output did not change with routing inputs") else: - for _ in range(5): - run() + output = run() torch.cuda.synchronize() # Flush L2 cache @@ -511,11 +598,15 @@ def run(): for i in range(num_iters): prepare(i) + # Flush immediately before each timed sample so cache state from the + # previous sample cannot bias the comparison between configurations. + cache_flush.zero_() + torch.cuda.synchronize() start_events[i].record() if use_graph: graph.replay() else: - run() + output = run() end_events[i].record() torch.cuda.synchronize() @@ -548,6 +639,8 @@ def build_model_entries(args: argparse.Namespace) -> List[ModelEntry]: entry.shard_intermediate_size = params["shard_intermediate_size"] entry.topk = params["topk"] entry.num_fused_shared_experts = params.get("num_fused_shared_experts", 0) + entry.activation = params.get("activation", "silu") + entry.is_gated = params.get("is_gated", True) entry.dtype_str = _resolve_dtype_str(entry.dtype_str, params) entry.dtype = _resolve_torch_dtype(entry.dtype_str, params) entry.block_shape = tuple(params["block_shape"]) if params["block_shape"] else None @@ -603,6 +696,8 @@ def build_model_entries(args: argparse.Namespace) -> List[ModelEntry]: entry.shard_intermediate_size = params["shard_intermediate_size"] entry.topk = params["topk"] entry.num_fused_shared_experts = params.get("num_fused_shared_experts", 0) + entry.activation = params.get("activation", "silu") + entry.is_gated = params.get("is_gated", True) entry.dtype_str = _resolve_dtype_str(entry.dtype_str, params) entry.dtype = _resolve_torch_dtype(entry.dtype_str, params) entry.block_shape = ( @@ -657,6 +752,8 @@ def _tune_worker( entry.per_channel_quant, list(entry.block_shape) if entry.block_shape else None, entry.num_fused_shared_experts, + activation=entry.activation, + is_gated=entry.is_gated, num_iters=10, ) except (triton.runtime.autotuner.OutOfResources, RuntimeError, AssertionError): @@ -767,6 +864,7 @@ def run_tuning(entries: List[ModelEntry], batch_sizes: List[int], args: argparse entry.use_int4, entry.per_channel_quant, entry.block_shape, + is_gated=entry.is_gated, ) sorted_batches = sorted(bs_to_config.keys()) best_configs = {bs: sort_config(bs_to_config[bs]) for bs in sorted_batches} @@ -808,7 +906,11 @@ def _benchmark_worker( ) block_n = entry.block_shape[0] if entry.block_shape else 0 block_k = entry.block_shape[1] if entry.block_shape else 0 - N = entry.shard_intermediate_size // 2 + N = ( + entry.shard_intermediate_size // 2 + if entry.is_gated + else entry.shard_intermediate_size + ) if entry.use_int4: N = N // 2 op_config = get_moe_configs( @@ -852,6 +954,8 @@ def _benchmark_worker( entry.per_channel_quant, list(entry.block_shape) if entry.block_shape else None, entry.num_fused_shared_experts, + activation=entry.activation, + is_gated=entry.is_gated, ) result_queue.put((entry, batch_size, kernel_time, None)) except Exception as e: @@ -975,7 +1079,16 @@ def main(args: argparse.Namespace) -> None: parser.add_argument( "--dtype", type=str, - choices=["auto", "fp8_w8a8", "int8_w8a16", "int8_w8a8", "int4_w4a16"], + choices=[ + "auto", + "bf16", + "bfloat16", + "fp16", + "fp8_w8a8", + "int8_w8a16", + "int8_w8a8", + "int4_w4a16", + ], default="auto", help="Quantization dtype.", ) 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/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 diff --git a/src/torchada/triton/runtime/fused_moe/fused_moe.py b/src/torchada/triton/runtime/fused_moe/fused_moe.py index 44728b2..dcdbb29 100644 --- a/src/torchada/triton/runtime/fused_moe/fused_moe.py +++ b/src/torchada/triton/runtime/fused_moe/fused_moe.py @@ -21,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() @@ -75,7 +93,7 @@ def moe_align_block_size( ) return sorted_ids, expert_ids, num_tokens_post_pad - except ImportError: + except (ImportError, AttributeError): pass # Try to import from vllm._custom_ops @@ -93,10 +111,39 @@ 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): + 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. + 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( @@ -206,6 +253,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 = ( @@ -272,11 +323,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]), @@ -300,7 +360,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, @@ -327,6 +391,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, @@ -358,6 +435,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 @@ -497,32 +578,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=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, + 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=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, + filter_expert=filter_expert, ) diff --git a/tests/test_moe_pipeline.py b/tests/test_moe_pipeline.py new file mode 100644 index 0000000..48f63e7 --- /dev/null +++ b/tests/test_moe_pipeline.py @@ -0,0 +1,216 @@ +"""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. +""" + +import sys +from types import SimpleNamespace + +import pytest +import torch + +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] + # 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 + # [tokens * topk, intermediate]. sorted_ids is identity in this test. + out_rows = out.reshape(-1, out.shape[-1]) + nrows = out_rows.shape[0] + + def write(row, value): + out_rows[row].copy_(value) + + if call_index == 0: + 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: + for row in range(nrows): + 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): + 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) + / 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.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, + 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, + # 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, + 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, 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) + + +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 new file mode 100644 index 0000000..f93cb96 --- /dev/null +++ b/tests/test_nemotron_moe_config.py @@ -0,0 +1,94 @@ +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 + + +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)