refactor(moe_ep)!: move the MoE all-to-all ops out of flashinfer.comm into their moe_ep backends; retire backend="trtllm" and the moe_comm bench; SM107 one-sided build fix

benchmarks/README.md

@@ -54,8 +54,7 @@ Currently supports testing attention, gemm, fused MOE, normalization, quantizati
     - `cutlass_fused_moe` - CUTLASS fused MoE (base/fp8/nvfp4 variants with optional TP/EP)
     - `cute_dsl_bf16_moe` - CuTe-DSL BF16/FP16 fused MoE for Hopper.
     - `unified_moe` - Unified MoE API comparison between the CUTLASS, cuTile, b12x and cuDNN grouped-GEMM backends. It supports BF16, NVFP4 and MXFP4 W4A4/W4A16, per-tensor FP8 and MXFP8 W8A8/W8A16, and MXFP4 W4A8 with gated SwiGLU, SwiGLU-Step, GeGLU, GeGLU-Tanh, and SiTU or non-gated GELU, ReLU, SiLU, ReLU2, and Identity; filters unsupported backends at runtime; and can autotune each backend independently.
-- MOE Communication:
-    - `moe_a2a_dispatch_combine` - MoE All-to-All dispatch + combine benchmark for multi-GPU expert-parallel inference. Requires `mpirun` for multi-GPU execution. Supports optional quantization (FP8, NVFP4, FP8 block-scale) and real MoE kernel computation.
+- MOE Communication: no routine yet; standalone scripts cover it. `benchmarks/moe_ep/core/comm/bench_moe_ep_comm.py` times dispatch + combine of the `flashinfer.moe_ep` communication backends (launched with `torchrun`), and `benchmarks/moe_ep/bench_moe_ep.py` times dispatch + MoE + combine.
 - AllReduce Communication:
     - `allreduce_fusion` - AllReduce fusion benchmark for multi-GPU inference. Requires `mpirun` for multi-GPU execution. Supports TRTLLM and TRTLLM MNNVL backends with multiple fusion patterns (plain allreduce, allreduce + residual + RMSNorm).
 - Norm:
@@ -350,57 +349,6 @@ Notes:
 - FP8 MOE kernels require integer values for group parameters, while FP4 MOE kernels accept optional values.
 - CUTLASS fused MoE (`cutlass_fused_moe`) ignores `--routing_method`, `--n_group`, and `--topk_group`; it computes routing via softmax+top-k internally from the provided logits.
 
-### MoE Communication Flags (moe_a2a_dispatch_combine)
-The `moe_a2a_dispatch_combine` routine benchmarks MoE All-to-All communication for multi-GPU expert-parallel inference. It must be launched with `mpirun`.
-
-| Flag                     | Description                                                                                                 |
-|--------------------------|-------------------------------------------------------------------------------------------------------------|
-| `--num_tokens`           | Number of tokens per rank (local batch size)                                                               |
-| `--hidden_size`          | Hidden dimension size                                                                                      |
-| `--num_experts`          | Total number of experts across all ranks                                                                   |
-| `--top_k`                | Number of experts to route each token to                                                                   |
-| `--input_dtype`          | Data type for hidden states payload: `bfloat16` (default) or `float16`                                     |
-| `--quant_dtype`          | Quantization format: `fp8` (per-tensor), `nvfp4` (block-scale FP4), `fp8_block_scale` (block-scale FP8)    |
-| `--real_math`            | Run actual MoE kernels instead of fake computation. Requires `--intermediate_size` and `--quant_dtype` to be `nvfp4` or `fp8_block_scale` |
-| `--intermediate_size`    | Intermediate FFN size. Required if `--real_math` is set                                                    |
-| `--max_num_tokens`       | Max tokens per rank for workspace allocation. Defaults to `--num_tokens`                                   |
-| `--validate`             | Run correctness validation before benchmarking using deterministic fake MoE                                |
-| `--per_phase_timing`     | Enable per-phase timing (dispatch/combine/moe_kernel). Adds slight overhead from CUDA events               |
-| `--nvtx`                 | Enable NVTX markers for Nsight Systems profiling                                                           |
-| `--use_lora`             | Carry a per-token int32 LoRA adapter ID through dispatch as an extra payload.                                                                                                                  |
-
-**Launch Examples:**
-```bash
-# Basic (no quantization)
-mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
-    --routine moe_a2a_dispatch_combine \
-    --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8
-
-# With FP8 quantization
-mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
-    --routine moe_a2a_dispatch_combine \
-    --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 \
-    --quant_dtype fp8
-
-# With NVFP4 quantization and real MoE kernel
-mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
-    --routine moe_a2a_dispatch_combine \
-    --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 \
-    --quant_dtype nvfp4 --real_math --intermediate_size 18432
-
-# With validation and per-phase timing
-mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
-    --routine moe_a2a_dispatch_combine \
-    --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 \
-    --validate --per_phase_timing
-
-# Multi-tenant LoRA: carry per-token adapter ID through dispatch
-mpirun -np 8 python benchmarks/flashinfer_benchmark.py \
-    --routine moe_a2a_dispatch_combine \
-    --num_tokens 2048 --hidden_size 7168 --num_experts 256 --top_k 8 \
-    --use_lora --validate
-```
-
 ### AllReduce Communication Flags (allreduce_fusion)
 The `allreduce_fusion` routine benchmarks AllReduce fusion operations for multi-GPU inference. It must be launched with `mpirun`. Both oneshot and twoshot strategies are benchmarked automatically and reported side by side.
 
