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 diff --git a/src/torchada/triton/autotune/fused_moe/tune_moe.py b/src/torchada/triton/autotune/fused_moe/tune_moe.py index 702c429..d110b73 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, @@ -188,7 +192,7 @@ def validate_and_log_entries(entries: List[ModelEntry]) -> None: for i, e in enumerate(entries): logger.info( "[%d] model=%s tp=%d ep=%d experts=%d hidden=%d " - "intermediate=%d topk=%d shared=%d dtype=%s block=%s", + "intermediate=%d topk=%d shared=%d activation=%s gated=%s dtype=%s block=%s", i, e.path, e.tp_size, @@ -198,6 +202,8 @@ def validate_and_log_entries(entries: List[ModelEntry]) -> None: e.shard_intermediate_size, e.topk, e.num_fused_shared_experts, + e.activation, + e.is_gated, e.dtype_str, e.block_shape, ) @@ -316,6 +322,8 @@ def benchmark_config( per_channel_quant: bool, block_shape: List[int] = None, num_fused_shared_experts: int = 0, + activation: str = "silu", + is_gated: bool = True, num_iters: int = 100, ) -> float: """Run the fused MoE kernel and return latency in microseconds.""" @@ -324,6 +332,10 @@ def benchmark_config( init_dtype = torch.float16 if use_fp8_w8a8 else dtype num_routed_experts = num_experts - num_fused_shared_experts assert num_routed_experts > 0 + # ``shard_intermediate_size`` is the first GEMM output width (w1.shape[1]). + # Gated models use half of it as the second GEMM input width because w1 + # contains gate and up projections; non-gated models use the full width. + w2_input_size = shard_intermediate_size // 2 if is_gated else shard_intermediate_size x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) # Create random weights based on quantization type @@ -338,7 +350,7 @@ def benchmark_config( w2 = torch.randint( -127, 127, - (num_experts, hidden_size, shard_intermediate_size // 2), + (num_experts, hidden_size, w2_input_size), dtype=torch.int8, device=device, ) @@ -353,7 +365,7 @@ def benchmark_config( w2 = torch.randint( 0, 255, - (num_experts, hidden_size, shard_intermediate_size // 4), + (num_experts, hidden_size, w2_input_size // 2), dtype=torch.uint8, device=device, ) @@ -368,7 +380,7 @@ def benchmark_config( w2 = torch.randn( num_experts, hidden_size, - shard_intermediate_size // 2, + w2_input_size, dtype=init_dtype, device=device, ) @@ -385,7 +397,7 @@ def benchmark_config( w1_scale = w2_scale = a1_scale = a2_scale = None if use_int8_w8a16: w1_scale = torch.randn( - (num_experts, 2 * shard_intermediate_size), + (num_experts, shard_intermediate_size), dtype=torch.float32, device=device, ) @@ -396,7 +408,7 @@ def benchmark_config( n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n n_tiles_w2 = (hidden_size + block_n - 1) // block_n k_tiles_w1 = (hidden_size + block_k - 1) // block_k - k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k + k_tiles_w2 = (w2_input_size + block_k - 1) // block_k w1_scale = torch.randn( (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.bfloat16, @@ -423,7 +435,7 @@ def benchmark_config( n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n n_tiles_w2 = (hidden_size + block_n - 1) // block_n k_tiles_w1 = (hidden_size + block_k - 1) // block_k - k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k + k_tiles_w2 = (w2_input_size + block_k - 1) // block_k w1_scale = torch.rand( (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.float32, @@ -462,6 +474,8 @@ def run(): top_k=topk, num_fused_shared_experts=num_fused_shared_experts, inplace=True, + activation=activation, + is_gated=is_gated, ) with override_config(config): fused_moe( @@ -548,6 +562,8 @@ def build_model_entries(args: argparse.Namespace) -> List[ModelEntry]: entry.shard_intermediate_size = params["shard_intermediate_size"] entry.topk = params["topk"] entry.num_fused_shared_experts = params.get("num_fused_shared_experts", 0) + entry.activation = params.get("activation", "silu") + entry.is_gated = params.get("is_gated", True) entry.dtype_str = _resolve_dtype_str(entry.dtype_str, params) entry.dtype = _resolve_torch_dtype(entry.dtype_str, params) entry.block_shape = tuple(params["block_shape"]) if params["block_shape"] else None @@ -603,6 +619,8 @@ def build_model_entries(args: argparse.Namespace) -> List[ModelEntry]: entry.shard_intermediate_size = params["shard_intermediate_size"] entry.topk = params["topk"] entry.num_fused_shared_experts = params.get("num_fused_shared_experts", 0) + entry.activation = params.get("activation", "silu") + entry.is_gated = params.get("is_gated", True) entry.dtype_str = _resolve_dtype_str(entry.dtype_str, params) entry.dtype = _resolve_torch_dtype(entry.dtype_str, params) entry.block_shape = ( @@ -657,6 +675,8 @@ def _tune_worker( entry.per_channel_quant, list(entry.block_shape) if entry.block_shape else None, entry.num_fused_shared_experts, + activation=entry.activation, + is_gated=entry.is_gated, num_iters=10, ) except (triton.runtime.autotuner.OutOfResources, RuntimeError, AssertionError): @@ -767,6 +787,7 @@ def run_tuning(entries: List[ModelEntry], batch_sizes: List[int], args: argparse entry.use_int4, entry.per_channel_quant, entry.block_shape, + is_gated=entry.is_gated, ) sorted_batches = sorted(bs_to_config.keys()) best_configs = {bs: sort_config(bs_to_config[bs]) for bs in sorted_batches} @@ -808,7 +829,11 @@ def _benchmark_worker( ) block_n = entry.block_shape[0] if entry.block_shape else 0 block_k = entry.block_shape[1] if entry.block_shape else 0 - N = entry.shard_intermediate_size // 2 + N = ( + entry.shard_intermediate_size // 2 + if entry.is_gated + else entry.shard_intermediate_size + ) if entry.use_int4: N = N // 2 op_config = get_moe_configs( @@ -852,6 +877,8 @@ def _benchmark_worker( entry.per_channel_quant, list(entry.block_shape) if entry.block_shape else None, entry.num_fused_shared_experts, + activation=entry.activation, + is_gated=entry.is_gated, ) result_queue.put((entry, batch_size, kernel_time, None)) except Exception as e: @@ -975,7 +1002,17 @@ def main(args: argparse.Namespace) -> None: parser.add_argument( "--dtype", type=str, - choices=["auto", "fp8_w8a8", "int8_w8a16", "int8_w8a8", "int4_w4a16"], + choices=[ + "auto", + "bf16", + "bfloat16", + "fp16", + "float16", + "fp8_w8a8", + "int8_w8a16", + "int8_w8a8", + "int4_w4a16", + ], default="auto", help="Quantization dtype.", ) diff --git a/src/torchada/triton/autotune/fused_moe/utils.py b/src/torchada/triton/autotune/fused_moe/utils.py index 469b362..16a9ab4 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,39 @@ class BenchmarkConfig(TypedDict): def calculate_shard_intermediate_size( - intermediate_size: int, tp_size: int, ep_size: int = 1 + intermediate_size: int, + tp_size: int, + ep_size: int = 1, + is_gated: bool = True, ) -> int: assert tp_size % ep_size == 0 moe_tp_size = tp_size // ep_size assert intermediate_size % moe_tp_size == 0 - return 2 * intermediate_size // moe_tp_size + # Gated projections (for example SwiGLU) store gate and up next to each + # other in w1. Non-gated projections such as Nemotron-H's relu2 use one + # projection, so the w1 output width is the model's intermediate size. + multiplier = 2 if is_gated else 1 + return multiplier * intermediate_size // moe_tp_size + + +def infer_moe_activation(config) -> Tuple[str, bool]: + """Infer the activation name and projection layout from a HF config. + + ``NemotronHForCausalLM`` passes ``activation_without_mul`` to its fused + MoE layer. The HF config advertises ``relu2`` while the checkpoint has a + single projection (``relu2_no_mul``). Keep all other architectures on + the historical gated-SiLU default unless their activation already carries + the explicit ``_no_mul`` suffix. + """ + raw = getattr(config, "mlp_hidden_act", None) + if raw is None: + raw = getattr(config, "hidden_act", "silu") + activation = str(raw).replace("torch.", "").lower() + architectures = getattr(config, "architectures", None) or [] + architecture = str(architectures[0]) if architectures else type(config).__name__ + if architecture == "NemotronHForCausalLM" and not activation.endswith("_no_mul"): + activation = f"{activation}_no_mul" + return activation, not activation.endswith("_no_mul") def get_num_shared_experts(config, disable_shared_experts_fusion: bool) -> int: @@ -146,6 +173,7 @@ def get_model_config( config = _load_model_config(model_name) architecture = config.architectures[0] + activation, is_gated = infer_moe_activation(config) quant_dtype_str = infer_quant_dtype_str(config) block_shape = None if hasattr(config, "quantization_config") and "weight_block_size" in config.quantization_config: @@ -247,7 +275,9 @@ def get_model_config( topk = config.num_experts_per_tok intermediate_size = config.intermediate_size - shard_intermediate_size = calculate_shard_intermediate_size(intermediate_size, tp_size, ep_size) + shard_intermediate_size = calculate_shard_intermediate_size( + intermediate_size, tp_size, ep_size, is_gated=is_gated + ) return { "num_experts": E, @@ -259,6 +289,8 @@ def get_model_config( "architecture": architecture, "num_fused_shared_experts": num_fused_shared_experts, "quant_dtype_str": quant_dtype_str, + "activation": activation, + "is_gated": is_gated, } @@ -333,6 +365,7 @@ def get_config_filename( use_int4_w4a16: bool, per_channel_quant: bool, block_shape: List[int], + is_gated: bool = True, ) -> str: dtype_str = get_config_dtype_str( dtype, @@ -344,7 +377,7 @@ def get_config_filename( # NOTE(woosuk): The current naming convention uses w2.shape[2], which # is the intermediate size after silu_and_mul. - N = shard_intermediate_size // 2 + N = shard_intermediate_size // 2 if is_gated else shard_intermediate_size if use_int4_w4a16: N = N // 2 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_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"], 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..ae4da4d --- /dev/null +++ b/tests/test_nemotron_moe_config.py @@ -0,0 +1,134 @@ +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, +) + + +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_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_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, + ) + non_gated = 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 non_gated + + +@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 + + seen = [] + + def fake_randn(*shape, **kwargs): + seen.append(tuple(shape)) + # 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)) + + 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, + } + + with pytest.raises(RuntimeError, match="stop after expert weights"): + tune_moe.benchmark_config( + config, + 1, + 128, + 3712 if is_gated else 1856, + 2688, + 6, + torch.bfloat16, + False, + False, + False, + False, + False, + activation="silu" if is_gated else "relu2_no_mul", + is_gated=is_gated, + ) + + assert seen[1] == (128, 3712 if is_gated else 1856, 2688) + assert seen[2] == (128, 2688, expected_w2)