diff --git a/benchmarks/README.md b/benchmarks/README.md index 76cb8b8045a..ab051cf1ab7 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -1,5 +1,38 @@ # FlashInfer Perf Benchmarking Framework -- `flashinfer_benchmark.py` +## VibeCUDA AlphaMoE router versus CAKE + +`bench_alphamoe_router.py` compares the VibeCUDA backend directly with the +optimized CAKE router from PR 4339. Torch is only a correctness reference, +never the performance denominator. Because both revisions provide the +`flashinfer` package, install the pinned CAKE checkout into an isolated virtual +environment and pass both its checkout and Python executable to the benchmark: + +```bash +BASELINE_WT=/tmp/flashinfer-pr4339-baseline +BASELINE_VENV=/tmp/flashinfer-pr4339-venv +CANDIDATE_VENV=/tmp/flashinfer-vibecuda-venv +git fetch https://github.com/flashinfer-ai/flashinfer.git \ + 0725744e58a9e338e8d315d82891878b07decd8f +git worktree add --detach "$BASELINE_WT" \ + 0725744e58a9e338e8d315d82891878b07decd8f +python3 -m pip install virtualenv +python3 -m virtualenv --system-site-packages "$BASELINE_VENV" +"$BASELINE_VENV/bin/python" -m pip install --no-build-isolation -e "$BASELINE_WT" -v + +python3 -m virtualenv --system-site-packages "$CANDIDATE_VENV" +"$CANDIDATE_VENV/bin/python" -m pip install --no-build-isolation -e "$PWD" -v + +PYTHONPATH=$PWD "$CANDIDATE_VENV/bin/python" benchmarks/bench_alphamoe_router.py \ + --candidate-python "$CANDIDATE_VENV/bin/python" \ + --baseline-root "$BASELINE_WT" \ + --baseline-python "$BASELINE_VENV/bin/python" +``` + +The command validates the immutable CAKE commit, runs both implementations in +isolated processes with the same four workloads and CUPTI protocol, and reports +CAKE/VibeCUDA per-workload, arithmetic-mean, and geometric-mean speedup. + The aim of `flashinfer_benchmark.py` is to provide a single framework for benchmarking any FlashInfer kernel and replace standalone benchmarking scripts. `bench_recurrent_kda_prefill.py --case-set h12` runs the six Kimi-K3 TP8 H12 diff --git a/benchmarks/bench_alphamoe_router.py b/benchmarks/bench_alphamoe_router.py new file mode 100644 index 00000000000..796a15d63e3 --- /dev/null +++ b/benchmarks/bench_alphamoe_router.py @@ -0,0 +1,281 @@ +# Copyright (c) 2026 by FlashInfer team. +# Licensed under the Apache License, Version 2.0 (the "License"). + +"""Compare the VibeCUDA AlphaMoE router with pinned CAKE PR 4339. + +The candidate and CAKE checkouts both provide the ``flashinfer`` package, so +the benchmark executes them in isolated Python processes and combines their +matched CUPTI measurements. Prepare the baseline checkout as documented in +``benchmarks/README.md``, then run:: + + python3 benchmarks/bench_alphamoe_router.py \ + --candidate-python /tmp/flashinfer-vibecuda-venv/bin/python \ + --baseline-root /tmp/flashinfer-pr4339-baseline \ + --baseline-python /tmp/flashinfer-pr4339-venv/bin/python + +Torch is only an independent correctness reference, never the denominator. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import statistics +import subprocess +import sys +import tempfile +from dataclasses import asdict, dataclass +from pathlib import Path + +CAKE_PR = "https://github.com/flashinfer-ai/flashinfer/pull/4339" +CAKE_SHA = "0725744e58a9e338e8d315d82891878b07decd8f" +DRY_RUN_ITERS = 5 +REPEAT_ITERS = 10 + + +@dataclass(frozen=True) +class RouterConfig: + name: str + num_tokens: int + num_experts: int + top_k: int + block_m: int + has_shared_expert: bool + + +CONFIGS = ( + RouterConfig("single-1tok-e512-shared", 1, 512, 2, 16, True), + RouterConfig("decode-8tok-e257-shared", 8, 257, 9, 8, True), + RouterConfig("batch-32tok-e512", 32, 512, 8, 16, False), + RouterConfig("batch-128tok-e512", 128, 512, 8, 16, False), +) + + +def _checkout_sha(root: Path) -> str: + return subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=root, text=True + ).strip() + + +def _validate_baseline(root: Path) -> None: + actual = _checkout_sha(root) + if actual != CAKE_SHA: + raise RuntimeError( + f"CAKE baseline must be {CAKE_SHA}, got {actual} at {root}" + ) + + +def _clean_pythonpath(root: Path) -> str: + candidate_root = Path(__file__).resolve().parents[1] + entries = [str(root)] + for entry in os.environ.get("PYTHONPATH", "").split(os.pathsep): + if not entry: + continue + resolved = Path(entry).resolve() + if resolved != candidate_root: + entries.append(str(resolved)) + return os.pathsep.join(entries) + + +def _run_worker(*, backend: str, root: Path, python: Path, output: Path) -> None: + env = os.environ.copy() + env["PYTHONPATH"] = _clean_pythonpath(root) + subprocess.run( + [ + str(python), + str(Path(__file__).resolve()), + "--worker", + backend, + "--output", + str(output), + ], + cwd=root, + env=env, + check=True, + ) + + +def _worker(backend: str, output: Path) -> None: + import numpy as np + import torch + + from flashinfer.fused_moe import ( + allocate_alphamoe_route_plan, + alphamoe_fused_router, + ) + try: + from flashinfer.testing import bench_gpu_time + except ImportError: + from flashinfer.testing.utils import bench_gpu_time + + capability = torch.cuda.get_device_capability() + if capability not in {(10, 0), (10, 3)}: + raise RuntimeError(f"CC 10.0 or 10.3 required, got {capability}") + + rows: list[dict[str, object]] = [] + for case_index, config in enumerate(CONFIGS): + generator = torch.Generator(device="cuda").manual_seed(29001 + case_index) + logits = torch.randn( + config.num_tokens, + config.num_experts, + generator=generator, + device="cuda", + dtype=torch.float32, + ) + plan = allocate_alphamoe_route_plan( + logits, + top_k=config.top_k, + block_m=config.block_m, + has_shared_expert=config.has_shared_expert, + ) + + if backend == "cake": + + def run() -> None: + alphamoe_fused_router( + logits, + top_k=config.top_k, + block_m=config.block_m, + has_shared_expert=config.has_shared_expert, + plan=plan, + ) + + else: + + def run() -> None: + alphamoe_fused_router(logits, plan=plan, backend="vibecuda") + + run() + torch.cuda.synchronize() + samples = bench_gpu_time( + run, + enable_cupti=True, + dry_run_iters=DRY_RUN_ITERS, + repeat_iters=REPEAT_ITERS, + cold_l2_cache=True, + use_cuda_graph=False, + ) + rows.append( + { + "config": asdict(config), + "median_us": float(np.median(samples)) * 1e3, + "samples": len(samples), + } + ) + + output.write_text( + json.dumps( + { + "backend": backend, + "device": torch.cuda.get_device_name(), + "compute_capability": list(capability), + "timing": { + "method": "CUPTI GPU activity", + "cold_l2": True, + "cuda_graph": False, + "dry_run_iters": DRY_RUN_ITERS, + "repeat_iters": REPEAT_ITERS, + "aggregation": "per-workload median", + }, + "rows": rows, + }, + indent=2, + ) + + "\n" + ) + + +def _aggregate(candidate: dict, cake: dict) -> dict: + if candidate["device"] != cake["device"]: + raise RuntimeError("candidate and CAKE measurements used different devices") + if candidate["timing"] != cake["timing"]: + raise RuntimeError("candidate and CAKE timing protocols differ") + rows = [] + for candidate_row, cake_row in zip(candidate["rows"], cake["rows"], strict=True): + if candidate_row["config"] != cake_row["config"]: + raise RuntimeError("candidate and CAKE workload manifests differ") + speedup = cake_row["median_us"] / candidate_row["median_us"] + rows.append( + { + "config": candidate_row["config"], + "cake_us": cake_row["median_us"], + "vibecuda_us": candidate_row["median_us"], + "speedup": speedup, + } + ) + speedups = [row["speedup"] for row in rows] + return { + "baseline": {"name": "CAKE AlphaMoE router", "pr": CAKE_PR, "sha": CAKE_SHA}, + "device": candidate["device"], + "timing": candidate["timing"], + "rows": rows, + "arithmetic_mean_speedup": statistics.fmean(speedups), + "geometric_mean_speedup": math.exp(statistics.fmean(map(math.log, speedups))), + } + + +def _print_result(result: dict) -> None: + print(f"VibeCUDA AlphaMoE router vs CAKE PR 4339 ({CAKE_SHA[:12]})") + print( + "Protocol: CUPTI, cold L2, no CUDA Graph, " + f"dry_run={DRY_RUN_ITERS}, repeats={REPEAT_ITERS}, median" + ) + for row in result["rows"]: + name = row["config"]["name"] + print( + f"{name:28s} CAKE {row['cake_us']:8.2f} us " + f"VibeCUDA {row['vibecuda_us']:8.2f} us {row['speedup']:6.2f}x" + ) + print(f"arithmetic mean: {result['arithmetic_mean_speedup']:.4f}x") + print(f"geometric mean: {result['geometric_mean_speedup']:.4f}x") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--candidate-python", type=Path, default=Path(sys.executable)) + parser.add_argument("--baseline-root", type=Path) + parser.add_argument("--baseline-python", type=Path) + parser.add_argument("--json", type=Path) + parser.add_argument("--worker", choices=("cake", "vibecuda"), help=argparse.SUPPRESS) + parser.add_argument("--output", type=Path, help=argparse.SUPPRESS) + args = parser.parse_args() + + if args.worker: + if args.output is None: + parser.error("--worker requires --output") + _worker(args.worker, args.output) + return + if args.baseline_root is None or args.baseline_python is None: + parser.error("--baseline-root and --baseline-python are required") + + baseline_root = args.baseline_root.resolve() + _validate_baseline(baseline_root) + candidate_root = Path(__file__).resolve().parents[1] + with tempfile.TemporaryDirectory(prefix="alphamoe-router-bench-") as tmp: + tmp_path = Path(tmp) + candidate_json = tmp_path / "candidate.json" + cake_json = tmp_path / "cake.json" + _run_worker( + backend="vibecuda", + root=candidate_root, + python=args.candidate_python.resolve(), + output=candidate_json, + ) + _run_worker( + backend="cake", + root=baseline_root, + python=args.baseline_python.resolve(), + output=cake_json, + ) + result = _aggregate( + json.loads(candidate_json.read_text()), json.loads(cake_json.read_text()) + ) + _print_result(result) + if args.json: + args.json.write_text(json.dumps(result, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/csrc/alphamoe_router/alphamoe_router.cu b/csrc/alphamoe_router/alphamoe_router.cu new file mode 100644 index 00000000000..e28c806debb --- /dev/null +++ b/csrc/alphamoe_router/alphamoe_router.cu @@ -0,0 +1,126 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include + +#include "flashinfer/fused_moe/alphamoe_router.cuh" +#include "tvm_ffi_utils.h" + +namespace flashinfer::alphamoe_router { + +using tvm::ffi::TensorView; + +/*! + * \brief TVM-FFI entry point for the fused AlphaMoE gating router. + * + * router_logits : [num_tokens, num_experts] float32 (input) + * topk_weights : [num_tokens, top_k] float32 (output) + * topk_ids : [num_tokens, top_k] int32 (output) + * sorted_token_ids : [max_blocks * block_m] int32 (output) + * expert_ids : [max_blocks] int32 (output) + * num_tokens_post_padded : [1] int32 (output) + * expert_counts : [num_experts] int32 (output) + * expert_offsets : [num_experts + 1] int32 (output) + * expert_scatter_offsets : [num_experts] int32 (output) + * scratch : [see alphamoe_router_scratch_ints] int32 (workspace) + * + * (top_k, block_m, has_shared_expert) come in as scalars; every output shape + * is derived from them and the input shape, so the launch is fixed under CUDA + * graph capture and no host synchronization happens anywhere on this path. + */ +void AlphaMoeFusedRouter(TensorView router_logits, TensorView topk_weights, + TensorView topk_ids, TensorView sorted_token_ids, + TensorView expert_ids, + TensorView num_tokens_post_padded, + TensorView expert_counts, TensorView expert_offsets, + TensorView expert_scatter_offsets, TensorView scratch, + int64_t top_k, int64_t block_m, + bool has_shared_expert) { + CHECK_INPUT_AND_TYPE(router_logits, dl_float32); + CHECK_INPUT_AND_TYPE(topk_weights, dl_float32); + CHECK_INPUT_AND_TYPE(topk_ids, dl_int32); + CHECK_INPUT_AND_TYPE(sorted_token_ids, dl_int32); + CHECK_INPUT_AND_TYPE(expert_ids, dl_int32); + CHECK_INPUT_AND_TYPE(num_tokens_post_padded, dl_int32); + CHECK_INPUT_AND_TYPE(expert_counts, dl_int32); + CHECK_INPUT_AND_TYPE(expert_offsets, dl_int32); + CHECK_INPUT_AND_TYPE(expert_scatter_offsets, dl_int32); + CHECK_INPUT_AND_TYPE(scratch, dl_int32); + + CHECK_DEVICE(topk_weights, router_logits); + CHECK_DEVICE(topk_ids, router_logits); + CHECK_DEVICE(sorted_token_ids, router_logits); + CHECK_DEVICE(expert_ids, router_logits); + CHECK_DEVICE(num_tokens_post_padded, router_logits); + CHECK_DEVICE(expert_counts, router_logits); + CHECK_DEVICE(expert_offsets, router_logits); + CHECK_DEVICE(expert_scatter_offsets, router_logits); + CHECK_DEVICE(scratch, router_logits); + + CHECK_DIM(2, router_logits); + CHECK_DIM(2, topk_weights); + CHECK_DIM(2, topk_ids); + CHECK_DIM(1, sorted_token_ids); + CHECK_DIM(1, expert_ids); + CHECK_DIM(1, num_tokens_post_padded); + CHECK_DIM(1, expert_counts); + CHECK_DIM(1, expert_offsets); + CHECK_DIM(1, expert_scatter_offsets); + CHECK_DIM(1, scratch); + + const int64_t num_tokens = router_logits.sizes()[0]; + const int64_t num_experts = router_logits.sizes()[1]; + + const auto params = flashinfer::fused_moe::make_alphamoe_router_params( + static_cast(num_tokens), static_cast(num_experts), + static_cast(top_k), static_cast(block_m), + has_shared_expert ? 1 : 0); + + TVM_FFI_ICHECK(topk_weights.sizes()[0] == num_tokens && + topk_weights.sizes()[1] == top_k && + topk_ids.sizes()[0] == num_tokens && + topk_ids.sizes()[1] == top_k) + << "topk outputs must be [num_tokens, top_k]"; + TVM_FFI_ICHECK(expert_counts.numel() == num_experts && + expert_scatter_offsets.numel() == num_experts) + << "expert counts/scatter offsets must be [num_experts]"; + TVM_FFI_ICHECK(expert_offsets.numel() == num_experts + 1) + << "expert_offsets must be [num_experts + 1]"; + TVM_FFI_ICHECK(num_tokens_post_padded.numel() == 1) + << "num_tokens_post_padded must be [1]"; + TVM_FFI_ICHECK(sorted_token_ids.numel() == params.slots) + << "sorted_token_ids must be [max_blocks * block_m] = " << params.slots; + TVM_FFI_ICHECK(expert_ids.numel() == params.max_blocks) + << "expert_ids must be [max_blocks] = " << params.max_blocks; + TVM_FFI_ICHECK(scratch.numel() >= + flashinfer::fused_moe::alphamoe_router_scratch_ints(params)) + << "scratch too small for the generic path"; + + auto stream = get_stream(router_logits.device()); + + flashinfer::fused_moe::alphamoe_router_forward( + params, static_cast(router_logits.data_ptr()), + static_cast(topk_weights.data_ptr()), + static_cast(topk_ids.data_ptr()), + static_cast(expert_counts.data_ptr()), + static_cast(expert_offsets.data_ptr()), + static_cast(expert_scatter_offsets.data_ptr()), + static_cast(num_tokens_post_padded.data_ptr()), + static_cast(expert_ids.data_ptr()), + static_cast(sorted_token_ids.data_ptr()), + static_cast(scratch.data_ptr()), stream); +} + +} // namespace flashinfer::alphamoe_router diff --git a/csrc/alphamoe_router/alphamoe_router_jit_binding.cu b/csrc/alphamoe_router/alphamoe_router_jit_binding.cu new file mode 100644 index 00000000000..62f71b9ab77 --- /dev/null +++ b/csrc/alphamoe_router/alphamoe_router_jit_binding.cu @@ -0,0 +1,31 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "tvm_ffi_utils.h" + +namespace flashinfer::alphamoe_router { +using tvm::ffi::TensorView; +void AlphaMoeFusedRouter(TensorView router_logits, TensorView topk_weights, + TensorView topk_ids, TensorView sorted_token_ids, + TensorView expert_ids, + TensorView num_tokens_post_padded, + TensorView expert_counts, TensorView expert_offsets, + TensorView expert_scatter_offsets, TensorView scratch, + int64_t top_k, int64_t block_m, + bool has_shared_expert); +} // namespace flashinfer::alphamoe_router + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(alphamoe_fused_router, + flashinfer::alphamoe_router::AlphaMoeFusedRouter); diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 699adee5d54..60b75d8a77f 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -105,6 +105,13 @@ trtllm_fp8_per_tensor_scale_routed_moe, ) +# AlphaMoE fused gating router ("vibecuda" backend); plain CUDA, always available. +from .fused_moe import ( # noqa: F401 + AlphaMoeRoutePlan as AlphaMoeRoutePlan, + allocate_alphamoe_route_plan as allocate_alphamoe_route_plan, + alphamoe_fused_router as alphamoe_fused_router, +) + # CuteDSL high-level APIs (conditionally if cute_dsl available) with contextlib.suppress(ImportError): from .fused_moe import ( diff --git a/flashinfer/aot.py b/flashinfer/aot.py index 695dace4176..bb27af9b02d 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -152,6 +152,7 @@ from .jit.spdlog import gen_spdlog_module from .jit.moe_utils import gen_moe_utils_module from .jit.hash_topk import gen_hash_topk_module +from .jit.alphamoe_router import gen_alphamoe_router_module from .jit.tllm_utils import gen_trtllm_utils_module from .jit.topk import gen_topk_module from .jit.xqa import gen_xqa_module, gen_xqa_module_mla @@ -673,6 +674,8 @@ def gen_all_modules( jit_specs.append(gen_bgmv_moe_module()) # DSv4 hash-based MoE routing (SM-portable) jit_specs.append(gen_hash_topk_module()) + # AlphaMoE fused gating router (SM-portable) + jit_specs.append(gen_alphamoe_router_module()) if has_sm90: jit_specs.append(gen_gemm_sm90_module()) # fp8 blockscale GEMM (SM90) diff --git a/flashinfer/fused_moe/__init__.py b/flashinfer/fused_moe/__init__.py index c1e8d4dfac0..d3c59fe8514 100644 --- a/flashinfer/fused_moe/__init__.py +++ b/flashinfer/fused_moe/__init__.py @@ -137,6 +137,12 @@ trtllm_gen_routing as trtllm_gen_routing, ) +from .alphamoe_router import ( # noqa: F401 + AlphaMoeRoutePlan as AlphaMoeRoutePlan, + allocate_alphamoe_route_plan as allocate_alphamoe_route_plan, + alphamoe_fused_router as alphamoe_fused_router, +) + from .bgmv_moe import ( # noqa: F401 bgmv_moe as bgmv_moe, bgmv_moe_shrink as bgmv_moe_shrink, @@ -269,6 +275,9 @@ "hash_topk", "TrtllmGenRoutingResult", "trtllm_gen_routing", + "AlphaMoeRoutePlan", + "allocate_alphamoe_route_plan", + "alphamoe_fused_router", "bgmv_moe", "bgmv_moe_shrink", "bgmv_moe_expand", diff --git a/flashinfer/fused_moe/alphamoe_router.py b/flashinfer/fused_moe/alphamoe_router.py new file mode 100644 index 00000000000..b1fc8befd8e --- /dev/null +++ b/flashinfer/fused_moe/alphamoe_router.py @@ -0,0 +1,440 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import functools +from types import SimpleNamespace +from typing import Optional, Tuple + +import torch + +from flashinfer.api_logging import flashinfer_api +from flashinfer.jit import gen_alphamoe_router_module +from flashinfer.trace.templates.moe import alphamoe_fused_router_trace +from flashinfer.utils import ( + backend_requirement, + register_custom_op, + supported_compute_capability, +) + +# alphamoe_fused_router is plain CUDA (warp reductions, shared-memory scans, +# Programmatic Dependent Launch on SM90+ with a regular-launch fallback on +# older architectures), portable across all tensor-core capable GPUs. +_ALPHAMOE_ROUTER_SUPPORTED_CC = [80, 86, 89, 90, 100, 103, 107, 110, 120, 121] + +# Canonical output order of the routing bundle (matches ``AlphaMoeRoutePlan`` +# iteration order and the ``alphamoe_fused_router`` return tuple). +_PLAN_FIELDS = ( + "topk_weights", + "topk_ids", + "sorted_token_ids", + "expert_ids", + "num_tokens_post_padded", + "expert_counts", + "expert_offsets", + "expert_scatter_offsets", +) + + +def _alphamoe_router_geometry( + num_tokens: int, num_experts: int, top_k: int, block_m: int +) -> Tuple[int, int]: + """(max_blocks, slots); depends only on the shape configuration.""" + pairs = num_tokens * top_k + nonempty = min(num_experts, pairs) + max_blocks = nonempty + (pairs - nonempty) // block_m + return max_blocks, max_blocks * block_m + + +class AlphaMoeRoutePlan: + """Persistent, shape-stable routing-plan buffers for the AlphaMoE router. + + The plan owns every output tensor of the fused routing bundle (plus the + internal scratch workspace). Allocating the buffers once before warmup + and reusing them for every routing call keeps tensor addresses fixed, so + a forward that routes through a plan can be captured into a CUDA graph + and replayed with new logit *values* in the same (or the captured) input + buffer; the replay then rewrites buffer contents only and no + re-allocation, re-planning, or host synchronization is required. + + Iterating or unpacking a plan yields the eight public output tensors in + the canonical order documented under :func:`alphamoe_fused_router`, and + the same tensors are also available as named attributes. + + Attributes + ---------- + num_tokens, num_experts, top_k, block_m : int + Geometry configuration the plan was allocated for. + has_shared_expert : bool + Shared-expert configuration the plan was allocated for. + """ + + def __init__( + self, + num_tokens: int, + num_experts: int, + top_k: int, + block_m: int, + has_shared_expert: bool = False, + device: Optional[torch.device] = None, + ) -> None: + self.num_tokens = int(num_tokens) + self.num_experts = int(num_experts) + self.top_k = int(top_k) + self.block_m = int(block_m) + self.has_shared_expert = bool(has_shared_expert) + max_blocks, slots = _alphamoe_router_geometry( + self.num_tokens, self.num_experts, self.top_k, self.block_m + ) + scratch_elems = ( + self.num_tokens * ((self.num_experts + 31) // 32) + if self.num_experts > 1024 + else 0 + ) + f32 = dict(dtype=torch.float32, device=device) + i32 = dict(dtype=torch.int32, device=device) + self.topk_weights = torch.empty((self.num_tokens, self.top_k), **f32) + self.topk_ids = torch.empty((self.num_tokens, self.top_k), **i32) + self.sorted_token_ids = torch.empty((slots,), **i32) + self.expert_ids = torch.empty((max_blocks,), **i32) + self.num_tokens_post_padded = torch.empty((1,), **i32) + self.expert_counts = torch.empty((self.num_experts,), **i32) + self.expert_offsets = torch.empty((self.num_experts + 1,), **i32) + self.expert_scatter_offsets = torch.empty((self.num_experts,), **i32) + # Internal workspace for the generic large-num_experts path; not part + # of the public bundle. + self.scratch = torch.empty((scratch_elems,), **i32) + self._tensors = tuple(getattr(self, name) for name in _PLAN_FIELDS) + + # Tuple emulation over the eight public outputs, so a plan can be used + # interchangeably with the return value of ``alphamoe_fused_router``. + def __iter__(self): + return iter(self._tensors) + + def __len__(self) -> int: + return len(self._tensors) + + def __getitem__(self, index): + return self._tensors[index] + + +@supported_compute_capability(_ALPHAMOE_ROUTER_SUPPORTED_CC) +def _check_alphamoe_router_vibecuda( + router_logits: torch.Tensor, + plan: Optional[AlphaMoeRoutePlan] = None, + top_k: Optional[int] = None, + block_m: Optional[int] = None, + has_shared_expert: bool = False, + backend: str = "vibecuda", +) -> bool: + """Validate dtypes, shapes, and configuration for the fused AlphaMoE router. + + Returns ``True`` when all inputs are valid and raises ``ValueError`` + otherwise, so direct FFI callers and the Python API share one contract. + """ + if router_logits.dim() != 2: + raise ValueError( + f"router_logits must be 2D [num_tokens, num_experts], got " + f"{tuple(router_logits.shape)}" + ) + if router_logits.dtype != torch.float32: + raise ValueError( + f"router_logits must be float32, got {router_logits.dtype}" + ) + if not router_logits.is_cuda: + raise ValueError("router_logits must be a CUDA tensor") + if not router_logits.is_contiguous(): + raise ValueError("router_logits must be contiguous") + + num_tokens, num_experts = router_logits.shape + if num_tokens < 1: + raise ValueError(f"num_tokens must be >= 1, got {num_tokens}") + if num_experts < 2: + raise ValueError(f"num_experts must be >= 2, got {num_experts}") + + if plan is not None: + if not isinstance(plan, AlphaMoeRoutePlan): + raise ValueError( + f"plan must be an AlphaMoeRoutePlan from " + f"allocate_alphamoe_route_plan, got {type(plan).__name__}" + ) + top_k = plan.top_k + block_m = plan.block_m + has_shared_expert = plan.has_shared_expert + if (plan.num_tokens, plan.num_experts) != (num_tokens, num_experts): + raise ValueError( + f"plan geometry ({plan.num_tokens} tokens, " + f"{plan.num_experts} experts) does not match router_logits " + f"({num_tokens} tokens, {num_experts} experts)" + ) + elif top_k is None or block_m is None: + raise ValueError( + "top_k and block_m are required when no route plan is provided" + ) + if not 1 <= int(top_k): + raise ValueError(f"top_k must be >= 1, got {top_k}") + if not 1 <= int(block_m): + raise ValueError(f"block_m must be >= 1, got {block_m}") + has_shared = int(bool(has_shared_expert)) + routed_top_k = int(top_k) - has_shared + routed_experts = num_experts - has_shared + if routed_top_k < 0 or routed_top_k > routed_experts: + raise ValueError( + f"invalid top_k ({top_k})/num_experts ({num_experts}) for the " + f"shared-expert configuration ({has_shared_expert})" + ) + return True + + +@functools.cache +def get_alphamoe_router_module(): + """Build, load, and cache the AlphaMoE fused-router JIT module as a custom op.""" + module = gen_alphamoe_router_module().build_and_load() + + @register_custom_op( + "flashinfer::alphamoe_fused_router", + mutates_args=[ + "topk_weights", + "topk_ids", + "sorted_token_ids", + "expert_ids", + "num_tokens_post_padded", + "expert_counts", + "expert_offsets", + "expert_scatter_offsets", + "scratch", + ], + ) + def alphamoe_fused_router( + router_logits: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + expert_counts: torch.Tensor, + expert_offsets: torch.Tensor, + expert_scatter_offsets: torch.Tensor, + scratch: torch.Tensor, + top_k: int, + block_m: int, + has_shared_expert: bool, + ) -> None: + """Custom-op wrapper that writes the routing bundle into the outputs.""" + module.alphamoe_fused_router( + router_logits, + topk_weights, + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + expert_counts, + expert_offsets, + expert_scatter_offsets, + scratch, + top_k, + block_m, + has_shared_expert, + ) + + return SimpleNamespace(alphamoe_fused_router=alphamoe_fused_router) + + +def allocate_alphamoe_route_plan( + router_logits: torch.Tensor, + top_k: int, + block_m: int, + has_shared_expert: bool = False, +) -> AlphaMoeRoutePlan: + """Allocate the persistent routing-plan buffers for one routing geometry. + + All output shapes of the fused AlphaMoE router depend only on + ``(num_tokens, num_experts, top_k, block_m)``, so the entire routing + bundle can be pre-allocated from the logits *shape* before any routing + call runs. Call this once outside the hot path (e.g. before CUDA graph + warmup/capture) and pass the returned plan to + :func:`alphamoe_fused_router`; every call then reuses the same buffer + addresses and only rewrites their contents. + + Parameters + ---------- + router_logits : torch.Tensor + Per-token expert logits of shape ``(num_tokens, num_experts)``, + ``float32``, on CUDA. Only the shape, dtype, and device are used; + the values are never read here. + top_k : int + Number of selected experts per token, including the shared expert + column when ``has_shared_expert`` is true. + block_m : int + Token-block alignment of the block-sparse MoE backend. + has_shared_expert : bool + Whether the last expert id is a shared expert that every token also + routes to. Default ``False``. + + Returns + ------- + AlphaMoeRoutePlan + The persistent routing-plan buffers, iterable/unpackable as the + eight-tuple documented under :func:`alphamoe_fused_router`. + """ + if not isinstance(router_logits, torch.Tensor) or router_logits.dim() != 2: + raise ValueError( + "router_logits must be a 2D [num_tokens, num_experts] tensor" + ) + if not router_logits.is_cuda: + raise ValueError("router_logits must be a CUDA tensor") + num_tokens, num_experts = router_logits.shape + return AlphaMoeRoutePlan( + num_tokens, + num_experts, + top_k, + block_m, + has_shared_expert, + device=router_logits.device, + ) + + +@backend_requirement({"vibecuda": _check_alphamoe_router_vibecuda}) +@flashinfer_api(trace=alphamoe_fused_router_trace) +def alphamoe_fused_router( + router_logits: torch.Tensor, + plan: Optional[AlphaMoeRoutePlan] = None, + top_k: Optional[int] = None, + block_m: Optional[int] = None, + has_shared_expert: bool = False, + backend: str = "vibecuda", +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + r"""Fused AlphaMoE gating router for a block-sparse MoE backend. + + One fused call consumes the per-token expert logits and emits the complete + routing metadata bundle: stable top-k selection with an optional shared + expert, softmax weights, the expert histogram, block-m-aligned padded + expert offsets with the padded extent, expert-grouped sorted route ids, + per-block expert ids, and per-expert scatter offsets. + + Selection for each token picks ``routed_top_k = top_k - + int(has_shared_expert)`` experts from ``router_logits[:, :num_experts - + int(has_shared_expert)]`` in descending logit order with a stable + tie-break (equal logits keep the lower expert index first, i.e. the + result of a stable descending sort over the row). When a shared + expert is present, expert id ``num_experts - 1`` is appended as the last + selected column with ``router_logits[:, -1]`` as its value. + ``topk_weights`` is the max-subtracted fp32 softmax over the ``top_k`` + selected logits per token. + + Route ids in ``sorted_token_ids`` are flat ``token * top_k + slot`` + indices in increasing order within each expert segment; in-segment + padding slots carry the sentinel ``num_tokens * top_k`` and slots outside + every expert segment stay zero. ``expert_ids`` is zero outside real + expert blocks. + + All output shapes depend only on + ``(num_tokens, num_experts, top_k, block_m)``, every output is fully + written on each call, and no host synchronization happens on the launch + path, so the operator is safe under CUDA graph capture: replaying a + captured graph with updated ``router_logits`` values in the captured + buffer reproduces the eagerly computed routing bundle for those values. + Pass a plan from :func:`allocate_alphamoe_route_plan` (allocated before + warmup/capture) to keep every output address fixed across capture and + replay. + + Parameters + ---------- + router_logits : torch.Tensor + Per-token expert logits of shape ``(num_tokens, num_experts)``, + ``float32``, contiguous, on CUDA. + plan : AlphaMoeRoutePlan, optional + Persistent routing-plan buffers from + :func:`allocate_alphamoe_route_plan`. When provided, the routing + geometry is taken from the plan and results are written into the + plan's buffers; ``top_k``/``block_m``/``has_shared_expert`` must be + left at their defaults. When omitted, ``top_k`` and ``block_m`` are + required and fresh output tensors are allocated per call. + top_k : int, optional + Number of selected experts per token, including the shared expert + column when ``has_shared_expert`` is true. Required when no ``plan`` + is provided. + block_m : int, optional + Token-block alignment of the block-sparse MoE backend. Required when + no ``plan`` is provided. + has_shared_expert : bool + Whether the last expert id is a shared expert that every token also + routes to. Default ``False``. + backend : str + Backend selector. Only ``"vibecuda"`` (the custom CUDA + implementation) is available; selecting it on an unsupported + architecture raises instead of silently rerouting. + + Returns + ------- + topk_weights : torch.Tensor + Softmaxed routing weights, ``(num_tokens, top_k)`` ``float32``. + topk_ids : torch.Tensor + Selected expert ids, ``(num_tokens, top_k)`` ``int32``. + sorted_token_ids : torch.Tensor + Expert-grouped flat route ids, ``(max_blocks * block_m,)`` ``int32``. + expert_ids : torch.Tensor + Per-block expert id, ``(max_blocks,)`` ``int32``. + num_tokens_post_padded : torch.Tensor + Total padded token extent, ``(1,)`` ``int32``. + expert_counts : torch.Tensor + Selected routes per expert, ``(num_experts,)`` ``int32``. + expert_offsets : torch.Tensor + Exclusive prefix of the padded counts, ``(num_experts + 1,)`` + ``int32``. + expert_scatter_offsets : torch.Tensor + Per-expert scatter offsets, ``(num_experts,)`` ``int32``; identical + to ``expert_counts`` (the upstream routing plan returns + ``counts.clone()`` for this output). + """ + if plan is None: + if top_k is None or block_m is None: + raise ValueError( + "top_k and block_m are required when no plan is provided" + ) + num_tokens, num_experts = router_logits.shape + plan = AlphaMoeRoutePlan( + num_tokens, + num_experts, + int(top_k), + int(block_m), + has_shared_expert, + device=router_logits.device, + ) + get_alphamoe_router_module().alphamoe_fused_router( + router_logits, + plan.topk_weights, + plan.topk_ids, + plan.sorted_token_ids, + plan.expert_ids, + plan.num_tokens_post_padded, + plan.expert_counts, + plan.expert_offsets, + plan.expert_scatter_offsets, + plan.scratch, + plan.top_k, + plan.block_m, + plan.has_shared_expert, + ) + return tuple(plan) diff --git a/flashinfer/jit/__init__.py b/flashinfer/jit/__init__.py index c92bacd0979..d5100ee1355 100644 --- a/flashinfer/jit/__init__.py +++ b/flashinfer/jit/__init__.py @@ -97,6 +97,9 @@ from .tinygemm2 import gen_tinygemm2_sm100_module as gen_tinygemm2_sm100_module from .moe_utils import gen_moe_utils_module as gen_moe_utils_module from .hash_topk import gen_hash_topk_module as gen_hash_topk_module +from .alphamoe_router import ( + gen_alphamoe_router_module as gen_alphamoe_router_module, +) from .fp4_kv_dequantization import ( gen_fp4_kv_dequantization_module as gen_fp4_kv_dequantization_module, ) diff --git a/flashinfer/jit/alphamoe_router.py b/flashinfer/jit/alphamoe_router.py new file mode 100644 index 00000000000..3187f0270e3 --- /dev/null +++ b/flashinfer/jit/alphamoe_router.py @@ -0,0 +1,35 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import functools + +from . import env as jit_env +from .core import JitSpec, gen_jit_spec + + +@functools.cache +def gen_alphamoe_router_module() -> JitSpec: + """Build the JIT spec for the fused AlphaMoE gating router module.""" + return gen_jit_spec( + "alphamoe_router", + [ + jit_env.FLASHINFER_CSRC_DIR / "alphamoe_router" / "alphamoe_router.cu", + jit_env.FLASHINFER_CSRC_DIR + / "alphamoe_router" + / "alphamoe_router_jit_binding.cu", + ], + extra_cuda_cflags=["-lineinfo"], + ) diff --git a/flashinfer/trace/templates/moe.py b/flashinfer/trace/templates/moe.py index d907de1761a..3ff9a6ef36f 100644 --- a/flashinfer/trace/templates/moe.py +++ b/flashinfer/trace/templates/moe.py @@ -4617,3 +4617,158 @@ def _trtllm_gen_routing_init( tags=["status:verified", "moe", "moe:routing"], init=_trtllm_gen_routing_init, ) + + +# --------------------------------------------------------------------------- +# Fused AlphaMoE gating router ("vibecuda" backend) +# --------------------------------------------------------------------------- + + +@torch.no_grad() +def _alphamoe_fused_router_reference( + router_logits: torch.Tensor, + top_k: int, + block_m: int, + has_shared_expert: bool = False, + **_unused, +): + """Exact torch reference for the fused AlphaMoE gating router. + + Stable descending top-k over the routed experts (ties keep the lower + expert index), optional shared-expert column, fp32 max-subtracted + softmax over the selected logits, then the block-sparse routing + metadata: expert histogram, block_m-aligned padded offsets and extent, + per-expert scatter offsets (identical to the expert histogram, mirroring + the upstream ``counts.clone()`` routing-plan output), expert-grouped + flat route ids, and per-block expert ids. + """ + num_tokens, num_experts = router_logits.shape + routed_experts = num_experts - int(has_shared_expert) + routed_top_k = top_k - int(has_shared_expert) + order = torch.argsort( + router_logits[:, :routed_experts], dim=-1, descending=True, stable=True + )[:, :routed_top_k] + selected = torch.gather(router_logits, 1, order) + if has_shared_expert: + shared = torch.full( + (num_tokens, 1), + num_experts - 1, + dtype=torch.int64, + device=router_logits.device, + ) + order = torch.cat((order, shared), dim=-1) + selected = torch.cat((selected, router_logits[:, -1:]), dim=-1) + topk_ids = order.to(torch.int32) + topk_weights = torch.softmax(selected, dim=-1) + flat = topk_ids.flatten().to(torch.int64) + counts = torch.bincount(flat, minlength=num_experts).to(torch.int32) + padded = (counts + block_m - 1) // block_m * block_m + offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=router_logits.device) + offsets[1:] = torch.cumsum(padded, dim=0) + scatter_offsets = counts.clone() + pairs = num_tokens * top_k + nonempty = min(num_experts, pairs) + max_blocks = nonempty + (pairs - nonempty) // block_m + sorted_ids = torch.zeros( + max_blocks * block_m, dtype=torch.int32, device=router_logits.device + ) + expert_ids = torch.zeros(max_blocks, dtype=torch.int32, device=router_logits.device) + sentinel = pairs + for expert in range(num_experts): + start = int(offsets[expert].item()) + count = int(counts[expert].item()) + end = int(offsets[expert + 1].item()) + if count: + routes = torch.nonzero(flat == expert).flatten().to(torch.int32) + sorted_ids[start : start + count] = routes + sorted_ids[start + count : end] = sentinel + expert_ids[start // block_m : end // block_m] = expert + extent = offsets[-1:].clone() + return ( + topk_weights, + topk_ids, + sorted_ids, + expert_ids, + extent, + counts, + offsets, + scatter_offsets, + ) + + +def _alphamoe_fused_router_init( + *, + num_tokens: int, + num_experts: int = 256, + top_k: int = 8, + block_m: int = 16, + has_shared_expert: bool = False, + # Derived by the exact output geometry; accepted only so the signature + # carries every Var axis, and recomputed rather than used. + max_blocks: int = 0, + slots: int = 0, + one: int = 1, + num_experts_plus_one: int = 0, + device: str = "cuda", + seed: int = 0, +): + """Build inputs for the fused AlphaMoE gating router.""" + torch.manual_seed(seed) + router_logits = torch.randn( + num_tokens, num_experts, dtype=torch.float32, device=device + ) + return { + "router_logits": router_logits, + "top_k": int(top_k), + "block_m": int(block_m), + "has_shared_expert": bool(has_shared_expert), + } + + +alphamoe_fused_router_trace = TraceTemplate( + op_type="moe_routing", + name_prefix="alphamoe_fused_router", + description=( + "Fused AlphaMoE gating router: stable descending top-k with an " + "optional shared-expert column -> fp32 max-subtracted softmax -> " + "block_m-aligned routing metadata bundle (expert histogram, padded " + "expert offsets + extent, scatter offsets equal to the expert " + "counts, expert-grouped flat route ids with sentinel padding, " + "per-block expert ids). Route ids are flat token*top_k+slot indices; " + "in-segment padding slots carry the sentinel num_tokens*top_k." + ), + axes={ + "num_tokens": Var(), + "num_experts": Const(abbrev="e"), + "top_k": Const(abbrev="k"), + "block_m": Const(abbrev="b"), + "max_blocks": Var( + description="Block-bound from the output geometry " + "(min(num_experts, pairs) + (pairs - nonempty) // block_m)." + ), + "slots": Var(description="max_blocks * block_m."), + "one": Var(description="Placeholder for shape [1] output tensors."), + "num_experts_plus_one": Var( + description="num_experts + 1 for the inclusive padded offsets." + ), + }, + inputs={ + "router_logits": Tensor(["num_tokens", "num_experts"]), + "top_k": Scalar("int32"), + "block_m": Scalar("int32"), + "has_shared_expert": Scalar("bool"), + }, + outputs={ + "topk_weights": Tensor(["num_tokens", "top_k"], dtype="float32"), + "topk_ids": Tensor(["num_tokens", "top_k"], dtype="int32"), + "sorted_token_ids": Tensor(["slots"], dtype="int32"), + "expert_ids": Tensor(["max_blocks"], dtype="int32"), + "num_tokens_post_padded": Tensor(["one"], dtype="int32"), + "expert_counts": Tensor(["num_experts"], dtype="int32"), + "expert_offsets": Tensor(["num_experts_plus_one"], dtype="int32"), + "expert_scatter_offsets": Tensor(["num_experts"], dtype="int32"), + }, + tags=["status:verified", "moe", "moe:routing"], + reference=_alphamoe_fused_router_reference, + init=_alphamoe_fused_router_init, +) diff --git a/include/flashinfer/fused_moe/alphamoe_router.cuh b/include/flashinfer/fused_moe/alphamoe_router.cuh new file mode 100644 index 00000000000..109e10fe922 --- /dev/null +++ b/include/flashinfer/fused_moe/alphamoe_router.cuh @@ -0,0 +1,1881 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file alphamoe_router.cuh + * \brief Fused AlphaMoE gating router ("vibecuda" backend, plain CUDA). + * + * Consumes per-token expert logits and emits the complete block-sparse MoE + * routing metadata bundle in a small, shape-adaptive kernel pipeline: + * - stable descending top-k selection over the routed experts (equal logits + * keep the lower expert index first; stable sort order), with the + * optional shared expert appended as the last selected column, + * - max-subtracted fp32 softmax over the selected logits, + * - expert histogram, block_m-aligned padded expert offsets, the + * per-expert scatter offsets (identical to the expert histogram), and + * the padded extent, + * - deterministic expert-grouped sorted route ids (flat token*top_k+slot + * route indices, increasing token order) and per-block expert ids. + * + * Dispatch (all thresholds measured on Blackwell B300): + * small tiles (experts <= 512, tokens <= 16, pairs <= 256) run as a single + * fused_small_kernel: selection + histogram + scan + scatter + fills in one + * block behind block barriers, so a forward costs one launch instead of + * three. Medium small inputs (num_experts <= 1024) run as: + * select_kernel : warp-per-token stable top-k selection with + * register-cached ordering keys and fp32 max-subtracted softmax. + * finish_kernel (tokens <= 256) : every block redundantly builds the + * expert histogram, (expert x token) route bitmap and padded exclusive + * scan in smem, then all blocks grid-stride the bitmap-rank scatter and + * the sentinel/zero fills; block 0 stores the histogram/scan outputs. + * reduce_kernel + tail_kernel (tokens > 256) : single-block + * histogram/scan, then a per-expert-warp ballot rescan of the route + * stream. + * Dependent kernels overlap predecessor tails via Programmatic Dependent + * Launch on SM90+ (pdl_sync gates every global read; a plain launch is used + * on older architectures). Large inputs (num_experts > 1024) take a generic + * grid-parallel path (select with global histogram atomics + single-block + * scan + grid scatter) with identical semantics; it requires one scratch + * element per route (num_tokens * ceil(num_experts/32) ints). + * + * The file is framework-agnostic: no torch headers. Callers allocate all + * outputs (and the generic-path scratch) and pass raw pointers plus the + * current CUDA stream. All outputs are fully written on every call, shapes + * depend only on (num_tokens, num_experts, top_k, block_m), and there is no + * host synchronization, so the interface is CUDA-graph capture safe. + */ + +#pragma once + +#include + +#include +#include +#include + +namespace flashinfer::fused_moe { + + +constexpr int kSelectThreads = 128; +constexpr int kReduceThreads = 512; +constexpr int kScatterThreads = 128; +constexpr int kScanThreads = 1024; +constexpr int kReduceMaxExperts = 1024; +constexpr int kSbMaxTwords = 8; // route-bitmap scatter covers tokens <= 256 +constexpr int kSbBitmapInts = kReduceMaxExperts * kSbMaxTwords; // 32KB smem +constexpr int kSBlkCap = 1000; // finish_kernel block->expert flat-map capacity +// finish_kernel's scan phase (which already touches every (word, expert) +// pair with conflict-free column accesses) stamps each bitmap word's base +// slot (s_off[e] + column prefix) into a small dynamic-smem table. The +// route-parallel scatter then needs ONE table load + ONE bitmask popcount +// per route instead of TW random-column popcounts, keeping full grid +// coverage and zero extra global traffic. Dynamic smem keeps the static +// footprint under the 48KB compiler limit. +// Single-block fully-fused path (small tiles): the whole 5-stage pipeline in +// one launch, removing the launch-chain + PDL floor that dominates tiny +// forwards. Guards: num_experts <= kFusedMaxExperts, pairs <= kFusedPairsCap, +// tokens <= 16 (one token per select warp). +constexpr int kFusedThreads = 512; +constexpr int kFusedMaxExperts = 512; +constexpr int kFusedMaxTwords = 8; +constexpr int kFusedPairsCap = 256; + +// Programmatic Dependent Launch helpers (SM90+). Successor kernels sync +// against their stream predecessors before touching any global memory, so +// launch scaffolding and shared-memory setup overlap the predecessor's tail. +__device__ __forceinline__ void pdl_sync() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); +#endif +} + +__device__ __forceinline__ void pdl_trigger() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// Launch with the programmatic-stream-serialization attribute so the kernel +// may begin its (sync-guarded) prologue while the predecessor finishes. +template +void launch_pdl(Kernel kernel, int blocks, int threads, size_t smem_bytes, + cudaStream_t stream, Args... args) { + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3((unsigned)blocks); + cfg.blockDim = dim3((unsigned)threads); + cfg.dynamicSmemBytes = smem_bytes; + cfg.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + cfg.attrs = attrs; + cfg.numAttrs = 1; + const cudaError_t err = cudaLaunchKernelEx(&cfg, kernel, args...); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("alphamoe_router: PDL launch " + "failed: ") + + cudaGetErrorString(err)); + } +} + +// Monotonic map float -> uint32 so that descending float order matches +// descending integer order (handles the full IEEE ordering incl. negatives). +__device__ __forceinline__ unsigned int float_sortable_key(float f) { + unsigned int u = __float_as_uint(f); + return u ^ ((u & 0x80000000u) ? 0xFFFFFFFFu : 0x80000000u); +} + +// Inverse of float_sortable_key: recover the original float from its key. +__device__ __forceinline__ float float_from_sortable_key(unsigned int u) { + unsigned int v = (u & 0x80000000u) ? (u ^ 0x80000000u) : (u ^ 0xFFFFFFFFu); + return __uint_as_float(v); +} + +// Exact branch-free unsigned division by a kernel-constant divisor +// (Granlund-Montgomery round-up method, same family as libdivide's +// branchfree u32 path). top_k and block_m are Model-constructor constants, +// so the launcher builds the magic pair once per forward on the host and +// passes it by value; on device each division is one mulhi + 3 cheap ops +// instead of the compiler's ~15-instruction runtime s32 division sequence. +// Exact for every 32-bit dividend and divisor >= 1. +struct FastDivU32 { + unsigned int magic; // 0 marks the power-of-two (shift) path + unsigned int shift; +}; + +static inline FastDivU32 make_fast_div_u32(unsigned int d) { + FastDivU32 fd; + if ((d & (d - 1u)) == 0u) { // includes d == 1 (shift 0) + fd.magic = 0u; + fd.shift = (unsigned int)__builtin_ctz(d); + } else { + const unsigned int l = 31u - (unsigned int)__builtin_clz(d); + const unsigned long long two_pow = 1ull << (32 + l); + unsigned long long m = two_pow / d; + unsigned long long rem = two_pow - m * (unsigned long long)d; + m += m; + rem += rem; + if (rem >= (unsigned long long)d) { + m += 1; + } + m += 1; + fd.magic = (unsigned int)m; // stores the low 32 bits (33rd implied) + fd.shift = l; + } + return fd; +} + +__device__ __forceinline__ unsigned int fdiv_u32(unsigned int n, + FastDivU32 fd) { + if (fd.magic == 0u) { + return n >> fd.shift; + } + const unsigned int t = __umulhi(n, fd.magic); + return (t + ((n - t) >> 1)) >> fd.shift; +} + +// 64-bit ordering key: descending logit, ties keep the lower expert index +// first. Keys are unique per expert so iterating "argmax key below previous +// winner" reproduces exactly the stable descending index order. +__device__ __forceinline__ unsigned long long route_key(float f, int expert) { + unsigned long long hi = (unsigned long long)float_sortable_key(f); + unsigned long long lo = (unsigned long long)(0xFFFFFFFFu - (unsigned int)expert); + return (hi << 32) | lo; +} + +// Warp-wide stable top-k selection for one token. CPL is a compile-time cap +// on ceil(routed_experts / 32) so candidate keys live in registers. Lane +// (j & 31) owns selection round j: records id/logit. The per-round warp max +// uses two hardware redux ops instead of a 10-instruction shuffle. SHARED is +// compile-time so the shared-expert round stays out of the selection loop. +// EMIT additionally bumps the smem expert histogram and (token, expert) route +// bitmap from the owning lane's registers, removing the separate phase that +// re-reads the just-written topk ids from global memory (fused small path). +// Round j's selected expert is also returned through own_expert for lanes +// where lane == (j & 31) (only tracked while j < 32; -1 otherwise). +template +__device__ __forceinline__ void token_select( + const float* __restrict__ row, + float* __restrict__ w_row, + int* __restrict__ id_row, + int lane, + int num_experts, + int top_k, + int routed_top_k, + int* __restrict__ own_expert = nullptr, + int* __restrict__ s_hist = nullptr, + unsigned int* __restrict__ s_bitmap = nullptr, + int token = 0) { + const int routed_experts = num_experts - (SHARED ? 1 : 0); + + unsigned long long keys[CPL]; +#pragma unroll + for (int c = 0; c < CPL; ++c) { + const int e = lane + c * 32; + keys[c] = (e < routed_experts) ? route_key(row[e], e) : 0ull; + } + + int own = -1; + unsigned long long prev = 0xFFFFFFFFFFFFFFFFull; + // Register-distributed softmax (top_k <= 32): each lane keeps the value + // of its own selection round, avoiding a global round-trip through w_row. + const bool reg_softmax = top_k <= 32; + float v_acc = -INFINITY; + for (int j = 0; j < routed_top_k; ++j) { + // Mask out keys at/above the previous winner, then take the max + // as a pairwise tree (log2(CPL) compare depth instead of a + // CPL-long dependent chain -- this is the per-round critical path). + unsigned long long cand[CPL]; +#pragma unroll + for (int c = 0; c < CPL; ++c) { + cand[c] = (keys[c] < prev) ? keys[c] : 0ull; + } +#pragma unroll + for (int step = CPL / 2; step > 0; step >>= 1) { +#pragma unroll + for (int c = 0; c < step; ++c) { + cand[c] = cand[c] > cand[c + step] ? cand[c] : cand[c + step]; + } + } + const unsigned long long best = cand[0]; + const unsigned int lane_hi = (unsigned int)(best >> 32); + const unsigned int best_hi = __reduce_max_sync(0xFFFFFFFFu, lane_hi); + const unsigned int tag = + (lane_hi == best_hi) ? (unsigned int)best : 0u; + const unsigned int best_lo = __reduce_max_sync(0xFFFFFFFFu, tag); + prev = ((unsigned long long)best_hi << 32) | best_lo; + const int expert = (int)(0xFFFFFFFFu - best_lo); + const float value = float_from_sortable_key(best_hi); // logit in key + if (reg_softmax) { + // All lanes track their own round's value via a cheap select. + v_acc = (lane == (j & 31)) ? value : v_acc; + } + if (lane == (j & 31)) { + id_row[j] = expert; + if (!reg_softmax) { + w_row[j] = value; // scratch: selected logits, softmax below + } + if (j < 32) { + own = expert; + } + if (EMIT == 1) { + atomicAdd(s_hist + expert, 1); + atomicOr(s_bitmap + expert, 1u << token); + } else if (EMIT == 3) { + // Bitmap only: counts derivable as bitmap popcounts when + // num_tokens <= block_m (POPC fused path). + atomicOr(s_bitmap + expert, 1u << token); + } else if (EMIT == 2) { + // Single-token path: expert bitmap over words + plain route + // record (selected experts are distinct, so no contention). + atomicOr(s_bitmap + (expert >> 5), 1u << (expert & 31)); + s_hist[expert] = j; + } + } + } + // Shared-expert round (compile-time presence): last logit column, id + // num_experts-1, appended as the final topk column. + if (SHARED) { + const int j = routed_top_k; + const float value = row[num_experts - 1]; + if (reg_softmax) { + v_acc = (lane == (j & 31)) ? value : v_acc; + } + if (lane == (j & 31)) { + id_row[j] = num_experts - 1; + if (!reg_softmax) { + w_row[j] = value; + } + if (j < 32) { + own = num_experts - 1; + } + if (EMIT == 1) { + atomicAdd(s_hist + (num_experts - 1), 1); + atomicOr(s_bitmap + (num_experts - 1), 1u << token); + } else if (EMIT == 3) { + atomicOr(s_bitmap + (num_experts - 1), 1u << token); + } else if (EMIT == 2) { + atomicOr(s_bitmap + ((num_experts - 1) >> 5), + 1u << ((num_experts - 1) & 31)); + s_hist[num_experts - 1] = j; + } + } + } + if (own_expert != nullptr) { + *own_expert = own; + } + __syncwarp(0xFFFFFFFFu); + + const int rounds = routed_top_k + (SHARED ? 1 : 0); + if (reg_softmax) { + float m = v_acc; +#pragma unroll + for (int ofs = 16; ofs > 0; ofs >>= 1) { + const float other = __shfl_xor_sync(0xFFFFFFFFu, m, ofs); + m = fmaxf(m, other); + } + const float e = (lane < rounds) ? expf(v_acc - m) : 0.0f; + float sum = e; +#pragma unroll + for (int ofs = 16; ofs > 0; ofs >>= 1) { + sum += __shfl_xor_sync(0xFFFFFFFFu, sum, ofs); + } + if (lane < rounds) { + w_row[lane] = e / sum; + } + return; + } + + float sm = -INFINITY; + for (int j = lane; j < top_k; j += 32) { + sm = fmaxf(sm, w_row[j]); + } + for (int ofs = 16; ofs > 0; ofs >>= 1) { + float other = __shfl_xor_sync(0xFFFFFFFFu, sm, ofs); + sm = fmaxf(sm, other); + } + float ssum = 0.0f; + for (int j = lane; j < top_k; j += 32) { + ssum += expf(w_row[j] - sm); + } + for (int ofs = 16; ofs > 0; ofs >>= 1) { + ssum += __shfl_xor_sync(0xFFFFFFFFu, ssum, ofs); + } + const float inv = 1.0f / ssum; + for (int j = lane; j < top_k; j += 32) { + w_row[j] = expf(w_row[j] - sm) * inv; + } +} + + +// CPL=16 expert band with register-softmax top_k (top_k <= 32): four +// warps share one token's tournament (select_quad_kernel, 128-thread block, +// CPL4 keys per lane per warp): warp w owns experts [w*32*CPL4, +// (w+1)*32*CPL4) and each round's four local winners combine through shared +// memory behind one double-buffered block barrier (round j touches slot j&1 +// only, so one barrier suffices). Selection math is identical to the +// single-warp tournament -- the round max is the max of the four local +// maxima of keys below prev -- so the stable-tie output matches bit for +// bit. Narrower bands and the fused paths keep the single-warp tournament. +// Fully-fused single-block path for small tiles. One launch runs selection, +// histogram, padded scan and bitmap-rank scatter with block barriers in place +// of the multi-kernel PDL chain, so the whole forward costs one launch +// instead of three. One warp per token (up to 16); num_tokens <= 16 keeps +// the (expert x token) smem bitmap to one word per expert and warp 0's +// serial padded scan to per-lane chunks of at most 16 experts. +// POPC (num_tokens <= block_m): every expert count is bounded by num_tokens, +// so counts are recovered as bitmap popcounts and the smem histogram zero +// loop plus the per-round atomicAdds disappear entirely. +template +__global__ void __launch_bounds__(kFusedThreads) +fused_small_kernel(const float* __restrict__ logits, + float* __restrict__ topk_weights, + int* __restrict__ topk_ids, + int* __restrict__ expert_counts, + int* __restrict__ expert_offsets, + int* __restrict__ num_tokens_post_padded, + int* __restrict__ expert_scatter_offsets, + int* __restrict__ block_expert_ids, + int* __restrict__ sorted_token_ids, + int num_tokens, + int num_experts, + int top_k, + int block_m, + int routed_top_k, + int pairs, + int max_blocks, + int slots, + int twords, + FastDivU32 fd_topk, + FastDivU32 fd_bm) { + __shared__ int s_hist[kFusedMaxExperts]; + __shared__ unsigned int s_bitmap[kFusedMaxExperts * kFusedMaxTwords]; + __shared__ int s_off[kFusedMaxExperts + 1]; + __shared__ int s_wsum[kFusedThreads / 32]; + + const int tid = (int)threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + constexpr int kWarps = kFusedThreads / 32; + + // num_tokens <= 16 in this path, so the route bitmap is one word per + // expert (the twords param is always 1; see launch_fused). + for (int e = tid; e < num_experts; e += kFusedThreads) { + if (!POPC) { + s_hist[e] = 0; + } + s_bitmap[e] = 0u; + } + __syncthreads(); + + // Warp-per-token stable selection (num_tokens <= warps in this path: one + // token per warp). Emitting warps bump the smem histogram/bitmap from + // registers as rounds resolve, removing the idle reload phase over the + // just-written global ids. Warps without a token instead zero the two + // global outputs *concurrently* with selection, so the post-scan leg + // only writes real route ids, in-segment padding sentinels, and real + // per-block expert ids; the barrier below orders these plain stores + // before the later overlapping stores by other threads. + int own_expert = -1; + const int token = warp; + if (token < num_tokens) { + token_select( + logits + (long long)token * num_experts, + topk_weights + (long long)token * top_k, + topk_ids + (long long)token * top_k, lane, num_experts, top_k, + routed_top_k, &own_expert, s_hist, s_bitmap, token); + } else { + const int ptid = tid - num_tokens * 32; + const int pthreads = kFusedThreads - num_tokens * 32; + for (int s = ptid; s < slots; s += pthreads) { + sorted_token_ids[s] = 0; + } + for (int b = ptid; b < max_blocks; b += pthreads) { + block_expert_ids[b] = 0; + } + } + __syncthreads(); + + // Block_m-padded exclusive scan (single chunk: num_experts <= threads). + // Round-31 latency design: the zero tails were already filled by the + // non-emitter warps during selection, and the tail below collapses + // scatter + padding + ownership into one barrier-free final leg, so the + // old 6-barrier phase machine (scan / s_off / ownership+rendezvous / + // three strided fill passes) needs only 3 barriers total. Every warp + // redundantly scans the 16 warp sums in registers and shuffles out its + // own base (the finish_kernel round-11 pattern), so no warp-0-only + // serial section lengthens the scan barrier's arrival leg -- the + // r16-r18 single-block rule never violated here. (A round-31A + // single-warp serial scan over per-lane expert chunks was REVERTED on + // measurement: production 5.54us vs this hybrid's target ~5.1us at + // (8,257,9,8); the warp-0-only leg between the two remaining barriers + // simply moved the rendezvous cost onto the scan barrier.) + int count = 0; + if (tid < num_experts) { + count = __popc(s_bitmap[tid]); // one bitmap word per expert here + s_hist[tid] = count; // normalize for the segment fills + expert_counts[tid] = count; + expert_scatter_offsets[tid] = count; + } + const int value = + (int)(fdiv_u32((unsigned int)(count + block_m - 1), fd_bm) * + (unsigned int)block_m); + int scan = value; + for (int ofs = 1; ofs < 32; ofs <<= 1) { + const int other = __shfl_up_sync(0xFFFFFFFFu, scan, ofs); + if (lane >= ofs) { + scan += other; + } + } + if (lane == 31) { + s_wsum[warp] = scan; + } + __syncthreads(); + int wsum = (lane < kWarps) ? s_wsum[lane] : 0; + for (int ofs = 1; ofs < kWarps; ofs <<= 1) { + const int other = __shfl_up_sync(0xFFFFFFFFu, wsum, ofs); + if (lane >= ofs) { + wsum += other; + } + } + const int warp_base = + (warp > 0) ? __shfl_sync(0xFFFFFFFFu, wsum, warp - 1) : 0; + const int inclusive = warp_base + scan; + if (tid < num_experts) { + expert_offsets[tid + 1] = inclusive; + s_off[tid + 1] = inclusive; + } + const int extent = __shfl_sync(0xFFFFFFFFu, wsum, kWarps - 1); + if (tid == 0) { + expert_offsets[0] = 0; + s_off[0] = 0; + num_tokens_post_padded[0] = extent; + } + __syncthreads(); + + // Final leg (no barrier after): emitters scatter their kept route from + // the register-cached expert (rank = number of earlier tokens selecting + // the same expert), and every thread tid < num_experts writes its + // expert's in-segment padding sentinels and per-block expert ids. + if (top_k <= 32 && token < num_tokens) { + if (own_expert >= 0) { + const int e = own_expert; + const int rank = + __popc(s_bitmap[e] & ((1u << token) - 1u)); + sorted_token_ids[s_off[e] + rank] = token * top_k + lane; + } + } else if (top_k > 32) { + for (int r = tid; r < pairs; r += kFusedThreads) { + const int token = (int)fdiv_u32((unsigned int)r, fd_topk); + const int e = topk_ids[r]; + const int rank = __popc(s_bitmap[e] & ((1u << token) - 1u)); + sorted_token_ids[s_off[e] + rank] = r; + } + } + if (tid < num_experts) { + const int off = s_off[tid]; + const int end = s_off[tid + 1]; + for (int s = off + s_hist[tid]; s < end; ++s) { + sorted_token_ids[s] = pairs; // padding sentinel + } + for (int b = (int)fdiv_u32((unsigned int)off, fd_bm); + b < (int)fdiv_u32((unsigned int)end, fd_bm); ++b) { + block_expert_ids[b] = tid; + } + } + // num_tokens == kFusedThreads/32: no warp was free for the selection-leg + // prefill, so the zero tails are written here instead (segment interiors + // are fully covered by the scatter + sentinel stores above). + if (num_tokens * 32 == kFusedThreads) { + const int extent = s_off[num_experts]; + const int used_blocks = (int)fdiv_u32((unsigned int)extent, fd_bm); + for (int s = extent + tid; s < slots; s += kFusedThreads) { + sorted_token_ids[s] = 0; + } + for (int b = used_blocks + tid; b < max_blocks; b += kFusedThreads) { + block_expert_ids[b] = 0; + } + } +} + +// Single-warp fully-fused path for exactly one token. The 512-thread fused +// path pays phase-separation __syncthreads barriers for phases a single warp +// can run serially with __syncwarp only: histogram/scan/fills here collapse +// to one smem expert-bitmap pass plus one 32-expert-per-lane chunk sweep. +// Valid for num_experts <= kFusedMaxExperts and top_k <= 32 (one selection +// round per lane); count per expert is 1 (selection rounds choose distinct +// experts), so each padded segment is exactly block_m slots. +template +__global__ void fused_single_kernel(const float* __restrict__ logits, + float* __restrict__ topk_weights, + int* __restrict__ topk_ids, + int* __restrict__ expert_counts, + int* __restrict__ expert_offsets, + int* __restrict__ num_tokens_post_padded, + int* __restrict__ expert_scatter_offsets, + int* __restrict__ block_expert_ids, + int* __restrict__ sorted_token_ids, + int num_experts, + int top_k, + int block_m, + int routed_top_k, + int max_blocks, + int slots, + FastDivU32 fd_bm) { + __shared__ int s_seg[32]; // segment index -> expert id + __shared__ int s_route[kFusedMaxExperts]; + __shared__ unsigned int s_bits[kFusedMaxExperts / 32]; + + const int lane = (int)threadIdx.x; + const int words = (num_experts + 31) / 32; + for (int w = lane; w < words; w += 32) { + s_bits[w] = 0u; + } + __syncwarp(0xFFFFFFFFu); + + int own_expert = -1; + token_select(logits, topk_weights, topk_ids, lane, + num_experts, top_k, routed_top_k, &own_expert, + s_route, s_bits, 0); + + // Per-lane chunk of 32 experts: prefix over chunk bitmaps, then a + // popcount-parallel sweep (no serial carry chain). counts, scatter and + // offsets all go out as aligned int4 vectors: writing the UNSHIFTED + // offsets window (offsets[e] = selected-count prefix below e times + // block_m) keeps 16B alignment, with the offsets[num_experts] tail + // stored by the last-chunk lane. The earlier per-element store loop + // with a serial carry measured ~1us slower here: it serialized 100+ + // dependent stores per lane. + const int rounds = routed_top_k + (SHARED ? 1 : 0); + const int extent = rounds * block_m; + // Round-26 re-chunk experiment, retained on measured latency: 16-expert + // chunks. The original sweep maps one 32-expert chunk per lane, so at + // num_experts <= kFusedMaxExperts (512) only ceil(num_experts/32) lanes + // work (16 of 32 at 512 experts) while the other half of the warp sits + // idle through the 8-quad serial store chain. 16-expert chunks activate + // all 32 lanes for the full fused-single band (exactly 32 chunks at 512 + // experts, matching the guard) and halve the per-lane quad chain to 4 at + // that endpoint. Bitwise-identical outputs: chunk c covers experts + // [16c, 16c+16), i.e. half of bitmap word c/2 (bit position preserved); + // int4 stores stay 16B aligned since chunk bases are multiples of 16 + // experts (64B). Revert = set kSingleChunkExperts back to 32. + constexpr int kSingleChunkExperts = 16; + constexpr int kSingleQuads = kSingleChunkExperts / 4; + static_assert(kSingleChunkExperts == 16 || kSingleChunkExperts == 32); + const int nchunks = + (num_experts + kSingleChunkExperts - 1) / kSingleChunkExperts; + const int chunk = lane; // per-lane expert chunk + const int e0 = chunk * kSingleChunkExperts; + const unsigned int w = (chunk < nchunks) ? s_bits[e0 >> 5] : 0u; + const unsigned int mine = (kSingleChunkExperts == 16) + ? ((w >> (e0 & 31)) & 0xFFFFu) + : w; + const int cnt = __popc(mine); + int pre = cnt; + for (int ofs = 1; ofs < 32; ofs <<= 1) { + const int other = __shfl_up_sync(0xFFFFFFFFu, pre, ofs); + if (lane >= ofs) { + pre += other; + } + } + // pre is the inclusive prefix of selected-expert counts over chunks. + const int base_cnt = pre - cnt; // selected experts before this chunk + if (chunk < nchunks) { + const int avail = num_experts - e0; // >= 1 + const int vec_quads = + (avail >> 2) < kSingleQuads ? (avail >> 2) : kSingleQuads; + // qb tracks the selected-count prefix at each quad start (running + // add replaces the per-element popcount). + int qb = base_cnt; +#pragma unroll + for (int q = 0; q < kSingleQuads; ++q) { + const int i = q * 4; + if (i >= avail) { + break; + } + const int b0 = (int)((mine >> i) & 1u); + const int b1 = (int)((mine >> (i + 1)) & 1u); + const int b2 = (int)((mine >> (i + 2)) & 1u); + const int b3 = (int)((mine >> (i + 3)) & 1u); + const int p1 = b0; + const int p2 = b0 + b1; + const int p3 = p2 + b2; + if (q < vec_quads) { + const int4 sel4 = {b0, b1, b2, b3}; + *reinterpret_cast(expert_counts + e0 + i) = sel4; + *reinterpret_cast(expert_scatter_offsets + e0 + i) = + sel4; + const int4 off4 = {qb * block_m, (qb + p1) * block_m, + (qb + p2) * block_m, (qb + p3) * block_m}; + *reinterpret_cast(expert_offsets + e0 + i) = off4; + } else { + const int bits[4] = {b0, b1, b2, b3}; + const int pfx[4] = {0, p1, p2, p3}; +#pragma unroll + for (int t = 0; t < 4; ++t) { + if (i + t >= avail) { + break; + } + expert_counts[e0 + i + t] = bits[t]; + expert_scatter_offsets[e0 + i + t] = bits[t]; + expert_offsets[e0 + i + t] = (qb + pfx[t]) * block_m; + } + } + if (b0) { + s_seg[qb] = e0 + i; + } + if (b1) { + s_seg[qb + p1] = e0 + i + 1; + } + if (b2) { + s_seg[qb + p2] = e0 + i + 2; + } + if (b3) { + s_seg[qb + p3] = e0 + i + 3; + } + qb += p3 + b3; + } + // Tail of the unshifted window: offsets[num_experts] = extent = pre + // (inclusive prefix over all chunks) only on the last chunk lane. + if (chunk == nchunks - 1) { + expert_offsets[num_experts] = extent; + } + } + if (lane == 0) { + num_tokens_post_padded[0] = extent; + } + __syncwarp(0xFFFFFFFFu); + + // Sentinel padding / real route / zero tail, slot-parallel over the warp. + // Every selected segment is exactly block_m wide and holds one route at + // slot 0, so segment/block ids come straight from the segment table. + for (int s = lane; s < slots; s += 32) { + if (s >= extent) { + sorted_token_ids[s] = 0; + continue; + } + const int seg = (int)fdiv_u32((unsigned int)s, fd_bm); + // num_tokens == 1: the padding sentinel value is pairs == top_k. + sorted_token_ids[s] = + (s == seg * block_m) ? s_route[s_seg[seg]] : top_k; + } + + // Per-block expert ids: each nonempty segment spans exactly one block. + for (int b = lane; b < max_blocks; b += 32) { + block_expert_ids[b] = (b < rounds) ? s_seg[b] : 0; + } +} +// NOTE (round 4 experiment, reverted): a radix-histogram threshold selection +// (8-bit prefix refinement levels over the 64-bit route keys + gather + an +// in-register bitonic ordering sort across lanes) was measured ~2x SLOWER +// than the tournament above at the staged shapes (select 6.5 vs 3.66us at +// shape 3, 7.4 vs 3.9us at shape 4): each histogram level serializes 64-bit +// variable shifts, __match_any_sync aggregation, smem atomics and a per-lane +// 8-bin suffix scan (~1us/level on this latency-bound one-warp-per-block +// regime), and randn logits need ~3 levels, versus ~90cy x k serial rounds +// for the tournament. Correctness subtleties solved before reverting: the +// termination test must subtract the current level's above-count from k_rem +// BEFORE comparing against the bucket size, and the universal membership +// predicate is keys[c] >= (prefix << (64 - pl)). + +__device__ __forceinline__ void token_softmax_finish( + float* __restrict__ w_row, int lane, int top_k) { + float sm = -INFINITY; + for (int j = lane; j < top_k; j += 32) { + sm = fmaxf(sm, w_row[j]); + } + for (int ofs = 16; ofs > 0; ofs >>= 1) { + float other = __shfl_xor_sync(0xFFFFFFFFu, sm, ofs); + sm = fmaxf(sm, other); + } + float ssum = 0.0f; + for (int j = lane; j < top_k; j += 32) { + ssum += expf(w_row[j] - sm); + } + for (int ofs = 16; ofs > 0; ofs >>= 1) { + ssum += __shfl_xor_sync(0xFFFFFFFFu, ssum, ofs); + } + const float inv = 1.0f / ssum; + for (int j = lane; j < top_k; j += 32) { + w_row[j] = expf(w_row[j] - sm) * inv; + } +} + +// Exclusive-scan helper: binary search for the expert whose block-aligned +// segment contains sorted-slot v over offs[1..num_experts]. +__device__ __forceinline__ int segment_owner(const int* offs, int num_experts, + int v) { + int lo = 1, hi = num_experts; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (offs[mid] <= v) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo - 1; +} + +// ----------------------- small/medium two-kernel path ---------------------- + +template +__global__ void select_kernel(const float* __restrict__ logits, + float* __restrict__ topk_weights, + int* __restrict__ topk_ids, + int num_tokens, + int num_experts, + int top_k, + int routed_top_k) { + const int token = (int)((blockIdx.x * blockDim.x + threadIdx.x) >> 5); + const int lane = (int)(threadIdx.x & 31); + if (token < num_tokens) { + token_select(logits + (long long)token * num_experts, + topk_weights + (long long)token * top_k, + topk_ids + (long long)token * top_k, lane, + num_experts, top_k, routed_top_k); + } + pdl_trigger(); +} + +// Four-warps-per-token selection. Warp w owns experts [w*32*CPL4, +// (w+1)*32*CPL4): 4 warps cover the full CPL=16 band with CPL4=4 keys per +// lane, one warp per SM sub-partition scheduler. Each round's four local +// winners meet in shared memory behind a single double-buffered block +// barrier (slot j&1): a write to slot p at round j+2 is separated from +// every read of slot p in round j by barrier j+1. Selection math is +// IDENTICAL to the single-warp tournament -- the round max is the max of +// the four local maxima of keys below prev -- so the stable-tie output +// matches bit for bit. Register-softmax regime only (top_k <= 32, rounds +// owned by warp-0 threads). +template +__global__ void __launch_bounds__(128) +select_quad_kernel(const float* __restrict__ logits, + float* __restrict__ topk_weights, + int* __restrict__ topk_ids, + int num_tokens, + int num_experts, + int top_k, + int routed_top_k) { + __shared__ unsigned long long s_win[2][4]; // [parity][warp] + + const int token = (int)blockIdx.x; + const int tid = (int)threadIdx.x; // 0..127 + const int warp = tid >> 5; + const int lane = tid & 31; + const float* row = logits + (long long)token * num_experts; + float* w_row = topk_weights + (long long)token * top_k; + int* id_row = topk_ids + (long long)token * top_k; + const int routed_experts = num_experts - (SHARED ? 1 : 0); + + unsigned long long keys[CPL4]; +#pragma unroll + for (int c = 0; c < CPL4; ++c) { + const int e = warp * (32 * CPL4) + lane + c * 32; + keys[c] = (e < routed_experts) ? route_key(row[e], e) : 0ull; + } + + unsigned long long prev = 0xFFFFFFFFFFFFFFFFull; + float v_acc = -INFINITY; + for (int j = 0; j < routed_top_k; ++j) { + unsigned long long cand[CPL4]; +#pragma unroll + for (int c = 0; c < CPL4; ++c) { + cand[c] = (keys[c] < prev) ? keys[c] : 0ull; + } +#pragma unroll + for (int step = CPL4 / 2; step > 0; step >>= 1) { +#pragma unroll + for (int c = 0; c < step; ++c) { + cand[c] = cand[c] > cand[c + step] ? cand[c] : cand[c + step]; + } + } + const unsigned long long best = cand[0]; + const unsigned int lane_hi = (unsigned int)(best >> 32); + const unsigned int best_hi = __reduce_max_sync(0xFFFFFFFFu, lane_hi); + const unsigned int tag = + (lane_hi == best_hi) ? (unsigned int)best : 0u; + const unsigned int best_lo = __reduce_max_sync(0xFFFFFFFFu, tag); + if (lane == 0) { + s_win[j & 1][warp] = + ((unsigned long long)best_hi << 32) | best_lo; + } + __syncthreads(); + const unsigned long long win0 = s_win[j & 1][0]; + const unsigned long long win1 = s_win[j & 1][1]; + const unsigned long long win2 = s_win[j & 1][2]; + const unsigned long long win3 = s_win[j & 1][3]; + const unsigned long long win01 = win0 > win1 ? win0 : win1; + const unsigned long long win23 = win2 > win3 ? win2 : win3; + const unsigned long long win = win01 > win23 ? win01 : win23; + prev = win; + const unsigned int win_hi = (unsigned int)(win >> 32); + const unsigned int win_lo = (unsigned int)(win & 0xFFFFFFFFu); + // round owner lives in warp 0 (rounds <= 32 in this regime) + if (tid == (j & 31)) { + id_row[j] = (int)(0xFFFFFFFFu - win_lo); + v_acc = float_from_sortable_key(win_hi); + } + } + if (SHARED) { + // rounds <= 32 here, so routed_top_k <= 31 and the owner lands in + // warp 0 with no aliasing of an earlier round. + const int j = routed_top_k; + if (tid == (j & 31)) { + id_row[j] = num_experts - 1; + v_acc = row[num_experts - 1]; + } + } + + // Register-distributed softmax: one round value per warp-0 thread. + if (warp == 0) { + const int rounds = routed_top_k + (SHARED ? 1 : 0); + float m = v_acc; +#pragma unroll + for (int ofs = 16; ofs > 0; ofs >>= 1) { + const float other = __shfl_xor_sync(0xFFFFFFFFu, m, ofs); + m = fmaxf(m, other); + } + const float e = (tid < rounds) ? expf(v_acc - m) : 0.0f; + float sum = e; +#pragma unroll + for (int ofs = 16; ofs > 0; ofs >>= 1) { + sum += __shfl_xor_sync(0xFFFFFFFFu, sum, ofs); + } + if (tid < rounds) { + w_row[tid] = e / sum; + } + } + pdl_trigger(); +} + +// Single block: expert histogram, then block_m-padded exclusive prefix +// scan (the >256-token companion of finish_kernel). +__global__ void __launch_bounds__(kReduceThreads) +reduce_kernel(const int* __restrict__ topk_ids, + int* __restrict__ expert_counts, + int* __restrict__ expert_offsets, + int* __restrict__ num_tokens_post_padded, + int* __restrict__ expert_scatter_offsets, + int num_experts, + int block_m, + int pairs) { + __shared__ int s_hist[kReduceMaxExperts]; + __shared__ int s_wsum[kReduceThreads / 32]; + + const int tid = (int)threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + + for (int e = tid; e < num_experts; e += kReduceThreads) { + s_hist[e] = 0; + } + // Global reads of topk_ids must wait for the predecessor; the smem-zero + // prologue above runs overlapped with it under PDL. + pdl_sync(); + __syncthreads(); + + // One pass over the routes builds the expert histogram. + for (int r = tid; r < pairs; r += kReduceThreads) { + const int e = topk_ids[r]; + atomicAdd(s_hist + e, 1); + } + __syncthreads(); + + // block_m-padded exclusive prefix scan over the histogram. Each thread + // owns one expert index per 512-wide chunk; warp-shuffle scans plus a + // one-warp scan of the warp sums keep it to three barriers per chunk. + int carry = 0; + for (int base = 0; base < num_experts; base += kReduceThreads) { + const int e = base + tid; + int count = 0; + if (e < num_experts) { + count = s_hist[e]; + expert_counts[e] = count; + expert_scatter_offsets[e] = count; + } + int value = ((count + block_m - 1) / block_m) * block_m; + int scan = value; + for (int ofs = 1; ofs < 32; ofs <<= 1) { + const int other = __shfl_up_sync(0xFFFFFFFFu, scan, ofs); + if (lane >= ofs) { + scan += other; + } + } + if (lane == 31) { + s_wsum[warp] = scan; + } + __syncthreads(); + if (warp == 0) { + // All 32 lanes must execute the sync-shuffle; only the first + // kReduceThreads/32 lanes carry meaningful warp sums. + int w = (lane < kReduceThreads / 32) ? s_wsum[lane] : 0; + for (int ofs = 1; ofs < kReduceThreads / 32; ofs <<= 1) { + const int other = __shfl_up_sync(0xFFFFFFFFu, w, ofs); + if (lane >= ofs) { + w += other; + } + } + if (lane < kReduceThreads / 32) { + s_wsum[lane] = w; + } + } + __syncthreads(); + const int warp_base = (warp > 0) ? s_wsum[warp - 1] : 0; + const int inclusive = carry + warp_base + scan; + if (e < num_experts) { + expert_offsets[e + 1] = inclusive; + } + carry += s_wsum[kReduceThreads / 32 - 1]; + // Same last-chunk skip as finish_kernel: no later chunk writes + // s_wsum, and only thread-local carry feeds the loop-exit stores. + if (base + kReduceThreads < num_experts) { + __syncthreads(); + } + } + // carry is uniform across the block after the final chunk and equals the + // total padded extent (works for multi-chunk num_experts > kReduceThreads). + if (tid == 0) { + expert_offsets[0] = 0; + num_tokens_post_padded[0] = carry; + } + pdl_trigger(); +} + +// Parallel tail: one warp per expert. Each warp replays the route stream with +// ballot-scan to scatter its expert's routes in ascending route order, fills +// its sentinel padding, stamps its block ids, and (with the trailing loop) +// writes zeros/zeros for everything at/after the extent. Spreads ~9K output +// stores over the whole GPU instead of one SM. +__global__ void tail_kernel(const int* __restrict__ topk_ids, + const int* __restrict__ expert_counts, + const int* __restrict__ expert_offsets, + int* __restrict__ block_expert_ids, + int* __restrict__ sorted_token_ids, + int num_experts, + int block_m, + int pairs, + int max_blocks, + int slots) { + const int e = (int)blockIdx.x; + const int lane = (int)(threadIdx.x & 31); + pdl_sync(); + const int start_off = expert_offsets[e]; + const int end_off = expert_offsets[e + 1]; + const int count = expert_counts[e]; + + if (count > 0) { + // Scatter: with r scanned in ascending order, the i-th route matching + // this expert lands at start_off + i. + int base = 0; + for (int c0 = 0; c0 < pairs; c0 += 32) { + const int r = c0 + lane; + const bool match = r < pairs && topk_ids[r] == e; + const unsigned int ball = __ballot_sync(0xFFFFFFFFu, match); + if (match) { + const int rank = base + __popc(ball & ((1u << lane) - 1u)); + sorted_token_ids[start_off + rank] = r; + } + base += __popc(ball); + } + for (int j = start_off + count + lane; j < end_off; j += 32) { + sorted_token_ids[j] = pairs; // padding sentinel + } + const int first_block = start_off / block_m; + const int last_block = end_off / block_m; + for (int b = first_block + lane; b < last_block; b += 32) { + block_expert_ids[b] = e; + } + } + + // Zero everything at/after the extent: slots and per-block expert ids. + const int extent = expert_offsets[num_experts]; + const int stride = num_experts * 32; + for (int j = extent + e * 32 + lane; j < slots; j += stride) { + sorted_token_ids[j] = 0; + } + const int used_blocks = extent / block_m; + for (int b = used_blocks + e * 32 + lane; b < max_blocks; b += stride) { + block_expert_ids[b] = 0; + } +} + +// -------------------- two-kernel finish (small/medium path) ---------------- +// +// Replaces the reduce_kernel + scatter_bitmap_kernel pair: every block +// rebuilds the histogram, (expert x token) route bitmap and padded exclusive +// scan redundantly in its own shared memory (a ~1.5us prologue that runs in +// parallel across all blocks), then all blocks grid-stride the scatter and +// the sentinel/zero fills. Block 0 alone stores the histogram/scan outputs. +// This trades a small amount of redundant smem compute for one less kernel +// launch plus one less PDL dependency in the chain, which is worth ~3us on +// the 17..256-token shapes. +template +__global__ void __launch_bounds__(kReduceThreads) +finish_kernel(const int* __restrict__ topk_ids, + int* __restrict__ expert_counts, + int* __restrict__ expert_offsets, + int* __restrict__ num_tokens_post_padded, + int* __restrict__ expert_scatter_offsets, + int* __restrict__ block_expert_ids, + int* __restrict__ sorted_token_ids, + int num_experts, + int top_k, + int block_m, + int pairs, + int max_blocks, + int slots, + int twords, + FastDivU32 fd_topk, + FastDivU32 fd_bm) { + __shared__ int s_hist[kReduceMaxExperts]; + __shared__ __align__(16) unsigned int s_bitmap[kSbBitmapInts]; + __shared__ int s_off[kReduceMaxExperts + 1]; + __shared__ int s_wsum[kReduceThreads / 32]; + __shared__ int s_blk[kSBlkCap]; // block -> expert map (guarded) + // Scan-stamped per-word base slots (launched with + // twords*num_experts*4 dynamic bytes, <= 32KB worst case). + extern __shared__ int s_wslot[]; + + const int tid = (int)threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int gtid = (int)(blockIdx.x * blockDim.x + threadIdx.x); + const int total = (int)(gridDim.x * blockDim.x); + const bool store_outputs = blockIdx.x == 0; + // Word-major bitmap layout: token-word i, expert e lives at + // i * ne_words + e. The scan phase's count popcounts (thread e reads its + // whole column) are then bank-conflict-free for any num_experts; the + // build atomics and scatter rank reads keep their original random-column + // pattern. + const int ne_words = num_experts; + + const int words = num_experts * TW; + // Only the route bitmap needs zeroing: expert counts are recovered in + // the scan phase as bitmap column popcounts (each token selects an + // expert at most once, so the bit count equals the route count), which + // removes one smem atomic per route from the redundant prologue. + // Vectorized when the (16B-aligned) word count is a multiple of 4. + if ((words & 3) == 0) { + uint4* z = reinterpret_cast(s_bitmap); + for (int w = tid; w < (words >> 2); w += kReduceThreads) { + z[w] = make_uint4(0u, 0u, 0u, 0u); + } + } else { + for (int w = tid; w < words; w += kReduceThreads) { + s_bitmap[w] = 0u; + } + } + // Global reads of topk_ids must wait for the predecessor; the smem-zero + // prologue above runs overlapped with it under PDL. + pdl_sync(); + __syncthreads(); + { + // Peel the first two strided routes so both topk_ids loads are + // issued back-to-back and their latencies overlap: pairs <= + // 2*kReduceThreads covers every staged shape, and this prologue is + // latency-bound, not bandwidth-bound. The generic loop below keeps + // every larger input. + const int r0 = tid; + if (r0 < pairs) { + const int r1 = r0 + kReduceThreads; + const bool two = r1 < pairs; + const int e0 = topk_ids[r0]; + const int e1 = two ? topk_ids[r1] : 0; + const int t0 = (int)fdiv_u32((unsigned int)r0, fd_topk); + atomicOr(s_bitmap + (t0 >> 5) * ne_words + e0, 1u << (t0 & 31)); + if (two) { + const int t1 = (int)fdiv_u32((unsigned int)r1, fd_topk); + atomicOr(s_bitmap + (t1 >> 5) * ne_words + e1, + 1u << (t1 & 31)); + } + } + } + for (int r = 2 * kReduceThreads + tid; r < pairs; r += kReduceThreads) { + const int token = (int)fdiv_u32((unsigned int)r, fd_topk); + const int e = topk_ids[r]; + atomicOr(s_bitmap + (token >> 5) * ne_words + e, 1u << (token & 31)); + } + __syncthreads(); + + // block_m-padded exclusive scan (loop over chunks for num_experts > + // kReduceThreads). Offsets stay in smem; only block 0 touches the global + // histogram/scan outputs. Every padded offset is a multiple of block_m, + // so when max_blocks fits kSBlkCap thread e also stamps its own + // [exclusive, inclusive) block span into s_blk (spans are disjoint and + // contiguous), giving the fill phase a 1-load expert lookup instead of a + // 9-level binary search. + const bool use_blkmap = max_blocks <= kSBlkCap; + int carry = 0; + for (int base = 0; base < num_experts; base += kReduceThreads) { + const int e = base + tid; + int count = 0; + int wcnt[TW]; // per-word column popcounts, reused by the slot stamp + if (e < num_experts) { +#pragma unroll + for (int i = 0; i < TW; ++i) { + wcnt[i] = __popc(s_bitmap[i * ne_words + e]); + count += wcnt[i]; + } + s_hist[e] = count; + if (store_outputs) { + expert_counts[e] = count; + expert_scatter_offsets[e] = count; + } + } + const int value = + (int)(fdiv_u32((unsigned int)(count + block_m - 1), fd_bm) * + (unsigned int)block_m); + int scan = value; + for (int ofs = 1; ofs < 32; ofs <<= 1) { + const int other = __shfl_up_sync(0xFFFFFFFFu, scan, ofs); + if (lane >= ofs) { + scan += other; + } + } + if (lane == 31) { + s_wsum[warp] = scan; + } + __syncthreads(); + // Every warp redundantly scans the 16 warp sums in registers and + // shuffles out its own base, removing the round-11 warp-0-only + // serial section and its second barrier (single-chunk shapes then + // need only the one barrier above; the conditional barrier below + // still guards s_wsum against the next chunk's writers). + int wsum = (lane < kReduceThreads / 32) ? s_wsum[lane] : 0; + for (int ofs = 1; ofs < kReduceThreads / 32; ofs <<= 1) { + const int other = __shfl_up_sync(0xFFFFFFFFu, wsum, ofs); + if (lane >= ofs) { + wsum += other; + } + } + const int warp_base = + (warp > 0) ? __shfl_sync(0xFFFFFFFFu, wsum, warp - 1) : 0; + const int inclusive = carry + warp_base + scan; + if (e < num_experts) { + s_off[e + 1] = inclusive; + // Stamp each bitmap word's base slot + // s_off[e] + (# set bits in lower words of this column). + int run = inclusive - value; +#pragma unroll + for (int i = 0; i < TW; ++i) { + s_wslot[i * ne_words + e] = run; + run += wcnt[i]; + } + if (use_blkmap) { + // NOTE: a hoisted endpoint-divide (b0, b1) form measured + // 0.1-0.18us SLOWER than this per-block divide-and-store + // loop on both 512-expert cases; keep the o-loop. + for (int o = inclusive - value; o < inclusive; + o += block_m) { + s_blk[o / block_m] = e; + } + } + if (store_outputs) { + expert_offsets[e + 1] = inclusive; + } + } + carry += __shfl_sync(0xFFFFFFFFu, wsum, kReduceThreads / 32 - 1); + // Only another chunk can clobber s_wsum underneath this chunk's + // readers; the trailing __syncthreads after the loop covers the rest, + // so the single-chunk case (num_experts <= kReduceThreads) skips it. + if (base + kReduceThreads < num_experts) { + __syncthreads(); + } + } + const int extent = carry; // uniform across the block + if (tid == 0) { + s_off[0] = 0; + if (store_outputs) { + expert_offsets[0] = 0; + num_tokens_post_padded[0] = extent; + } + } + __syncthreads(); + + // Scatter: rank inside the expert segment = number of earlier tokens + // selecting the same expert -> ascending flattened route order. + // Route-parallel: the per-route rank is a single table load (stamped by + // the scan phase) plus one bitmask popcount of the route's own word. + for (int r = gtid; r < pairs; r += total) { + const int token = (int)fdiv_u32((unsigned int)r, fd_topk); + const int e = topk_ids[r]; + const int w = (token >> 5) * ne_words + e; + const int rank = + s_wslot[w] + __popc(s_bitmap[w] & ((1u << (token & 31)) - 1u)); + sorted_token_ids[rank] = r; + } + + // Sentinel padding + zero tail + per-block expert ids in one slot-parallel + // pass; the scatter phase owns the in-extent non-pad slots. Every expert + // offset is block_m-aligned, so the enclosing expert for a slot is exactly + // s_blk[s/block_m] (stamped during the scan, ~1 store per expert) when the + // map fits kSBlkCap; larger block counts keep the binary search. The + // slot-parallel form keeps stores fully coalesced; expert-segment- + // parallel variants (thread-per-expert and warp-per-expert) were measured + // 0.6-3us SLOWER here (scattered half-full store transactions). Round 9's + // flat-map attempt built its map with a separate redundant search pass + // and measured ~0.1us slower at 544 blocks; stamping the same map from + // the scan loop itself removes that build cost. + for (int s = gtid; s < slots; s += total) { + const int b = (int)fdiv_u32((unsigned int)s, fd_bm); + const bool block_start = s == b * block_m; + if (s >= extent) { + sorted_token_ids[s] = 0; + if (block_start) { + block_expert_ids[b] = 0; + } + continue; + } + int e; + if (use_blkmap) { + e = s_blk[b]; + } else { + // General path (max_blocks > kSBlkCap): 9-level binary search + // for the enclosing expert segment. + int lo = 1, hi = num_experts; + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (s_off[mid] <= s) { + lo = mid + 1; + } else { + hi = mid; + } + } + e = lo - 1; + } + if (s - s_off[e] >= s_hist[e]) { + sorted_token_ids[s] = pairs; + } + if (block_start) { + block_expert_ids[b] = e; + } + } +} + +// ------------------------- generic large-input path ------------------------ + +__global__ void select_generic_kernel(const float* __restrict__ logits, + float* __restrict__ topk_weights, + int* __restrict__ topk_ids, + int* __restrict__ expert_counts, + unsigned int* __restrict__ route_bits, + int num_tokens, + int num_experts, + int top_k, + int routed_top_k, + int has_shared, + int words_per_token) { + const int token = (int)((blockIdx.x * blockDim.x + threadIdx.x) >> 5); + const int lane = (int)(threadIdx.x & 31); + if (token >= num_tokens) { + return; + } + + const float* row = logits + (long long)token * num_experts; + unsigned int* bits_row = route_bits + (long long)token * words_per_token; + for (int i = lane; i < words_per_token; i += 32) { + bits_row[i] = 0u; + } + __syncwarp(); + + const int routed_experts = num_experts - has_shared; + float* w_row = topk_weights + (long long)token * top_k; + int* id_row = topk_ids + (long long)token * top_k; + + unsigned long long prev = 0xFFFFFFFFFFFFFFFFull; + for (int j = 0; j < routed_top_k; ++j) { + unsigned long long best = 0ull; + for (int e = lane; e < routed_experts; e += 32) { + unsigned long long key = route_key(row[e], e); + if (key < prev && key > best) { + best = key; + } + } + for (int ofs = 16; ofs > 0; ofs >>= 1) { + unsigned long long other = __shfl_xor_sync(0xFFFFFFFFu, best, ofs); + best = other > best ? other : best; + } + prev = best; + const int expert = + (int)(0xFFFFFFFFu - (unsigned int)(best & 0xFFFFFFFFull)); + if (lane == (j & 31)) { + id_row[j] = expert; + w_row[j] = row[expert]; + atomicAdd(expert_counts + expert, 1); + atomicOr(bits_row + (expert >> 5), 1u << (expert & 31)); + } + } + if (has_shared) { + const int expert = num_experts - 1; + const int j = routed_top_k; + if (lane == (j & 31)) { + id_row[j] = expert; + w_row[j] = row[expert]; + atomicAdd(expert_counts + expert, 1); + atomicOr(bits_row + (expert >> 5), 1u << (expert & 31)); + } + } + + token_softmax_finish(w_row, lane, top_k); +} + +__global__ void scan_kernel(const int* __restrict__ expert_counts, + int* __restrict__ expert_offsets, + int* __restrict__ num_tokens_post_padded, + int* __restrict__ expert_scatter_offsets, + int* __restrict__ block_expert_ids, + int num_experts, + int block_m, + int max_blocks) { + __shared__ int tile[kScanThreads]; + const int tid = (int)threadIdx.x; + + int carry = 0; + for (int base = 0; base < num_experts; base += kScanThreads) { + const int e = base + tid; + int count = 0; + if (e < num_experts) { + count = expert_counts[e]; + expert_scatter_offsets[e] = count; + } + const int padded = ((count + block_m - 1) / block_m) * block_m; + tile[tid] = padded; + __syncthreads(); + for (int ofs = 1; ofs < kScanThreads; ofs <<= 1) { + int v = (tid >= ofs) ? tile[tid - ofs] : 0; + __syncthreads(); + tile[tid] += v; + __syncthreads(); + } + if (e < num_experts) { + expert_offsets[e + 1] = carry + tile[tid]; + } + carry += tile[kScanThreads - 1]; + __syncthreads(); + } + if (tid == 0) { + expert_offsets[0] = 0; + num_tokens_post_padded[0] = carry; + } + + const int extent = carry; // uniform across the block + for (int b = tid; b < max_blocks; b += kScanThreads) { + const int v = b * block_m; + block_expert_ids[b] = + (v < extent) ? segment_owner(expert_offsets, num_experts, v) : 0; + } +} + +__global__ void scatter_kernel(const int* __restrict__ topk_ids, + const int* __restrict__ expert_counts, + const int* __restrict__ expert_offsets, + const unsigned int* __restrict__ route_bits, + int* __restrict__ sorted_token_ids, + int num_tokens, + int top_k, + int block_m, + int num_experts, + int words_per_token, + int pairs, + int slots) { + const int r = (int)(blockIdx.x * blockDim.x + threadIdx.x); + + if (r < pairs) { + const int token = r / top_k; + const int expert = topk_ids[r]; + const unsigned int mask = 1u << (expert & 31); + const unsigned int* col = route_bits + (expert >> 5); + int rank = 0; + for (int t = 0; t < token; ++t) { + rank += (col[(long long)t * words_per_token] & mask) ? 1 : 0; + } + sorted_token_ids[expert_offsets[expert] + rank] = r; + } + + if (r < slots) { + const int extent = expert_offsets[num_experts]; + if (r >= extent) { + sorted_token_ids[r] = 0; + } else { + const int expert = segment_owner(expert_offsets, num_experts, r); + if (r - expert_offsets[expert] >= expert_counts[expert]) { + sorted_token_ids[r] = pairs; // padding sentinel + } + } + } +} + +template +void launch_fused(const float* logits_ptr, + float* weights_ptr, + int* ids_ptr, + int* counts_ptr, + int* offsets_ptr, + int* extent_ptr, + int* scatter_ptr, + int* blocks_ptr, + int* sorted_ptr, + int num_tokens, + int num_experts, + int top_k, + int block_m, + int routed_top_k, + int has_shared, + int pairs, + int max_blocks, + int slots, + int twords, + cudaStream_t stream) { + // Exact magic-division pairs for the two constructor-constant divisors + // (built on the host once per forward; pure CPU math, no input + // dependence). + const FastDivU32 fd_topk = make_fast_div_u32((unsigned int)top_k); + const FastDivU32 fd_bm = make_fast_div_u32((unsigned int)block_m); + // One token: every other warp in the block would be idle and the phase + // barriers dominate; a single warp runs the whole pipeline with + // __syncwarp only. top_k <= 32 keeps one selection round per lane. + if (num_tokens == 1 && top_k <= 32) { + if (has_shared) { + fused_single_kernel<<<1, 32, 0, stream>>>( + logits_ptr, weights_ptr, ids_ptr, counts_ptr, + offsets_ptr, extent_ptr, scatter_ptr, blocks_ptr, sorted_ptr, + num_experts, top_k, block_m, routed_top_k, max_blocks, slots, + fd_bm); + } else { + fused_single_kernel<<<1, 32, 0, stream>>>( + logits_ptr, weights_ptr, ids_ptr, counts_ptr, + offsets_ptr, extent_ptr, scatter_ptr, blocks_ptr, sorted_ptr, + num_experts, top_k, block_m, routed_top_k, max_blocks, slots, + fd_bm); + } + return; + } + const bool popc_ok = num_tokens <= block_m; + // (Round-31 experiment -- REVERTED on measurement: widening the fused + // path to 17..32 tokens with a 1024-thread instantiation measured 15.5us + // on (32,512,8,16) vs the 7.0us select+finish pair; launch_bounds(1024) + // caps the selector at ~55 regs/thread and the per-lane key/id + // tournament spills. The multi-kernel dispatch for >16 tokens stays.) + if (has_shared) { + if (popc_ok) { + fused_small_kernel<<<1, kFusedThreads, 0, stream>>>( + logits_ptr, weights_ptr, ids_ptr, counts_ptr, + offsets_ptr, extent_ptr, scatter_ptr, blocks_ptr, sorted_ptr, + num_tokens, num_experts, top_k, block_m, routed_top_k, pairs, + max_blocks, slots, twords, fd_topk, fd_bm); + } else { + fused_small_kernel<<<1, kFusedThreads, 0, + stream>>>( + logits_ptr, weights_ptr, ids_ptr, counts_ptr, + offsets_ptr, extent_ptr, scatter_ptr, blocks_ptr, sorted_ptr, + num_tokens, num_experts, top_k, block_m, routed_top_k, pairs, + max_blocks, slots, twords, fd_topk, fd_bm); + } + } else { + if (popc_ok) { + fused_small_kernel<<<1, kFusedThreads, 0, + stream>>>( + logits_ptr, weights_ptr, ids_ptr, counts_ptr, + offsets_ptr, extent_ptr, scatter_ptr, blocks_ptr, sorted_ptr, + num_tokens, num_experts, top_k, block_m, routed_top_k, pairs, + max_blocks, slots, twords, fd_topk, fd_bm); + } else { + fused_small_kernel<<<1, kFusedThreads, 0, + stream>>>( + logits_ptr, weights_ptr, ids_ptr, counts_ptr, + offsets_ptr, extent_ptr, scatter_ptr, blocks_ptr, sorted_ptr, + num_tokens, num_experts, top_k, block_m, routed_top_k, pairs, + max_blocks, slots, twords, fd_topk, fd_bm); + } + } +} + +template +void launch_small(const float* logits_ptr, + float* weights_ptr, + int* ids_ptr, + int* counts_ptr, + int* offsets_ptr, + int* extent_ptr, + int* scatter_ptr, + int* blocks_ptr, + int* sorted_ptr, + bool use_finish, + int num_tokens, + int num_experts, + int top_k, + int block_m, + int routed_top_k, + int has_shared, + int words_per_token, + int pairs, + int max_blocks, + int slots, + int twords, + cudaStream_t stream) { + const int select_blocks = + (int)((num_tokens * 32 + kSelectThreads - 1) / kSelectThreads); + // CPL=16 expert band with register-softmax top_k: four warps share + // one token's tournament (select_quad_kernel). Everything else keeps + // warp-per-token selection. + const bool split_ok = CPL == 16 && top_k <= 32 && num_tokens > 0; + if (split_ok) { + if (has_shared) { + select_quad_kernel<4, true><<>>( + logits_ptr, weights_ptr, ids_ptr, num_tokens, + num_experts, top_k, routed_top_k); + } else { + select_quad_kernel<4, false><<>>( + logits_ptr, weights_ptr, ids_ptr, num_tokens, + num_experts, top_k, routed_top_k); + } + } else if (has_shared) { + select_kernel<<>>( + logits_ptr, weights_ptr, ids_ptr, num_tokens, num_experts, + top_k, routed_top_k); + } else { + select_kernel<<>>( + logits_ptr, weights_ptr, ids_ptr, num_tokens, num_experts, + top_k, routed_top_k); + } + // NOTE: a single-block "finish everything" kernel was measured strictly + // slower than the reduce+tail/scatter pair at every staged shape across + // three independent implementations (rounds 2-3, up to 17us at shape 3): + // serial warp sections dominate once grid parallelism is removed. The + // finish_kernel below keeps grid parallelism (every block redundantly + // rebuilds the smem histogram/bitmap/scan, then grid-strides the fills) + // and replaces reduce+scatter with a single launch chain link. + if (use_finish) { + // finish_kernel's static smem is ~45KB, so the word-slot dynamic + // table needs the >48KB opt-in exactly once per process. + static const bool smem_opt_in = [] { + constexpr int kDynMax = kSbMaxTwords * kReduceMaxExperts * 4; + for (auto kern : + {finish_kernel<1>, finish_kernel<2>, finish_kernel<4>, + finish_kernel<8>}) { + const cudaError_t err = cudaFuncSetAttribute( + kern, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynMax); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("alphamoe_router: smem opt-in failed: ") + + cudaGetErrorString(err)); + } + } + return true; + }(); + (void)smem_opt_in; + // Measured grid sweep: finish_kernel plateaus at ~16 blocks (the + // redundant prologue runs in parallel and the fills are strip-thin); + // keep a floor of 16 and cap at 64 to bound redundant traffic. + const int work = pairs > slots ? pairs : slots; + int blocks = (work + kReduceThreads - 1) / kReduceThreads; + if (blocks < 16) { + blocks = 16; + } + if (blocks > 64) { + blocks = 64; + } + // finish is templated on the route-bitmap word count so the popcount + // rank/count loops unroll and drop their dynamic bounds checks. + // Exact magic-division pairs for the two constructor-constant + // divisors; built on the host once per forward (pure CPU math, no + // input dependence). + const FastDivU32 fd_topk = make_fast_div_u32((unsigned int)top_k); + const FastDivU32 fd_bm = make_fast_div_u32((unsigned int)block_m); + switch (twords) { + case 1: + launch_pdl(finish_kernel<1>, blocks, kReduceThreads, + (size_t)1 * num_experts * 4, stream, + ids_ptr, counts_ptr, offsets_ptr, extent_ptr, + scatter_ptr, blocks_ptr, sorted_ptr, num_experts, top_k, + block_m, pairs, max_blocks, slots, twords, fd_topk, + fd_bm); + break; + case 2: + launch_pdl(finish_kernel<2>, blocks, kReduceThreads, + (size_t)2 * num_experts * 4, stream, + ids_ptr, counts_ptr, offsets_ptr, extent_ptr, + scatter_ptr, blocks_ptr, sorted_ptr, num_experts, top_k, + block_m, pairs, max_blocks, slots, twords, fd_topk, + fd_bm); + break; + case 3: + case 4: + launch_pdl(finish_kernel<4>, blocks, kReduceThreads, + (size_t)4 * num_experts * 4, stream, + ids_ptr, counts_ptr, offsets_ptr, extent_ptr, + scatter_ptr, blocks_ptr, sorted_ptr, num_experts, top_k, + block_m, pairs, max_blocks, slots, twords, fd_topk, + fd_bm); + break; + default: + launch_pdl(finish_kernel<8>, blocks, kReduceThreads, + (size_t)8 * num_experts * 4, stream, + ids_ptr, counts_ptr, offsets_ptr, extent_ptr, + scatter_ptr, blocks_ptr, sorted_ptr, num_experts, top_k, + block_m, pairs, max_blocks, slots, twords, fd_topk, + fd_bm); + break; + } + } else { + launch_pdl(reduce_kernel, 1, kReduceThreads, 0, stream, ids_ptr, + counts_ptr, offsets_ptr, extent_ptr, scatter_ptr, + num_experts, block_m, pairs); + launch_pdl(tail_kernel, num_experts, 32, 0, stream, ids_ptr, + counts_ptr, offsets_ptr, blocks_ptr, sorted_ptr, + num_experts, block_m, pairs, max_blocks, slots); + } +} + +} // namespace flashinfer::fused_moe + +namespace flashinfer::fused_moe { + +// --------------------------------------------------------------------------- +// Host-side geometry + launch interface (raw pointers, caller-allocated). +// --------------------------------------------------------------------------- + +struct AlphaMoeRouterParams { + int num_tokens = 0; + int num_experts = 0; + int top_k = 0; + int block_m = 0; + int has_shared = 0; + int routed_top_k = 0; + int routed_experts = 0; + int64_t pairs = 0; + int64_t max_blocks = 0; + int64_t slots = 0; + int64_t words_per_token = 0; + int64_t twords = 0; +}; + +// Pure arithmetic; every field depends only on the shape configuration, so +// the geometry is fixed under CUDA graph capture for a fixed captured shape. +inline AlphaMoeRouterParams make_alphamoe_router_params(int num_tokens, + int num_experts, + int top_k, + int block_m, + int has_shared) { + if (num_experts < 2) { + throw std::invalid_argument("alphamoe_router: num_experts must be >= 2"); + } + if (top_k < 1) { + throw std::invalid_argument("alphamoe_router: top_k must be >= 1"); + } + if (block_m < 1) { + throw std::invalid_argument("alphamoe_router: block_m must be >= 1"); + } + if (num_tokens < 1) { + throw std::invalid_argument("alphamoe_router: num_tokens must be >= 1"); + } + AlphaMoeRouterParams p; + p.num_tokens = num_tokens; + p.num_experts = num_experts; + p.top_k = top_k; + p.block_m = block_m; + p.has_shared = has_shared ? 1 : 0; + p.routed_top_k = top_k - p.has_shared; + p.routed_experts = num_experts - p.has_shared; + if (p.routed_top_k < 0 || p.routed_top_k > p.routed_experts) { + throw std::invalid_argument( + "alphamoe_router: invalid top_k/num_experts for shared-expert " + "configuration"); + } + p.pairs = static_cast(num_tokens) * top_k; + const int64_t nonempty = static_cast(num_experts) < p.pairs + ? num_experts + : p.pairs; + p.max_blocks = nonempty + (p.pairs - nonempty) / block_m; + p.slots = p.max_blocks * block_m; + p.words_per_token = (static_cast(num_experts) + 31) / 32; + p.twords = (static_cast(num_tokens) + 31) / 32; + return p; +} + +// Number of int32 scratch elements the generic path needs (zero otherwise). +inline int64_t alphamoe_router_scratch_ints(const AlphaMoeRouterParams& p) { + return p.num_experts > kReduceMaxExperts ? p.num_tokens * p.words_per_token + : 0; +} + +// Launch the routing pipeline. All outputs must be sized per +// make_alphamoe_router_params. `scratch` must hold +// alphamoe_router_scratch_ints(p) int32 elements when that is non-zero. +inline void alphamoe_router_forward(const AlphaMoeRouterParams& p, + const float* logits_ptr, + float* topk_weights_ptr, int* topk_ids_ptr, + int* expert_counts_ptr, + int* expert_offsets_ptr, + int* expert_scatter_offsets_ptr, + int* num_tokens_post_padded_ptr, + int* block_expert_ids_ptr, + int* sorted_token_ids_ptr, int* scratch_ptr, + cudaStream_t stream) { + const int num_tokens = p.num_tokens; + const int num_experts = p.num_experts; + const int top_k = p.top_k; + const int block_m = p.block_m; + const int has_shared = p.has_shared; + const int routed_top_k = p.routed_top_k; + const int pairs = static_cast(p.pairs); + const int max_blocks = static_cast(p.max_blocks); + const int slots = static_cast(p.slots); + const int twords = static_cast(p.twords); + const int words_per_token = static_cast(p.words_per_token); + + const bool small_ok = num_experts <= kReduceMaxExperts; + const int64_t per_lane = + (static_cast(p.routed_experts) + 31) / 32; + + if (small_ok) { + const bool fused_ok = num_experts <= kFusedMaxExperts && + pairs <= kFusedPairsCap && + num_tokens <= kFusedThreads / 32; + if (fused_ok) { + if (per_lane <= 4) { + launch_fused<4>(logits_ptr, topk_weights_ptr, topk_ids_ptr, + expert_counts_ptr, expert_offsets_ptr, + num_tokens_post_padded_ptr, + expert_scatter_offsets_ptr, block_expert_ids_ptr, + sorted_token_ids_ptr, num_tokens, num_experts, top_k, + block_m, routed_top_k, has_shared, pairs, max_blocks, + slots, twords, stream); + } else if (per_lane <= 8) { + launch_fused<8>(logits_ptr, topk_weights_ptr, topk_ids_ptr, + expert_counts_ptr, expert_offsets_ptr, + num_tokens_post_padded_ptr, + expert_scatter_offsets_ptr, block_expert_ids_ptr, + sorted_token_ids_ptr, num_tokens, num_experts, top_k, + block_m, routed_top_k, has_shared, pairs, max_blocks, + slots, twords, stream); + } else { + launch_fused<16>(logits_ptr, topk_weights_ptr, topk_ids_ptr, + expert_counts_ptr, expert_offsets_ptr, + num_tokens_post_padded_ptr, + expert_scatter_offsets_ptr, block_expert_ids_ptr, + sorted_token_ids_ptr, num_tokens, num_experts, + top_k, block_m, routed_top_k, has_shared, pairs, + max_blocks, slots, twords, stream); + } + return; + } + const bool use_finish = twords <= kSbMaxTwords; + if (per_lane <= 4) { + launch_small<4>(logits_ptr, topk_weights_ptr, topk_ids_ptr, + expert_counts_ptr, expert_offsets_ptr, + num_tokens_post_padded_ptr, expert_scatter_offsets_ptr, + block_expert_ids_ptr, sorted_token_ids_ptr, use_finish, + num_tokens, num_experts, top_k, block_m, routed_top_k, + has_shared, words_per_token, pairs, max_blocks, slots, + twords, stream); + } else if (per_lane <= 8) { + launch_small<8>(logits_ptr, topk_weights_ptr, topk_ids_ptr, + expert_counts_ptr, expert_offsets_ptr, + num_tokens_post_padded_ptr, expert_scatter_offsets_ptr, + block_expert_ids_ptr, sorted_token_ids_ptr, use_finish, + num_tokens, num_experts, top_k, block_m, routed_top_k, + has_shared, words_per_token, pairs, max_blocks, slots, + twords, stream); + } else if (per_lane <= 16) { + launch_small<16>(logits_ptr, topk_weights_ptr, topk_ids_ptr, + expert_counts_ptr, expert_offsets_ptr, + num_tokens_post_padded_ptr, expert_scatter_offsets_ptr, + block_expert_ids_ptr, sorted_token_ids_ptr, use_finish, + num_tokens, num_experts, top_k, block_m, routed_top_k, + has_shared, words_per_token, pairs, max_blocks, slots, + twords, stream); + } else { + launch_small<32>(logits_ptr, topk_weights_ptr, topk_ids_ptr, + expert_counts_ptr, expert_offsets_ptr, + num_tokens_post_padded_ptr, expert_scatter_offsets_ptr, + block_expert_ids_ptr, sorted_token_ids_ptr, use_finish, + num_tokens, num_experts, top_k, block_m, routed_top_k, + has_shared, words_per_token, pairs, max_blocks, slots, + twords, stream); + } + return; + } + + // Generic path: grid-parallel select with global histogram atomics + + // single-block scan + grid scatter over the route bitmap in scratch. + unsigned int* bits_ptr = reinterpret_cast(scratch_ptr); + cudaMemsetAsync(expert_counts_ptr, 0, sizeof(int) * num_experts, stream); + + const int select_blocks = + static_cast((num_tokens * 32 + kSelectThreads - 1) / + kSelectThreads); + select_generic_kernel<<>>( + logits_ptr, topk_weights_ptr, topk_ids_ptr, expert_counts_ptr, bits_ptr, + num_tokens, num_experts, top_k, routed_top_k, has_shared, + words_per_token); + + scan_kernel<<<1, kScanThreads, 0, stream>>>( + expert_counts_ptr, expert_offsets_ptr, num_tokens_post_padded_ptr, + expert_scatter_offsets_ptr, block_expert_ids_ptr, num_experts, block_m, + max_blocks); + + const int64_t work = pairs > slots ? pairs : slots; + const int scatter_blocks = + static_cast((work + kScatterThreads - 1) / kScatterThreads); + scatter_kernel<<>>( + topk_ids_ptr, expert_counts_ptr, expert_offsets_ptr, bits_ptr, + sorted_token_ids_ptr, num_tokens, top_k, block_m, num_experts, + words_per_token, pairs, slots); +} + +} // namespace flashinfer::fused_moe diff --git a/tests/moe/test_alphamoe_fused_router.py b/tests/moe/test_alphamoe_fused_router.py new file mode 100644 index 00000000000..514775e9b4d --- /dev/null +++ b/tests/moe/test_alphamoe_fused_router.py @@ -0,0 +1,263 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Tests for the fused AlphaMoE gating router (``alphamoe_fused_router`` + +``allocate_alphamoe_route_plan``, "vibecuda" backend). + +The oracle recomputes the complete block-sparse routing metadata bundle in +plain torch: stable descending top-k over the routed experts (ties keep the +lower expert index), optional shared-expert column, fp32 max-subtracted +softmax over the selected logits, expert histogram, block_m-aligned padded +offsets and extent, scatter offsets equal to the expert counts, expert- +grouped flat route ids with sentinel padding, and per-block expert ids. +""" + +import pytest +import torch + +from flashinfer.fused_moe import ( + AlphaMoeRoutePlan, + allocate_alphamoe_route_plan, + alphamoe_fused_router, +) +from flashinfer.utils import BackendSupportedError + + +def _router_reference(logits, top_k, block_m, has_shared_expert): + """Exact torch oracle for the routing bundle (see module docstring).""" + num_tokens, num_experts = logits.shape + routed_experts = num_experts - int(has_shared_expert) + routed_top_k = top_k - int(has_shared_expert) + order = torch.argsort( + logits[:, :routed_experts], dim=-1, descending=True, stable=True + )[:, :routed_top_k] + selected = torch.gather(logits, 1, order) + if has_shared_expert: + shared = torch.full( + (num_tokens, 1), + num_experts - 1, + dtype=torch.int64, + device=logits.device, + ) + order = torch.cat((order, shared), dim=-1) + selected = torch.cat((selected, logits[:, -1:]), dim=-1) + topk_ids = order.to(torch.int32) + topk_weights = torch.softmax(selected, dim=-1) + + flat = topk_ids.cpu().reshape(-1).to(torch.int64) + counts = torch.bincount(flat, minlength=num_experts).to(torch.int32) + padded = (counts + block_m - 1) // block_m * block_m + offsets = torch.empty(num_experts + 1, dtype=torch.int32) + offsets[0] = 0 + offsets[1:] = torch.cumsum(padded, dim=0) + extent = int(offsets[-1]) + + pairs = num_tokens * top_k + nonempty = min(num_experts, pairs) + max_blocks = nonempty + (pairs - nonempty) // block_m + sorted_ids = torch.zeros(max_blocks * block_m, dtype=torch.int32) + expert_ids = torch.zeros(max_blocks, dtype=torch.int32) + sentinel = pairs + for expert in range(num_experts): + start = int(offsets[expert]) + count = int(counts[expert]) + end = int(offsets[expert + 1]) + if count == 0: + continue + routes = torch.nonzero(flat == expert).flatten().to(torch.int32) + sorted_ids[start : start + count] = routes + sorted_ids[start + count : end] = sentinel + expert_ids[start // block_m : end // block_m] = expert + + return ( + topk_weights.cpu(), + topk_ids.cpu(), + sorted_ids, + expert_ids, + torch.tensor([extent], dtype=torch.int32), + counts, + offsets, + counts.clone(), + ) + + +def _max_blocks(num_tokens, num_experts, top_k, block_m): + pairs = num_tokens * top_k + nonempty = min(num_experts, pairs) + return nonempty + (pairs - nonempty) // block_m + + +def _assert_bundle(out, ref): + names = ( + "topk_weights", + "topk_ids", + "sorted_token_ids", + "expert_ids", + "num_tokens_post_padded", + "expert_counts", + "expert_offsets", + "expert_scatter_offsets", + ) + assert len(out) == len(names) + for name, actual, expected in zip(names, out, ref): + if name == "topk_weights": + torch.testing.assert_close( + actual.cpu(), expected, rtol=3e-4, atol=3e-4, msg=name + ) + else: + torch.testing.assert_close( + actual.cpu(), expected, rtol=0, atol=0, msg=name + ) + + +@pytest.mark.parametrize( + "num_tokens,num_experts,top_k,block_m,has_shared_expert", + [ + (1, 512, 2, 16, True), + (8, 32, 4, 8, False), + (8, 257, 9, 8, True), + (32, 512, 8, 16, False), + (128, 512, 8, 16, False), + (33, 65, 5, 4, True), + (256, 128, 8, 32, False), + ], +) +def test_alphamoe_fused_router_correctness( + num_tokens, num_experts, top_k, block_m, has_shared_expert +): + torch.manual_seed(num_tokens * 1000 + num_experts + top_k) + logits = torch.randn(num_tokens, num_experts, dtype=torch.float32, device="cuda") + ref = _router_reference(logits, top_k, block_m, has_shared_expert) + + # Fresh-allocation path. + out = alphamoe_fused_router( + logits, + top_k=top_k, + block_m=block_m, + has_shared_expert=has_shared_expert, + ) + _assert_bundle(out, ref) + + # Plan path: buffers preallocated once, refilled per call. + plan = allocate_alphamoe_route_plan( + logits, top_k=top_k, block_m=block_m, has_shared_expert=has_shared_expert + ) + out = alphamoe_fused_router(logits, plan) + _assert_bundle(out, ref) + + +@pytest.mark.parametrize("has_shared_expert", [False, True]) +def test_alphamoe_fused_router_stable_ties(has_shared_expert): + """Equal logits must keep the lower expert index first (stable order).""" + torch.manual_seed(7) + num_tokens, num_experts, top_k, block_m = 16, 64, 8, 8 + # Small integer alphabet forces many exact ties per row. + logits = torch.randint( + -2, 3, (num_tokens, num_experts), dtype=torch.float32, device="cuda" + ) + ref = _router_reference(logits, top_k, block_m, has_shared_expert) + out = alphamoe_fused_router( + logits, top_k=top_k, block_m=block_m, has_shared_expert=has_shared_expert + ) + _assert_bundle(out, ref) + + +def test_alphamoe_route_plan_contract(): + num_tokens, num_experts, top_k, block_m = 8, 32, 4, 8 + logits = torch.randn(num_tokens, num_experts, dtype=torch.float32, device="cuda") + plan = allocate_alphamoe_route_plan(logits, top_k=top_k, block_m=block_m) + assert isinstance(plan, AlphaMoeRoutePlan) + max_blocks = _max_blocks(num_tokens, num_experts, top_k, block_m) + assert plan.topk_weights.shape == (num_tokens, top_k) + assert plan.topk_ids.shape == (num_tokens, top_k) + assert plan.sorted_token_ids.shape == (max_blocks * block_m,) + assert plan.expert_ids.shape == (max_blocks,) + assert plan.num_tokens_post_padded.shape == (1,) + assert plan.expert_counts.shape == (num_experts,) + assert plan.expert_offsets.shape == (num_experts + 1,) + assert plan.expert_scatter_offsets.shape == (num_experts,) + # Tuple emulation covers the whole public bundle in canonical order. + assert len(plan) == 8 + names = ( + "topk_weights", + "topk_ids", + "sorted_token_ids", + "expert_ids", + "num_tokens_post_padded", + "expert_counts", + "expert_offsets", + "expert_scatter_offsets", + ) + for tensor, name in zip(plan, names): + assert tensor is getattr(plan, name) + # A plan writes into its own persistent buffers across calls. + out1 = alphamoe_fused_router(logits, plan) + assert out1[0] is plan.topk_weights + logits2 = torch.randn_like(logits) + out2 = alphamoe_fused_router(logits2, plan) + assert out2[0] is plan.topk_weights + ref2 = _router_reference(logits2, top_k, block_m, False) + _assert_bundle(out2, ref2) + + +def test_alphamoe_fused_router_validation(): + logits = torch.randn(8, 32, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError): + alphamoe_fused_router(logits) # missing top_k/block_m without a plan + with pytest.raises(ValueError): + alphamoe_fused_router(logits.cpu(), top_k=4, block_m=8) # CPU input + with pytest.raises(ValueError): + alphamoe_fused_router( + logits.to(torch.float16), top_k=4, block_m=8 + ) # non-fp32 input + plan = allocate_alphamoe_route_plan(logits, top_k=4, block_m=8) + mismatched = torch.randn(16, 64, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError): + alphamoe_fused_router(mismatched, plan) # geometry mismatch + with pytest.raises(ValueError): + alphamoe_fused_router(logits, top_k=0, block_m=8) + with pytest.raises(BackendSupportedError): + alphamoe_fused_router(logits, top_k=4, block_m=8, has_shared_expert=False, + backend="tensorrt_llm") + + +def test_alphamoe_fused_router_cuda_graph_replay(): + """Pinned upstream lifecycle: allocate the route plan before warmup and + capture, capture the routing call, change logits in place at the same + address, replay, and validate the complete route plan for the new values. + No replanning and no metadata changes.""" + num_tokens, num_experts, top_k, block_m = 8, 32, 4, 8 + torch.manual_seed(1234) + logits = torch.randn(num_tokens, num_experts, dtype=torch.float32, device="cuda") + plan = allocate_alphamoe_route_plan(logits, top_k=top_k, block_m=block_m) + + out = alphamoe_fused_router(logits, plan) + torch.cuda.synchronize() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + out = alphamoe_fused_router(logits, plan) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = alphamoe_fused_router(logits, plan) + stream.synchronize() + + # Runtime data changes without replanning: new values, same addresses. + logits.copy_(torch.randn_like(logits)) + graph.replay() + torch.cuda.synchronize() + ref = _router_reference(logits, top_k, block_m, False) + _assert_bundle(out, ref) diff --git a/tests/trace/example.py b/tests/trace/example.py index 6820994fa5b..e275e6d5813 100644 --- a/tests/trace/example.py +++ b/tests/trace/example.py @@ -14,6 +14,8 @@ Results: - We would get these example json files under fi_trace_out directory: +alphamoe_fused_router_e256_k8_b16.json +alphamoe_fused_router_e257_k9_b8.json bmm_mxfp8_N128_K128.json fused_add_rmsnorm_h5120.json fused_add_rmsnorm_quant_h7168.json @@ -1074,6 +1076,24 @@ ) +# ── Fused AlphaMoE gating router ("vibecuda" backend, SM100+) ──────────────── +# Stable descending top-k + softmax + block-sparse routing metadata bundle in +# one fused call. Two geometries: plain top-k-8 and shared-expert top-k-9. +with contextlib.suppress(Exception): + flashinfer.fused_moe.alphamoe_fused_router( + torch.randn(128, 256, dtype=torch.float32, device=device), + top_k=8, + block_m=16, + ) +with contextlib.suppress(Exception): + flashinfer.fused_moe.alphamoe_fused_router( + torch.randn(8, 257, dtype=torch.float32, device=device), + top_k=9, + block_m=8, + has_shared_expert=True, + ) + + # ── Routed MoE FP8 per-tensor scale (packed expert ids + weights) ──────────── _routed_topk_ids = ( torch.arange(T_moe * 8, dtype=torch.int32, device=device).reshape(T_moe, 8) % E_loc diff --git a/tests/trace/fi_trace_out/alphamoe_fused_router_e256_k8_b16.json b/tests/trace/fi_trace_out/alphamoe_fused_router_e256_k8_b16.json new file mode 100644 index 00000000000..951af729192 --- /dev/null +++ b/tests/trace/fi_trace_out/alphamoe_fused_router_e256_k8_b16.json @@ -0,0 +1,120 @@ +{ + "name": "alphamoe_fused_router_e256_k8_b16", + "description": "Fused AlphaMoE gating router: stable descending top-k with an optional shared-expert column -> fp32 max-subtracted softmax -> block_m-aligned routing metadata bundle (expert histogram, padded expert offsets + extent, scatter offsets equal to the expert counts, expert-grouped flat route ids with sentinel padding, per-block expert ids). Route ids are flat token*top_k+slot indices; in-segment padding slots carry the sentinel num_tokens*top_k.", + "op_type": "moe_routing", + "tags": [ + "fi_api:flashinfer.fused_moe.alphamoe_router.alphamoe_fused_router", + "status:verified", + "moe", + "moe:routing" + ], + "axes": { + "num_tokens": { + "type": "var" + }, + "num_experts": { + "type": "const", + "value": 256 + }, + "top_k": { + "type": "const", + "value": 8 + }, + "block_m": { + "type": "const", + "value": 16 + }, + "max_blocks": { + "type": "var", + "description": "Block-bound from the output geometry (min(num_experts, pairs) + (pairs - nonempty) // block_m)." + }, + "slots": { + "type": "var", + "description": "max_blocks * block_m." + }, + "one": { + "type": "var", + "description": "Placeholder for shape [1] output tensors." + }, + "num_experts_plus_one": { + "type": "var", + "description": "num_experts + 1 for the inclusive padded offsets." + } + }, + "inputs": { + "router_logits": { + "shape": [ + "num_tokens", + "num_experts" + ], + "dtype": "float32" + }, + "top_k": { + "shape": null, + "dtype": "int32" + }, + "block_m": { + "shape": null, + "dtype": "int32" + }, + "has_shared_expert": { + "shape": null, + "dtype": "bool" + } + }, + "outputs": { + "topk_weights": { + "shape": [ + "num_tokens", + "top_k" + ], + "dtype": "float32" + }, + "topk_ids": { + "shape": [ + "num_tokens", + "top_k" + ], + "dtype": "int32" + }, + "sorted_token_ids": { + "shape": [ + "slots" + ], + "dtype": "int32" + }, + "expert_ids": { + "shape": [ + "max_blocks" + ], + "dtype": "int32" + }, + "num_tokens_post_padded": { + "shape": [ + "one" + ], + "dtype": "int32" + }, + "expert_counts": { + "shape": [ + "num_experts" + ], + "dtype": "int32" + }, + "expert_offsets": { + "shape": [ + "num_experts_plus_one" + ], + "dtype": "int32" + }, + "expert_scatter_offsets": { + "shape": [ + "num_experts" + ], + "dtype": "int32" + } + }, + "reference": "from __future__ import annotations\nimport math\nimport torch\nimport torch.nn.functional as F\n\n@torch.no_grad()\ndef _alphamoe_fused_router_reference(\n router_logits: torch.Tensor,\n top_k: int,\n block_m: int,\n has_shared_expert: bool = False,\n **_unused,\n):\n \"\"\"Exact torch reference for the fused AlphaMoE gating router.\n\n Stable descending top-k over the routed experts (ties keep the lower\n expert index), optional shared-expert column, fp32 max-subtracted\n softmax over the selected logits, then the block-sparse routing\n metadata: expert histogram, block_m-aligned padded offsets and extent,\n per-expert scatter offsets (identical to the expert histogram, mirroring\n the upstream ``counts.clone()`` routing-plan output), expert-grouped\n flat route ids, and per-block expert ids.\n \"\"\"\n num_tokens, num_experts = router_logits.shape\n routed_experts = num_experts - int(has_shared_expert)\n routed_top_k = top_k - int(has_shared_expert)\n order = torch.argsort(\n router_logits[:, :routed_experts], dim=-1, descending=True, stable=True\n )[:, :routed_top_k]\n selected = torch.gather(router_logits, 1, order)\n if has_shared_expert:\n shared = torch.full(\n (num_tokens, 1),\n num_experts - 1,\n dtype=torch.int64,\n device=router_logits.device,\n )\n order = torch.cat((order, shared), dim=-1)\n selected = torch.cat((selected, router_logits[:, -1:]), dim=-1)\n topk_ids = order.to(torch.int32)\n topk_weights = torch.softmax(selected, dim=-1)\n flat = topk_ids.flatten().to(torch.int64)\n counts = torch.bincount(flat, minlength=num_experts).to(torch.int32)\n padded = (counts + block_m - 1) // block_m * block_m\n offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=router_logits.device)\n offsets[1:] = torch.cumsum(padded, dim=0)\n scatter_offsets = counts.clone()\n pairs = num_tokens * top_k\n nonempty = min(num_experts, pairs)\n max_blocks = nonempty + (pairs - nonempty) // block_m\n sorted_ids = torch.zeros(\n max_blocks * block_m, dtype=torch.int32, device=router_logits.device\n )\n expert_ids = torch.zeros(max_blocks, dtype=torch.int32, device=router_logits.device)\n sentinel = pairs\n for expert in range(num_experts):\n start = int(offsets[expert].item())\n count = int(counts[expert].item())\n end = int(offsets[expert + 1].item())\n if count:\n routes = torch.nonzero(flat == expert).flatten().to(torch.int32)\n sorted_ids[start : start + count] = routes\n sorted_ids[start + count : end] = sentinel\n expert_ids[start // block_m : end // block_m] = expert\n extent = offsets[-1:].clone()\n return (\n topk_weights,\n topk_ids,\n sorted_ids,\n expert_ids,\n extent,\n counts,\n offsets,\n scatter_offsets,\n )\n", + "check": "def standard_check(\n reference_outputs: Any,\n actual_outputs: Any,\n *,\n rtol: Optional[float] = None,\n atol: Optional[float] = None,\n max_mismatch_pct: float = 0.0,\n min_cos_sim: Optional[float] = 1.0 - 1e-3,\n) -> bool:\n \"\"\"Default trace correctness check used when a template does not override it.\"\"\"\n from flashinfer.trace import default_check\n\n return default_check(\n reference_outputs,\n actual_outputs,\n rtol=rtol,\n atol=atol,\n max_mismatch_pct=max_mismatch_pct,\n min_cos_sim=min_cos_sim,\n )\n", + "init": "from __future__ import annotations\nimport math\nimport torch\n\n# ----- shared init helpers -----\n# Copyright (c) 2025 by FlashInfer team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Shared helpers used by ``TraceTemplate.init`` functions.\n\nThis module contains the small set of input-construction patterns that\nrecur across many templates (paged-KV cache index arrays, ragged indptr,\nRoPE pos_ids and cos/sin caches, sampling probs). Each helper is short and\ndocumented; init functions in ``templates/.py`` call into here so\nthe per-template init bodies stay focused on shape/dtype, not boilerplate.\n\nThe full source of this module is **inlined into every dumped JSON's\n``\"init\"`` field** by ``flashinfer/trace/template.py:_render_init_source``,\nso downstream consumers don't need flashinfer installed to re-run the init\nsnippets.\n\"\"\"\n\n\nfrom typing import Optional, Tuple\n\nimport torch\n\n\ndef make_paged_kv_indices(\n batch_size: int,\n num_pages_per_seq: int,\n page_size: int,\n *,\n device: str = \"cuda\",\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Return ``(kv_indptr, kv_indices, kv_last_page_len)`` for a uniform batch.\n\n Every sequence is assigned exactly ``num_pages_per_seq`` pages, fully\n populated (last-page length == page_size).\n\n Invariants\n ----------\n - ``kv_indptr.shape == (batch_size + 1,)``, dtype int32, monotonic, [0]=0.\n - ``kv_indices == arange(0, batch_size * num_pages_per_seq)``, int32.\n - ``kv_last_page_len == full(batch_size, page_size)``, int32.\n \"\"\"\n total_pages = batch_size * num_pages_per_seq\n kv_indptr = (\n torch.arange(batch_size + 1, dtype=torch.int32, device=device)\n * num_pages_per_seq\n )\n kv_indices = torch.arange(total_pages, dtype=torch.int32, device=device)\n kv_last_page_len = torch.full(\n (batch_size,), page_size, dtype=torch.int32, device=device\n )\n return kv_indptr, kv_indices, kv_last_page_len\n\n\ndef make_ragged_indptr(\n seg_lens,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.int32,\n) -> torch.Tensor:\n \"\"\"Return cumulative-sum ``indptr`` of length ``len(seg_lens)+1``.\n\n ``seg_lens`` may be a list / tuple / 1-D tensor of segment lengths.\n \"\"\"\n if isinstance(seg_lens, torch.Tensor):\n lens = seg_lens.to(device=device, dtype=dtype)\n else:\n lens = torch.tensor(list(seg_lens), dtype=dtype, device=device)\n indptr = torch.zeros(lens.numel() + 1, dtype=dtype, device=device)\n indptr[1:] = torch.cumsum(lens, dim=0).to(dtype)\n return indptr\n\n\ndef make_uniform_qo_indptr(\n batch_size: int,\n qo_len: int,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, qo_len, 2*qo_len, ..., batch_size*qo_len]`` int32.\"\"\"\n return torch.arange(batch_size + 1, dtype=torch.int32, device=device) * qo_len\n\n\ndef make_pos_ids(\n nnz: int,\n max_seq_len: Optional[int] = None,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, 1, ..., nnz-1] (% max_seq_len)`` as int32 on ``device``.\n\n If ``max_seq_len`` is None, no wrapping is applied.\n \"\"\"\n pos = torch.arange(nnz, dtype=torch.int32, device=device)\n if max_seq_len is not None:\n pos = pos % max_seq_len\n return pos\n\n\ndef make_rope_cos_sin_cache(\n max_seq_len: int,\n rope_dim: int,\n *,\n base: float = 1e4,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return concatenated ``[cos | sin]`` cache of shape ``[max_seq_len, rope_dim]``.\"\"\"\n t = torch.arange(max_seq_len, dtype=torch.float32, device=device)\n inv = 1.0 / (\n base\n ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)\n )\n freqs = t.unsqueeze(-1) * inv.unsqueeze(0)\n cache = torch.cat([torch.cos(freqs), torch.sin(freqs)], dim=-1)\n return cache.to(dtype)\n\n\ndef make_probs(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return a ``[batch_size, vocab_size]`` probability distribution.\n\n Uses ``softmax(randn(...))`` so each row sums to 1.0. This mirrors the\n pattern used throughout ``tests/utils/test_sampling.py``.\n \"\"\"\n return torch.softmax(\n torch.randn(batch_size, vocab_size, dtype=torch.float32, device=device),\n dim=-1,\n ).to(dtype)\n\n\ndef make_logits(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return ``randn(batch_size, vocab_size)`` logits.\"\"\"\n return torch.randn(batch_size, vocab_size, dtype=dtype, device=device)\n\n\ndef fp8_safe_randn(\n *shape: int,\n scale: float = 0.1,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.bfloat16,\n) -> torch.Tensor:\n \"\"\"``randn(*shape) * scale`` \u2014 keeps values in the FP8/FP4 representable range.\n\n Tests for fp8/fp4 paths typically multiply ``randn`` by 0.1 to avoid\n saturation when quantizing. Use this helper to mirror that convention.\n \"\"\"\n return (torch.randn(*shape, dtype=dtype, device=device) * scale).to(dtype)\n\n\ndef per_tensor_fp8_quantize(\n x: torch.Tensor,\n *,\n fp8_dtype: torch.dtype = torch.float8_e4m3fn,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Per-tensor FP8 quantization, mirroring ``tests/utils_fp8.py:to_float8``.\n\n Returns ``(x_fp8, inv_scale)`` where ``inv_scale`` is the dequant\n multiplier (``float \u2248 fp8 * inv_scale``).\n \"\"\"\n finfo = torch.finfo(fp8_dtype)\n amax = x.abs().amax().clamp(min=1e-12)\n scale = finfo.max / amax\n x_q = (x.float() * scale).clamp(min=finfo.min, max=finfo.max).to(fp8_dtype)\n return x_q, scale.float().reciprocal()\n\n\ndef fp8_block_quant_1d(\n x_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize ``[T, H]`` activations into FP8 with per-``(token, block)``\n column-block scales. Returns ``(x_fp8, scales)`` where\n ``scales`` has shape ``[T, H // block]``.\n\n Mirrors ``_fp8_block_quant_1d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert x_bf16.dim() == 2\n T, H = x_bf16.shape\n assert H % block == 0\n nb = H // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n x_f32 = x_bf16.to(torch.float32)\n x_fp8 = torch.empty((T, H), dtype=torch.float8_e4m3fn, device=x_bf16.device)\n scales = torch.empty((T, nb), dtype=torch.float32, device=x_bf16.device)\n for j in range(nb):\n sl = slice(j * block, (j + 1) * block)\n blk = x_f32[:, sl]\n amax = torch.amax(torch.abs(blk), dim=1)\n s = torch.where(amax > 0, amax / max_fp8, torch.ones_like(amax))\n x_fp8[:, sl] = (blk / s.unsqueeze(1)).to(torch.float8_e4m3fn)\n scales[:, j] = s\n return x_fp8, scales\n\n\ndef fp8_block_quant_2d(\n w_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize weights ``[..., R, C]`` with 2-D ``block \u00d7 block`` scales.\n\n Returns ``(w_fp8, scales)`` where ``scales`` has shape\n ``[..., R // block, C // block]``. Mirrors ``_fp8_block_quant_2d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert w_bf16.dim() >= 2\n *prefix, R, C = w_bf16.shape\n assert R % block == 0 and C % block == 0\n nb_r, nb_c = R // block, C // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n w_f32 = w_bf16.to(torch.float32).contiguous()\n prefix_ndim = len(prefix)\n reshaped = w_f32.reshape(*prefix, nb_r, block, nb_c, block)\n permute_dims = tuple(range(prefix_ndim)) + (\n prefix_ndim,\n prefix_ndim + 2,\n prefix_ndim + 1,\n prefix_ndim + 3,\n )\n blocks = reshaped.permute(permute_dims).contiguous()\n amax = torch.amax(torch.abs(blocks), dim=(-1, -2))\n scales = torch.where(\n amax > 0, amax / max_fp8, torch.ones_like(amax, dtype=torch.float32)\n )\n q_blocks = (blocks / scales.unsqueeze(-1).unsqueeze(-1)).to(torch.float8_e4m3fn)\n inv_permute = [0] * (prefix_ndim + 4)\n for i, p in enumerate(permute_dims):\n inv_permute[p] = i\n w_fp8 = q_blocks.permute(*inv_permute).reshape(*prefix, R, C).contiguous()\n return w_fp8, scales\n\n\n__all__ = [\n \"make_paged_kv_indices\",\n \"make_ragged_indptr\",\n \"make_uniform_qo_indptr\",\n \"make_pos_ids\",\n \"make_rope_cos_sin_cache\",\n \"make_probs\",\n \"make_logits\",\n \"fp8_safe_randn\",\n \"per_tensor_fp8_quantize\",\n \"fp8_block_quant_1d\",\n \"fp8_block_quant_2d\",\n]\n\n# ----- init -----\ndef _alphamoe_fused_router_init(\n *,\n num_tokens: int,\n num_experts: int = 256,\n top_k: int = 8,\n block_m: int = 16,\n has_shared_expert: bool = False,\n # Derived by the exact output geometry; accepted only so the signature\n # carries every Var axis, and recomputed rather than used.\n max_blocks: int = 0,\n slots: int = 0,\n one: int = 1,\n num_experts_plus_one: int = 0,\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for the fused AlphaMoE gating router.\"\"\"\n torch.manual_seed(seed)\n router_logits = torch.randn(\n num_tokens, num_experts, dtype=torch.float32, device=device\n )\n return {\n \"router_logits\": router_logits,\n \"top_k\": int(top_k),\n \"block_m\": int(block_m),\n \"has_shared_expert\": bool(has_shared_expert),\n }\n" +} \ No newline at end of file diff --git a/tests/trace/fi_trace_out/alphamoe_fused_router_e257_k9_b8.json b/tests/trace/fi_trace_out/alphamoe_fused_router_e257_k9_b8.json new file mode 100644 index 00000000000..9d5b1e29068 --- /dev/null +++ b/tests/trace/fi_trace_out/alphamoe_fused_router_e257_k9_b8.json @@ -0,0 +1,120 @@ +{ + "name": "alphamoe_fused_router_e257_k9_b8", + "description": "Fused AlphaMoE gating router: stable descending top-k with an optional shared-expert column -> fp32 max-subtracted softmax -> block_m-aligned routing metadata bundle (expert histogram, padded expert offsets + extent, scatter offsets equal to the expert counts, expert-grouped flat route ids with sentinel padding, per-block expert ids). Route ids are flat token*top_k+slot indices; in-segment padding slots carry the sentinel num_tokens*top_k.", + "op_type": "moe_routing", + "tags": [ + "fi_api:flashinfer.fused_moe.alphamoe_router.alphamoe_fused_router", + "status:verified", + "moe", + "moe:routing" + ], + "axes": { + "num_tokens": { + "type": "var" + }, + "num_experts": { + "type": "const", + "value": 257 + }, + "top_k": { + "type": "const", + "value": 9 + }, + "block_m": { + "type": "const", + "value": 8 + }, + "max_blocks": { + "type": "var", + "description": "Block-bound from the output geometry (min(num_experts, pairs) + (pairs - nonempty) // block_m)." + }, + "slots": { + "type": "var", + "description": "max_blocks * block_m." + }, + "one": { + "type": "var", + "description": "Placeholder for shape [1] output tensors." + }, + "num_experts_plus_one": { + "type": "var", + "description": "num_experts + 1 for the inclusive padded offsets." + } + }, + "inputs": { + "router_logits": { + "shape": [ + "num_tokens", + "num_experts" + ], + "dtype": "float32" + }, + "top_k": { + "shape": null, + "dtype": "int32" + }, + "block_m": { + "shape": null, + "dtype": "int32" + }, + "has_shared_expert": { + "shape": null, + "dtype": "bool" + } + }, + "outputs": { + "topk_weights": { + "shape": [ + "num_tokens", + "top_k" + ], + "dtype": "float32" + }, + "topk_ids": { + "shape": [ + "num_tokens", + "top_k" + ], + "dtype": "int32" + }, + "sorted_token_ids": { + "shape": [ + "slots" + ], + "dtype": "int32" + }, + "expert_ids": { + "shape": [ + "max_blocks" + ], + "dtype": "int32" + }, + "num_tokens_post_padded": { + "shape": [ + "one" + ], + "dtype": "int32" + }, + "expert_counts": { + "shape": [ + "num_experts" + ], + "dtype": "int32" + }, + "expert_offsets": { + "shape": [ + "num_experts_plus_one" + ], + "dtype": "int32" + }, + "expert_scatter_offsets": { + "shape": [ + "num_experts" + ], + "dtype": "int32" + } + }, + "reference": "from __future__ import annotations\nimport math\nimport torch\nimport torch.nn.functional as F\n\n@torch.no_grad()\ndef _alphamoe_fused_router_reference(\n router_logits: torch.Tensor,\n top_k: int,\n block_m: int,\n has_shared_expert: bool = False,\n **_unused,\n):\n \"\"\"Exact torch reference for the fused AlphaMoE gating router.\n\n Stable descending top-k over the routed experts (ties keep the lower\n expert index), optional shared-expert column, fp32 max-subtracted\n softmax over the selected logits, then the block-sparse routing\n metadata: expert histogram, block_m-aligned padded offsets and extent,\n per-expert scatter offsets (identical to the expert histogram, mirroring\n the upstream ``counts.clone()`` routing-plan output), expert-grouped\n flat route ids, and per-block expert ids.\n \"\"\"\n num_tokens, num_experts = router_logits.shape\n routed_experts = num_experts - int(has_shared_expert)\n routed_top_k = top_k - int(has_shared_expert)\n order = torch.argsort(\n router_logits[:, :routed_experts], dim=-1, descending=True, stable=True\n )[:, :routed_top_k]\n selected = torch.gather(router_logits, 1, order)\n if has_shared_expert:\n shared = torch.full(\n (num_tokens, 1),\n num_experts - 1,\n dtype=torch.int64,\n device=router_logits.device,\n )\n order = torch.cat((order, shared), dim=-1)\n selected = torch.cat((selected, router_logits[:, -1:]), dim=-1)\n topk_ids = order.to(torch.int32)\n topk_weights = torch.softmax(selected, dim=-1)\n flat = topk_ids.flatten().to(torch.int64)\n counts = torch.bincount(flat, minlength=num_experts).to(torch.int32)\n padded = (counts + block_m - 1) // block_m * block_m\n offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=router_logits.device)\n offsets[1:] = torch.cumsum(padded, dim=0)\n scatter_offsets = counts.clone()\n pairs = num_tokens * top_k\n nonempty = min(num_experts, pairs)\n max_blocks = nonempty + (pairs - nonempty) // block_m\n sorted_ids = torch.zeros(\n max_blocks * block_m, dtype=torch.int32, device=router_logits.device\n )\n expert_ids = torch.zeros(max_blocks, dtype=torch.int32, device=router_logits.device)\n sentinel = pairs\n for expert in range(num_experts):\n start = int(offsets[expert].item())\n count = int(counts[expert].item())\n end = int(offsets[expert + 1].item())\n if count:\n routes = torch.nonzero(flat == expert).flatten().to(torch.int32)\n sorted_ids[start : start + count] = routes\n sorted_ids[start + count : end] = sentinel\n expert_ids[start // block_m : end // block_m] = expert\n extent = offsets[-1:].clone()\n return (\n topk_weights,\n topk_ids,\n sorted_ids,\n expert_ids,\n extent,\n counts,\n offsets,\n scatter_offsets,\n )\n", + "check": "def standard_check(\n reference_outputs: Any,\n actual_outputs: Any,\n *,\n rtol: Optional[float] = None,\n atol: Optional[float] = None,\n max_mismatch_pct: float = 0.0,\n min_cos_sim: Optional[float] = 1.0 - 1e-3,\n) -> bool:\n \"\"\"Default trace correctness check used when a template does not override it.\"\"\"\n from flashinfer.trace import default_check\n\n return default_check(\n reference_outputs,\n actual_outputs,\n rtol=rtol,\n atol=atol,\n max_mismatch_pct=max_mismatch_pct,\n min_cos_sim=min_cos_sim,\n )\n", + "init": "from __future__ import annotations\nimport math\nimport torch\n\n# ----- shared init helpers -----\n# Copyright (c) 2025 by FlashInfer team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Shared helpers used by ``TraceTemplate.init`` functions.\n\nThis module contains the small set of input-construction patterns that\nrecur across many templates (paged-KV cache index arrays, ragged indptr,\nRoPE pos_ids and cos/sin caches, sampling probs). Each helper is short and\ndocumented; init functions in ``templates/.py`` call into here so\nthe per-template init bodies stay focused on shape/dtype, not boilerplate.\n\nThe full source of this module is **inlined into every dumped JSON's\n``\"init\"`` field** by ``flashinfer/trace/template.py:_render_init_source``,\nso downstream consumers don't need flashinfer installed to re-run the init\nsnippets.\n\"\"\"\n\n\nfrom typing import Optional, Tuple\n\nimport torch\n\n\ndef make_paged_kv_indices(\n batch_size: int,\n num_pages_per_seq: int,\n page_size: int,\n *,\n device: str = \"cuda\",\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Return ``(kv_indptr, kv_indices, kv_last_page_len)`` for a uniform batch.\n\n Every sequence is assigned exactly ``num_pages_per_seq`` pages, fully\n populated (last-page length == page_size).\n\n Invariants\n ----------\n - ``kv_indptr.shape == (batch_size + 1,)``, dtype int32, monotonic, [0]=0.\n - ``kv_indices == arange(0, batch_size * num_pages_per_seq)``, int32.\n - ``kv_last_page_len == full(batch_size, page_size)``, int32.\n \"\"\"\n total_pages = batch_size * num_pages_per_seq\n kv_indptr = (\n torch.arange(batch_size + 1, dtype=torch.int32, device=device)\n * num_pages_per_seq\n )\n kv_indices = torch.arange(total_pages, dtype=torch.int32, device=device)\n kv_last_page_len = torch.full(\n (batch_size,), page_size, dtype=torch.int32, device=device\n )\n return kv_indptr, kv_indices, kv_last_page_len\n\n\ndef make_ragged_indptr(\n seg_lens,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.int32,\n) -> torch.Tensor:\n \"\"\"Return cumulative-sum ``indptr`` of length ``len(seg_lens)+1``.\n\n ``seg_lens`` may be a list / tuple / 1-D tensor of segment lengths.\n \"\"\"\n if isinstance(seg_lens, torch.Tensor):\n lens = seg_lens.to(device=device, dtype=dtype)\n else:\n lens = torch.tensor(list(seg_lens), dtype=dtype, device=device)\n indptr = torch.zeros(lens.numel() + 1, dtype=dtype, device=device)\n indptr[1:] = torch.cumsum(lens, dim=0).to(dtype)\n return indptr\n\n\ndef make_uniform_qo_indptr(\n batch_size: int,\n qo_len: int,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, qo_len, 2*qo_len, ..., batch_size*qo_len]`` int32.\"\"\"\n return torch.arange(batch_size + 1, dtype=torch.int32, device=device) * qo_len\n\n\ndef make_pos_ids(\n nnz: int,\n max_seq_len: Optional[int] = None,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, 1, ..., nnz-1] (% max_seq_len)`` as int32 on ``device``.\n\n If ``max_seq_len`` is None, no wrapping is applied.\n \"\"\"\n pos = torch.arange(nnz, dtype=torch.int32, device=device)\n if max_seq_len is not None:\n pos = pos % max_seq_len\n return pos\n\n\ndef make_rope_cos_sin_cache(\n max_seq_len: int,\n rope_dim: int,\n *,\n base: float = 1e4,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return concatenated ``[cos | sin]`` cache of shape ``[max_seq_len, rope_dim]``.\"\"\"\n t = torch.arange(max_seq_len, dtype=torch.float32, device=device)\n inv = 1.0 / (\n base\n ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)\n )\n freqs = t.unsqueeze(-1) * inv.unsqueeze(0)\n cache = torch.cat([torch.cos(freqs), torch.sin(freqs)], dim=-1)\n return cache.to(dtype)\n\n\ndef make_probs(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return a ``[batch_size, vocab_size]`` probability distribution.\n\n Uses ``softmax(randn(...))`` so each row sums to 1.0. This mirrors the\n pattern used throughout ``tests/utils/test_sampling.py``.\n \"\"\"\n return torch.softmax(\n torch.randn(batch_size, vocab_size, dtype=torch.float32, device=device),\n dim=-1,\n ).to(dtype)\n\n\ndef make_logits(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return ``randn(batch_size, vocab_size)`` logits.\"\"\"\n return torch.randn(batch_size, vocab_size, dtype=dtype, device=device)\n\n\ndef fp8_safe_randn(\n *shape: int,\n scale: float = 0.1,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.bfloat16,\n) -> torch.Tensor:\n \"\"\"``randn(*shape) * scale`` \u2014 keeps values in the FP8/FP4 representable range.\n\n Tests for fp8/fp4 paths typically multiply ``randn`` by 0.1 to avoid\n saturation when quantizing. Use this helper to mirror that convention.\n \"\"\"\n return (torch.randn(*shape, dtype=dtype, device=device) * scale).to(dtype)\n\n\ndef per_tensor_fp8_quantize(\n x: torch.Tensor,\n *,\n fp8_dtype: torch.dtype = torch.float8_e4m3fn,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Per-tensor FP8 quantization, mirroring ``tests/utils_fp8.py:to_float8``.\n\n Returns ``(x_fp8, inv_scale)`` where ``inv_scale`` is the dequant\n multiplier (``float \u2248 fp8 * inv_scale``).\n \"\"\"\n finfo = torch.finfo(fp8_dtype)\n amax = x.abs().amax().clamp(min=1e-12)\n scale = finfo.max / amax\n x_q = (x.float() * scale).clamp(min=finfo.min, max=finfo.max).to(fp8_dtype)\n return x_q, scale.float().reciprocal()\n\n\ndef fp8_block_quant_1d(\n x_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize ``[T, H]`` activations into FP8 with per-``(token, block)``\n column-block scales. Returns ``(x_fp8, scales)`` where\n ``scales`` has shape ``[T, H // block]``.\n\n Mirrors ``_fp8_block_quant_1d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert x_bf16.dim() == 2\n T, H = x_bf16.shape\n assert H % block == 0\n nb = H // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n x_f32 = x_bf16.to(torch.float32)\n x_fp8 = torch.empty((T, H), dtype=torch.float8_e4m3fn, device=x_bf16.device)\n scales = torch.empty((T, nb), dtype=torch.float32, device=x_bf16.device)\n for j in range(nb):\n sl = slice(j * block, (j + 1) * block)\n blk = x_f32[:, sl]\n amax = torch.amax(torch.abs(blk), dim=1)\n s = torch.where(amax > 0, amax / max_fp8, torch.ones_like(amax))\n x_fp8[:, sl] = (blk / s.unsqueeze(1)).to(torch.float8_e4m3fn)\n scales[:, j] = s\n return x_fp8, scales\n\n\ndef fp8_block_quant_2d(\n w_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize weights ``[..., R, C]`` with 2-D ``block \u00d7 block`` scales.\n\n Returns ``(w_fp8, scales)`` where ``scales`` has shape\n ``[..., R // block, C // block]``. Mirrors ``_fp8_block_quant_2d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert w_bf16.dim() >= 2\n *prefix, R, C = w_bf16.shape\n assert R % block == 0 and C % block == 0\n nb_r, nb_c = R // block, C // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n w_f32 = w_bf16.to(torch.float32).contiguous()\n prefix_ndim = len(prefix)\n reshaped = w_f32.reshape(*prefix, nb_r, block, nb_c, block)\n permute_dims = tuple(range(prefix_ndim)) + (\n prefix_ndim,\n prefix_ndim + 2,\n prefix_ndim + 1,\n prefix_ndim + 3,\n )\n blocks = reshaped.permute(permute_dims).contiguous()\n amax = torch.amax(torch.abs(blocks), dim=(-1, -2))\n scales = torch.where(\n amax > 0, amax / max_fp8, torch.ones_like(amax, dtype=torch.float32)\n )\n q_blocks = (blocks / scales.unsqueeze(-1).unsqueeze(-1)).to(torch.float8_e4m3fn)\n inv_permute = [0] * (prefix_ndim + 4)\n for i, p in enumerate(permute_dims):\n inv_permute[p] = i\n w_fp8 = q_blocks.permute(*inv_permute).reshape(*prefix, R, C).contiguous()\n return w_fp8, scales\n\n\n__all__ = [\n \"make_paged_kv_indices\",\n \"make_ragged_indptr\",\n \"make_uniform_qo_indptr\",\n \"make_pos_ids\",\n \"make_rope_cos_sin_cache\",\n \"make_probs\",\n \"make_logits\",\n \"fp8_safe_randn\",\n \"per_tensor_fp8_quantize\",\n \"fp8_block_quant_1d\",\n \"fp8_block_quant_2d\",\n]\n\n# ----- init -----\ndef _alphamoe_fused_router_init(\n *,\n num_tokens: int,\n num_experts: int = 256,\n top_k: int = 8,\n block_m: int = 16,\n has_shared_expert: bool = False,\n # Derived by the exact output geometry; accepted only so the signature\n # carries every Var axis, and recomputed rather than used.\n max_blocks: int = 0,\n slots: int = 0,\n one: int = 1,\n num_experts_plus_one: int = 0,\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for the fused AlphaMoE gating router.\"\"\"\n torch.manual_seed(seed)\n router_logits = torch.randn(\n num_tokens, num_experts, dtype=torch.float32, device=device\n )\n return {\n \"router_logits\": router_logits,\n \"top_k\": int(top_k),\n \"block_m\": int(block_m),\n \"has_shared_expert\": bool(has_shared_expert),\n }\n" +} \ No newline at end of file diff --git a/tests/trace/template_registry.py b/tests/trace/template_registry.py index 357ed1c1231..29cdd410d77 100644 --- a/tests/trace/template_registry.py +++ b/tests/trace/template_registry.py @@ -53,6 +53,7 @@ "flashinfer.cute_dsl.attention.wrappers.batch_prefill", "flashinfer.cute_dsl.rmsnorm_fp4quant", "flashinfer.decode", + "flashinfer.fused_moe.alphamoe_router", "flashinfer.fused_moe.core", "flashinfer.fused_moe.cute_dsl.b12x_moe", "flashinfer.fused_moe.cute_dsl.fused_moe", diff --git a/tests/trace/test_alphamoe_fused_router_trace.py b/tests/trace/test_alphamoe_fused_router_trace.py new file mode 100644 index 00000000000..9d2b63ecb11 --- /dev/null +++ b/tests/trace/test_alphamoe_fused_router_trace.py @@ -0,0 +1,217 @@ +"""Targeted trace coverage for the fused AlphaMoE gating router template. + +Covers definition-name encoding, committed fi_trace_out artifacts that must +exec and run standalone, the rendered init source, and end-to-end fi_trace +emission through the public ``fi_trace`` attribute. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch + +from flashinfer.trace.templates.moe import ( + alphamoe_fused_router_trace as ALPHAMOE, +) + +FI_TRACE_OUT = Path(__file__).parent / "fi_trace_out" +PLAIN_JSON = FI_TRACE_OUT / "alphamoe_fused_router_e256_k8_b16.json" +SHARED_JSON = FI_TRACE_OUT / "alphamoe_fused_router_e257_k9_b8.json" + + +def _exec_in_fresh_namespace(source: str) -> dict: + namespace: dict = {} + exec(source, namespace) # noqa: S102 — exercising the emitted source is the point + return namespace + + +def _axes(**overrides): + base = dict(num_experts=256, top_k=8, block_m=16) + base.update(overrides) + return base + + +# --------------------------------------------------------------------------- +# 1. Definition-name encoding +# --------------------------------------------------------------------------- + + +def test_alphamoe_name_encodes_const_axes(): + """The definition name is the prefix plus the const-axis abbreviation vector.""" + assert ( + ALPHAMOE.definition_name(_axes()) == "alphamoe_fused_router_e256_k8_b16" + ) + assert ( + ALPHAMOE.definition_name(_axes(num_experts=257, top_k=9, block_m=8)) + == "alphamoe_fused_router_e257_k9_b8" + ) + + +def test_alphamoe_routed_and_shared_geometries_have_distinct_names(): + """Shared-expert geometry (E+1 experts, top-k+1) never collides with plain.""" + plain = ALPHAMOE.definition_name(_axes()) + shared = ALPHAMOE.definition_name(_axes(num_experts=257, top_k=9, block_m=8)) + assert plain != shared + + +def test_alphamoe_var_axes_do_not_enter_the_name(): + """num_tokens is a var axis: a batch resize must not rename the definition.""" + name = ALPHAMOE.definition_name(_axes()) + named_with_batch = ALPHAMOE.definition_name({**_axes(), "num_tokens": 128}) + assert named_with_batch == name + + +# --------------------------------------------------------------------------- +# 2. Committed artifacts are runnable standalone +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not PLAIN_JSON.exists(), reason="alphamoe trace not generated") +def test_committed_alphamoe_init_runs_standalone(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + source = json.loads(PLAIN_JSON.read_text())["init"] + namespace = _exec_in_fresh_namespace(source) + init = namespace["_alphamoe_fused_router_init"] + + out = init(num_tokens=4, num_experts=32, top_k=4, block_m=8, device="cuda") + assert out["router_logits"].shape == (4, 32) + assert out["router_logits"].dtype == torch.float32 + assert out["router_logits"].is_cuda + assert out["top_k"] == 4 and out["block_m"] == 8 + assert out["has_shared_expert"] is False + + +@pytest.mark.skipif(not SHARED_JSON.exists(), reason="alphamoe trace not generated") +@pytest.mark.parametrize("has_shared_expert", [False, True]) +def test_committed_alphamoe_reference_runs_standalone(has_shared_expert): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + source = json.loads(SHARED_JSON.read_text())["reference"] + namespace = _exec_in_fresh_namespace(source) + reference = namespace["_alphamoe_fused_router_reference"] + + T, E, K, B = 6, 33, 5, 8 + logits = torch.randn(T, E, dtype=torch.float32, device="cuda") + ( + topk_weights, + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + expert_counts, + expert_offsets, + expert_scatter_offsets, + ) = reference(logits, K, B, has_shared_expert) + + # softmax over the selected logits: rows sum to one. + assert topk_weights.shape == (T, K) + assert torch.allclose( + topk_weights.sum(dim=-1), torch.ones(T, device="cuda"), atol=1e-6 + ) + # stable descending selection sanity: ids are in range and unique per token. + assert topk_ids.shape == (T, K) + assert int(topk_ids.min()) >= 0 and int(topk_ids.max()) < E + for row in topk_ids.tolist(): + assert len(set(row)) == K + # shared expert column is the last selected id for every token. + if has_shared_expert: + assert (topk_ids[:, -1] == E - 1).all() + # histogram and padded-inclusive offsets are consistent. + assert expert_counts.shape == (E,) + assert int(expert_counts.sum()) == T * K + assert expert_offsets.shape == (E + 1,) + assert int(expert_offsets[0]) == 0 + assert int(expert_offsets[-1]) == int(num_tokens_post_padded[0]) + padded_block = ((expert_counts + B - 1) // B) * B + assert int(padded_block.sum()) == int(expert_offsets[-1]) + # scatter offsets mirror the upstream counts.clone() routing plan output. + assert torch.equal(expert_scatter_offsets, expert_counts) + # per-block expert ids: every nonempty expert's blocks carry its id. + expert = 5 + start, end = int(expert_offsets[expert]), int(expert_offsets[expert + 1]) + if start < end: + assert (expert_ids[start // B : end // B] == expert).all() + # sentinel tail inside the segment. + count = int(expert_counts[expert]) + if start + count < end: + assert (sorted_token_ids[start + count : end] == T * K).all() + # valid routes reference (token, slot) pairs of the same expert. + routes = sorted_token_ids[start : start + count].to(torch.int64) + assert (topk_ids.flatten()[routes] == expert).all() + + +# --------------------------------------------------------------------------- +# 3. Rendered source is standalone (name-resolution trap) +# --------------------------------------------------------------------------- + + +def test_alphamoe_init_renders_standalone(): + """Dump-time globals must not leak into the committed init source.""" + from flashinfer.trace.template import _render_init_source + from flashinfer.trace.templates.moe import _alphamoe_fused_router_init + + namespace = _exec_in_fresh_namespace( + _render_init_source(_alphamoe_fused_router_init) + ) + inputs = namespace["_alphamoe_fused_router_init"]( + num_tokens=4, num_experts=8, top_k=2, block_m=4, device="cpu" + ) + assert inputs["router_logits"].shape == (4, 8) + assert inputs["router_logits"].dtype == torch.float32 + + +def test_alphamoe_reference_renders_standalone(): + from flashinfer.trace.template import _render_reference_source + from flashinfer.trace.templates.moe import ( + _alphamoe_fused_router_init, + _alphamoe_fused_router_reference, + ) + + namespace = _exec_in_fresh_namespace( + _render_reference_source(_alphamoe_fused_router_reference) + ) + inputs = _alphamoe_fused_router_init( + num_tokens=3, + num_experts=8, + top_k=3, + block_m=8, + has_shared_expert=True, + device="cpu", + ) + out = namespace["_alphamoe_fused_router_reference"]( + inputs["router_logits"], + inputs["top_k"], + inputs["block_m"], + inputs["has_shared_expert"], + ) + assert len(out) == 8 + assert out[0].shape == (3, 3) and out[1].shape == (3, 3) + assert (out[1][:, -1] == 7).all() + + +# --------------------------------------------------------------------------- +# 4. End-to-end: public fi_trace emits a complete definition +# --------------------------------------------------------------------------- + + +def test_fi_trace_emits_alphamoe_definition(): + """Exercise the symbolic trace through the decorated public API.""" + import flashinfer + + defn = flashinfer.fused_moe.alphamoe_fused_router.fi_trace( + router_logits=torch.zeros(8, 32, dtype=torch.float32), + top_k=4, + block_m=8, + has_shared_expert=True, + ) + assert defn["op_type"] == "moe_routing" + assert defn["name"] == "alphamoe_fused_router_e32_k4_b8" + assert defn["axes"]["num_experts"]["value"] == 32 + assert defn["axes"]["top_k"]["value"] == 4 + assert defn["axes"]["block_m"]["value"] == 8 + assert "unknown" not in str(defn["inputs"]) + assert "unknown" not in str(defn["outputs"])