[Refactor][MoE] Separate dispatched expert compute

benchmarks/ops/bench_deepep_dispatch.py

@@ -0,0 +1,163 @@
+"""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 tileops.ops.moe import DeepEPDispatchAdapter
+
+WARMUP = 10
+ITERS = 50
+
+
+def _time_ms(fn) -> 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 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 = torch.randn(
+            num_tokens,
+            args.hidden_size,
+            dtype=torch.bfloat16,
+            device="cuda",
+        )
+        topk_ids = (
+            torch.rand(num_tokens, args.num_experts, device="cuda")
+            .topk(args.top_k, dim=-1)
+            .indices.to(torch.int64)
+        )
+        topk_weights = torch.softmax(torch.randn(num_tokens, args.top_k, device="cuda"), dim=-1)
+        offsets = torch.empty(
+            args.num_experts // world_size + 1,
+            dtype=torch.int32,
+            device="cuda",
+        )
+
+        def run(
+            hidden=hidden,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            offsets=offsets,
+        ):
+            return adapter.dispatch(
+                hidden,
+                topk_ids,
+                topk_weights,
+                expert_offsets=offsets,
+            )
+
+        result = run()
+        dispatch_fresh_ms = _time_ms(run)
+        cached_offsets = torch.empty_like(offsets)
+
+        def run_cached(
+            hidden=hidden,
+            topk_weights=topk_weights,
+            offsets=cached_offsets,
+            cached_handle=result.combine_handle,
+        ):
+            return adapter.dispatch(
+                hidden,
+                None,
+                topk_weights,
+                expert_offsets=offsets,
+                cached_handle=cached_handle,
+            )
+
+        cached_result = run_cached()
+        dispatch_cached_ms = _time_ms(run_cached)
+        torch.testing.assert_close(
+            cached_result.batch.expert_offsets,
+            result.batch.expert_offsets,
+        )
+        valid_rows = int(result.batch.valid_rows.item())
+        physical_rows = result.batch.capacity
+        sent_pairs = num_tokens * args.top_k
+        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": sent_pairs,
+            "received_pairs": valid_rows,
+            "physical_rows": physical_rows,
+            "dispatch_fresh_allocating_ms": round(dispatch_fresh_ms, 4),
+            "dispatch_cached_allocating_ms": round(dispatch_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,205 @@
+"""H100 microbenchmark for the communication-independent expert MLP."""
+
+import torch
+
+from tileops.ops.moe import (
+    DispatchedExpertMLPFwdOp,
+    ExpertBatch,
+    FusedMoEExpertsNopadPersistent3WGFwdOp,
+)
+from tileops.ops.moe.routed_expert import MoePermuteNopadFwdOp
+
+DTYPE = torch.bfloat16
+E, TOP_K, H, F = 128, 8, 2048, 1024
+WARMUP, ITERS = 20, 100
+
+
+def _time_ms(fn) -> 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 _run_shape(num_tokens: int) -> tuple[float, ...]:
+    num_pairs = num_tokens * TOP_K
+    hidden = torch.randn(num_tokens, 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
+    # Select unique experts per token and materialize the exact expert-major
+    # batch used by the full path. Pure/full timings therefore differ only by
+    # routing work, not by expert-size distribution.
+    topk_ids = torch.rand(
+        num_tokens, E, device="cuda"
+    ).topk(TOP_K, dim=-1).indices.to(torch.int32)
+    topk_weights = torch.softmax(
+        torch.randn(num_tokens, TOP_K, device="cuda"), dim=-1
+    )
+    permute = MoePermuteNopadFwdOp(
+        total_tokens=num_tokens,
+        top_k=TOP_K,
+        num_experts=E,
+        hidden_size=H,
+        dtype=DTYPE,
+    )
+    expert_input, true_offsets, true_sizes, _, _ = permute(
+        hidden, topk_ids
+    )
+    output_unfused = torch.empty(
+        num_tokens, H, device="cuda", dtype=DTYPE
+    )
+    output_fused = torch.empty_like(output_unfused)
+    workspace = torch.empty(0, device="cuda", dtype=DTYPE)
+
+    pure_unfused = DispatchedExpertMLPFwdOp(
+        num_pairs, E, H, F, DTYPE, use_fused_activation=False
+    )
+    pure_fused = DispatchedExpertMLPFwdOp(
+        num_pairs, E, H, F, DTYPE, use_fused_activation=True
+    )
+    full_unfused = FusedMoEExpertsNopadPersistent3WGFwdOp(
+        num_tokens, E, TOP_K, H, F, dtype=DTYPE, use_fused_activation=False
+    )
+    full_fused = FusedMoEExpertsNopadPersistent3WGFwdOp(
+        num_tokens, E, TOP_K, H, F, dtype=DTYPE, use_fused_activation=True
+    )
+
+    def run_pure_unfused():
+        return pure_unfused(
+            expert_input, w_gate_up, w_down, true_sizes, true_offsets
+        )
+
+    def run_pure_fused():
+        return pure_fused(
+            expert_input, w_gate_up, w_down, true_sizes, true_offsets
+        )
+
+    def run_full_unfused():
+        return full_unfused.forward(
+            output_unfused,
+            hidden,
+            w_gate_up,
+            w_down,
+            topk_weights,
+            topk_ids,
+            None,
+            workspace,
+            workspace,
+            E,
+        )
+
+    def run_full_fused():
+        return full_fused.forward(
+            output_fused,
+            hidden,
+            w_gate_up,
+            w_down,
+            topk_weights,
+            topk_ids,
+            None,
+            workspace,
+            workspace,
+            E,
+        )
+
+    timings = tuple(
+        _time_ms(fn)
+        for fn in (
+            run_pure_unfused,
+            run_pure_fused,
+            run_full_unfused,
+            run_full_fused,
+        )
+    )
+    logical_flops = 6 * num_pairs * H * F
+    tflops = tuple(
+        logical_flops / (timing / 1e3) / 1e12
+        for timing in timings
+    )
+    return (*timings, *tflops)
+
+
+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 = _time_ms(
+            lambda batch=batch: unfused.forward_batch(
+                batch, w_gate_up, w_down
+            )
+        )
+        fused_ms = _time_ms(
+            lambda batch=batch: fused.forward_batch(
+                batch, w_gate_up, w_down
+            )
+        )
+        logical_flops = 6 * valid_rows * H * F
+        print(
+            f"{capacity},{valid_rows},{valid_rows / capacity:.6f},"
+            f"{unfused_ms:.4f},{fused_ms:.4f},"
+            f"{logical_flops / (unfused_ms / 1e3) / 1e12:.2f},"
+            f"{logical_flops / (fused_ms / 1e3) / 1e12:.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} "
+        f"E={E} K={TOP_K} H={H} F={F}"
+    )
+    print(
+        "T,M,pure_unfused_ms,pure_fused_ms,"
+        "full_unfused_ms,full_fused_ms,"
+        "pure_unfused_effective_TFLOPS,pure_fused_effective_TFLOPS,"
+        "full_unfused_effective_TFLOPS,full_fused_effective_TFLOPS"
+    )
+    for num_tokens in (128, 1024):
+        values = _run_shape(num_tokens)
+        formatted = ",".join(f"{value:.4f}" for value in values)
+        print(f"{num_tokens},{num_tokens * TOP_K},{formatted}", flush=True)
+        torch.cuda.empty_cache()
+    _run_capacity_sweep()
+
+
+if __name__ == "__main__":
+    main()