@@ -588,7 +536,6 @@ Legend:
 | **cutlass_fused_moe** |  |  |  |  |  | cutlass | cutlass |  |
 | **cute_dsl_bf16_moe** |  |  |  |  | cute-dsl |  |  |  |
 | **unified_moe** |  | cudnn (BF16) | cudnn (BF16) | cutlass (BF16), cutile (BF16, NVFP4/MXFP4 W4A16), cudnn (BF16, FP8) | cutlass (BF16, MXFP4 W4A16), cutile (BF16, NVFP4/MXFP4 W4A16), cudnn (BF16, FP8) | cutlass, cudnn | cutlass, cudnn | cutlass (BF16, NVFP4 W4A4), cutile (BF16, NVFP4/MXFP4 W4A4/W4A16), b12x (NVFP4 W4A4/W4A16), cudnn |
-| **moe_a2a_dispatch_combine** |  |  |  |  |  | moe_a2a | moe_a2a |  |
 | **allreduce_fusion** |  |  |  |  |  | allreduce | allreduce |  |
 | **rmsnorm** | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl |
 | **fused_add_rmsnorm** | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl | cute-dsl |

benchmarks/flashinfer_benchmark.py

@@ -36,10 +36,6 @@ def run_test(args):
         from routines.unified_moe import run_unified_moe_test
 
         res = run_unified_moe_test(args)
-    elif args.routine in benchmark_apis["moe_comm"]:
-        from routines.moe_comm import run_moe_comm_test
-
-        res = run_moe_comm_test(args)
     elif args.routine in benchmark_apis["allreduce_comm"]:
         from routines.allreduce_comm import run_allreduce_comm_test
 
@@ -129,7 +125,6 @@ def parse_args(line=sys.argv[1:]):
         + list(benchmark_apis["gemm"])
         + list(benchmark_apis["moe"])
         + list(benchmark_apis["unified_moe"])
-        + list(benchmark_apis["moe_comm"])
         + list(benchmark_apis["allreduce_comm"])
         + list(benchmark_apis["mixed_comm"])
         + list(benchmark_apis["norm"])
@@ -265,10 +260,6 @@ def parse_args(line=sys.argv[1:]):
         from routines.unified_moe import parse_unified_moe_args
 
         args = parse_unified_moe_args(line, parser)
-    elif args.routine in benchmark_apis["moe_comm"]:
-        from routines.moe_comm import parse_moe_comm_args
-
-        args = parse_moe_comm_args(line, parser)
     elif args.routine in benchmark_apis["allreduce_comm"]:
         from routines.allreduce_comm import parse_allreduce_comm_args
 

benchmarks/moe_ep/core/comm/bench_moe_ep_comm.py

@@ -51,6 +51,7 @@
 DISPATCH_FORMATS = {
     "bf16": None,
     "fp8": "FP8PerTensor",
+    "fp8_block_scale": "DeepSeekFp8",
     "mxfp8": "MXFP8",
     "nvfp4": "NVFP4",
 }
@@ -153,6 +154,17 @@ def quantize_fp8_per_tensor(x: torch.Tensor) -> Tuple[torch.Tensor, None]:
     return (x.float() * (FP8_E4M3_MAX / amax)).to(torch.float8_e4m3fn), None
 
 
+def quantize_fp8_block_scale(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """DeepSeek FP8: E4M3 values with one FP32 scale per 128 elements."""
+    num_tokens, hidden = x.shape
+    if hidden % 128:
+        raise ValueError(f"fp8_block_scale needs hidden_size % 128 == 0, got {hidden}")
+    blocks = x.float().view(num_tokens, hidden // 128, 128)
+    scales = blocks.abs().amax(-1, keepdim=True).clamp(min=1e-12) / FP8_E4M3_MAX
+    values = (blocks / scales).to(torch.float8_e4m3fn)
+    return values.view(num_tokens, hidden), scales.squeeze(-1)
+
+
 def quantize_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     """MXFP8: E4M3 values with one UE8M0 scale per 32 elements."""
     num_tokens, hidden = x.shape
@@ -208,6 +220,8 @@ def quantize_input(
         return x, None
     if quant == "fp8":
         return quantize_fp8_per_tensor(x)
+    if quant == "fp8_block_scale":
+        return quantize_fp8_block_scale(x)
     if quant == "mxfp8":
         return quantize_mxfp8(x)
     if quant == "nvfp4":

benchmarks/routines/flashinfer_benchmark_utils.py

@@ -84,13 +84,6 @@
         "ep_size",
         "ep_rank",
     ],
-    "moe_comm": [
-        "num_tokens",
-        "num_experts",
-        "top_k",
-        "ep_size",
-        "max_num_tokens",
-    ],
     "allreduce_comm": [
         "num_tokens",
         "ar_backend",
@@ -217,7 +210,6 @@
     + output_column_dict["dsv4_sparse_mla"]
     + output_column_dict["gemm"]
     + output_column_dict["moe"]
-    + output_column_dict["moe_comm"]
     + output_column_dict["allreduce_comm"]
     + output_column_dict["mixed_comm"]
     + output_column_dict["norm"]
@@ -270,9 +262,6 @@
     "unified_moe": [
         "unified_moe",
     ],
-    "moe_comm": [
-        "moe_a2a_dispatch_combine",
-    ],
     "allreduce_comm": [
         "allreduce_fusion",
     ],