[Refactor][MoE] Separate dispatched expert compute

benchmarks/ops/_moe_bench_utils.py

@@ -0,0 +1,167 @@
+"""Small shared primitives for MoE benchmark entry points."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+
+def cuda_time_ms(fn: Callable[[], Any], *, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    begin = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    begin.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return begin.elapsed_time(end) / iters
+
+
+def host_time_ms(fn: Callable[[], Any], *, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    begin = time.perf_counter()
+    for _ in range(iters):
+        fn()
+    torch.cuda.synchronize()
+    return (time.perf_counter() - begin) * 1e3 / iters
+
+
+def measure_dispatch(
+    run_fresh: Callable[[], Any],
+    *,
+    warmup: int,
+    iters: int,
+    make_cached: Callable[[Any], Callable[[], Any]] | None = None,
+    validate_cached: Callable[[Any, Any], None] | None = None,
+) -> tuple[Any, float, float | None]:
+    """Measure fresh dispatch and an optional cached-handle path."""
+    fresh_result = run_fresh()
+    fresh_ms = cuda_time_ms(run_fresh, warmup=warmup, iters=iters)
+    if make_cached is None:
+        return fresh_result, fresh_ms, None
+    run_cached = make_cached(fresh_result)
+    cached_result = run_cached()
+    cached_ms = cuda_time_ms(run_cached, warmup=warmup, iters=iters)
+    if validate_cached is not None:
+        validate_cached(fresh_result, cached_result)
+    return fresh_result, fresh_ms, cached_ms
+
+
+def make_routing_inputs(
+    num_tokens: int,
+    hidden_size: int,
+    top_k: int,
+    num_experts: int,
+    *,
+    distribution: str = "uniform",
+    topk_dtype: torch.dtype = torch.int32,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    noise = torch.rand(num_tokens, num_experts, device="cuda")
+    if distribution == "uniform":
+        scores = noise
+    elif distribution == "hotspot":
+        scores = torch.linspace(4.0, -4.0, num_experts, device="cuda").unsqueeze(0) + noise
+    elif distribution == "longtail":
+        ranks = torch.arange(1, num_experts + 1, device="cuda")
+        scores = -torch.log(ranks).unsqueeze(0) + noise
+    else:
+        raise ValueError(f"unknown distribution {distribution!r}")
+    hidden = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda")
+    topk_ids = scores.topk(top_k, dim=-1).indices.to(topk_dtype)
+    weights = torch.softmax(torch.randn(num_tokens, top_k, device="cuda"), dim=-1)
+    return hidden, topk_ids, weights
+
+
+def make_expert_sizes(
+    num_pairs: int,
+    num_experts: int,
+    distribution: str,
+) -> list[int]:
+    if distribution == "uniform":
+        return [
+            num_pairs // num_experts + (expert < num_pairs % num_experts)
+            for expert in range(num_experts)
+        ]
+    generator = torch.Generator(device="cpu").manual_seed(42)
+    if distribution == "longtail":
+        probabilities = 1 / torch.arange(1, num_experts + 1, dtype=torch.float64).pow(1.2)
+    elif distribution == "hotspot":
+        hot_experts = max(1, num_experts // 8)
+        probabilities = torch.full((num_experts,), 0.2 / num_experts)
+        probabilities[:hot_experts] += 0.8 / hot_experts
+    elif distribution == "router-like":
+        probabilities = torch.softmax(torch.randn(num_experts, generator=generator) * 1.5, dim=0)
+    else:
+        raise ValueError(f"unknown distribution {distribution!r}")
+    assignments = torch.multinomial(
+        probabilities,
+        num_samples=num_pairs,
+        replacement=True,
+        generator=generator,
+    )
+    return torch.bincount(assignments, minlength=num_experts).tolist()
+
+
+def expert_mlp_reference(
+    hidden: torch.Tensor,
+    w_gate_up: torch.Tensor,
+    w_down: torch.Tensor,
+    sizes: list[int],
+) -> torch.Tensor:
+    output = torch.empty(
+        hidden.shape[0],
+        hidden.shape[1],
+        device=hidden.device,
+        dtype=torch.float32,
+    )
+    ffn_size = w_down.shape[-1]
+    start = 0
+    for expert, size in enumerate(sizes):
+        if size == 0:
+            continue
+        rows = hidden[start : start + size].float()
+        gate_up = rows @ w_gate_up[expert].float().t()
+        activated = F.silu(gate_up[:, :ffn_size]) * gate_up[:, ffn_size:]
+        output[start : start + size] = activated @ w_down[expert].float().t()
+        start += size
+    return output
+
+
+def check_output(
+    name: str,
+    actual: torch.Tensor,
+    reference: torch.Tensor,
+    *,
+    max_range_relative_tolerance: float,
+    relative_l2_tolerance: float,
+) -> tuple[float, float, float, float]:
+    actual = actual.float()
+    if not torch.isfinite(actual).all():
+        raise AssertionError(f"{name}: output contains NaN or Inf")
+    error = actual - reference
+    max_abs = error.abs().max().item()
+    rmse = error.square().mean().sqrt().item()
+    relative_l2 = (
+        torch.linalg.vector_norm(error) / torch.linalg.vector_norm(reference).clamp_min(1e-12)
+    ).item()
+    max_range_relative = max_abs / reference.abs().max().clamp_min(1e-12).item()
+    if max_range_relative > max_range_relative_tolerance or relative_l2 > relative_l2_tolerance:
+        raise AssertionError(
+            f"{name}: max_abs={max_abs}, rmse={rmse}, "
+            f"relative_l2={relative_l2}, "
+            f"max_range_relative={max_range_relative}"
+        )
+    return max_abs, rmse, relative_l2, max_range_relative
+
+
+def effective_tflops(logical_flops: int, elapsed_ms: float) -> float:
+    return logical_flops / (elapsed_ms / 1e3) / 1e12

benchmarks/ops/bench_deepep_dispatch.py

@@ -0,0 +1,153 @@
+"""Single-node multi-GPU DeepEP V2 dispatch benchmark for the M6 adapter.
+
+Example:
+    torchrun --standalone --nproc-per-node=8 \
+        benchmarks/ops/bench_deepep_dispatch.py \
+        --tokens 1 8 32 128 --hidden-size 7168 --num-experts 256 --top-k 8
+
+Only dispatch is timed.  Expert compute and combine are deliberately excluded.
+The script requires DeepEP V2, but TileOps itself does not depend on DeepEP.
+"""
+
+import argparse
+import json
+import os
+
+import torch
+import torch.distributed as dist
+from _moe_bench_utils import (
+    make_routing_inputs,
+    measure_dispatch,
+)
+
+from tileops.ops.moe import DeepEPDispatchAdapter
+
+WARMUP = 10
+ITERS = 50
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--hidden-size", type=int, default=7168)
+    parser.add_argument("--num-experts", type=int, default=256)
+    parser.add_argument("--top-k", type=int, default=8)
+    parser.add_argument("--tokens", type=int, nargs="+", default=[1, 8, 32, 128, 512])
+    parser.add_argument("--num-sms", type=int, default=0)
+    args = parser.parse_args()
+
+    try:
+        from deep_ep import ElasticBuffer
+    except ImportError as exc:
+        raise RuntimeError("bench_deepep_dispatch.py requires a DeepEP V2 installation") from exc
+
+    dist.init_process_group("nccl")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    group = dist.group.WORLD
+    world_size = dist.get_world_size(group)
+    rank = dist.get_rank(group)
+    if args.num_experts % world_size != 0:
+        raise ValueError(
+            f"num_experts must be divisible by world_size; got {args.num_experts} and {world_size}"
+        )
+    if args.top_k > args.num_experts:
+        raise ValueError("top_k cannot exceed num_experts")
+
+    max_tokens = max(args.tokens)
+    buffer = ElasticBuffer(
+        group,
+        num_max_tokens_per_rank=max_tokens,
+        hidden=args.hidden_size,
+        num_topk=args.top_k,
+        use_fp8_dispatch=False,
+    )
+    adapter = DeepEPDispatchAdapter(
+        buffer,
+        num_experts=args.num_experts,
+        num_local_experts=args.num_experts // world_size,
+        num_max_tokens_per_rank=max_tokens,
+        num_sms=args.num_sms,
+    )
+
+    for num_tokens in args.tokens:
+        torch.manual_seed(2026 + rank)
+        hidden, topk_ids, topk_weights = make_routing_inputs(
+            num_tokens,
+            args.hidden_size,
+            args.top_k,
+            args.num_experts,
+            topk_dtype=torch.int64,
+        )
+        offsets = torch.empty(
+            args.num_experts // world_size + 1,
+            dtype=torch.int32,
+            device="cuda",
+        )
+
+        def run_fresh(
+            hidden=hidden,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            offsets=offsets,
+        ):
+            return adapter.dispatch(
+                hidden,
+                topk_ids,
+                topk_weights,
+                expert_offsets=offsets,
+            )
+
+        def make_cached(
+            result,
+            offsets=offsets,
+            cached_hidden=hidden,
+            cached_weights=topk_weights,
+        ):
+            cached_offsets = torch.empty_like(offsets)
+
+            def run_cached():
+                return adapter.dispatch(
+                    cached_hidden,
+                    None,
+                    cached_weights,
+                    expert_offsets=cached_offsets,
+                    cached_handle=result.combine_handle,
+                )
+
+            return run_cached
+
+        def validate_cached(result, cached_result):
+            torch.testing.assert_close(
+                cached_result.batch.expert_offsets,
+                result.batch.expert_offsets,
+            )
+
+        result, fresh_ms, cached_ms = measure_dispatch(
+            run_fresh,
+            warmup=WARMUP,
+            iters=ITERS,
+            make_cached=make_cached,
+            validate_cached=validate_cached,
+        )
+        assert cached_ms is not None
+        record = {
+            "rank": rank,
+            "world_size": world_size,
+            "tokens": num_tokens,
+            "top_k": args.top_k,
+            "num_experts": args.num_experts,
+            "num_local_experts": args.num_experts // world_size,
+            "hidden_size": args.hidden_size,
+            "sent_pairs": num_tokens * args.top_k,
+            "received_pairs": int(result.batch.valid_rows.item()),
+            "physical_rows": result.batch.capacity,
+            "dispatch_fresh_allocating_ms": round(fresh_ms, 4),
+            "dispatch_cached_allocating_ms": round(cached_ms, 4),
+        }
+        print(json.dumps(record), flush=True)
+        dist.barrier(group)
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

benchmarks/ops/bench_dispatched_expert.py

@@ -0,0 +1,62 @@
+"""Dynamic-capacity benchmark for the dispatched expert MLP."""
+
+import torch
+from _moe_bench_utils import cuda_time_ms, effective_tflops
+
+from tileops.ops.moe import DispatchedExpertMLPFwdOp, ExpertBatch
+
+DTYPE = torch.bfloat16
+E, H, F = 128, 2048, 1024
+WARMUP, ITERS = 20, 100
+
+
+def _run_capacity_sweep(capacity: int = 16_384) -> None:
+    hidden = torch.randn(capacity, H, device="cuda", dtype=DTYPE)
+    w_gate_up = torch.randn(E, 2 * F, H, device="cuda", dtype=DTYPE) * 0.02
+    w_down = torch.randn(E, H, F, device="cuda", dtype=DTYPE) * 0.02
+    unfused = DispatchedExpertMLPFwdOp(capacity, E, H, F, DTYPE, use_fused_activation=False)
+    fused = DispatchedExpertMLPFwdOp(capacity, E, H, F, DTYPE, use_fused_activation=True)
+    print(
+        "capacity,valid_rows,utilization,"
+        "capacity_unfused_ms,capacity_fused_ms,"
+        "unfused_effective_TFLOPS,fused_effective_TFLOPS"
+    )
+    for valid_rows in (0, capacity // 64, capacity // 8, capacity):
+        sizes = [valid_rows // E + (expert < valid_rows % E) for expert in range(E)]
+        offsets = [0]
+        for size in sizes:
+            offsets.append(offsets[-1] + size)
+        batch = ExpertBatch(
+            hidden=hidden,
+            expert_offsets=torch.tensor(offsets, device="cuda", dtype=torch.int32),
+        )
+        unfused_ms = cuda_time_ms(
+            lambda batch=batch: unfused.forward_batch(batch, w_gate_up, w_down),
+            warmup=WARMUP,
+            iters=ITERS,
+        )
+        fused_ms = cuda_time_ms(
+            lambda batch=batch: fused.forward_batch(batch, w_gate_up, w_down),
+            warmup=WARMUP,
+            iters=ITERS,
+        )
+        logical_flops = 6 * valid_rows * H * F
+        print(
+            f"{capacity},{valid_rows},{valid_rows / capacity:.6f},"
+            f"{unfused_ms:.4f},{fused_ms:.4f},"
+            f"{effective_tflops(logical_flops, unfused_ms):.2f},"
+            f"{effective_tflops(logical_flops, fused_ms):.2f}",
+            flush=True,
+        )
+
+
+def main() -> None:
+    assert torch.cuda.is_available()
+    torch.manual_seed(42)
+    torch.set_grad_enabled(False)
+    print(f"GPU={torch.cuda.get_device_name()} dtype={DTYPE} E={E} H={H} F={F}")
+    _run_capacity_sweep()
+
+
+if __name__ == "__main__":
+    main()