From e9abf452b0439fd189473a7bcc6ed2c160c40be6 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Sat, 17 Oct 2026 06:19:49 +0000 Subject: [PATCH 01/17] fix(moe_ep): build the NVLink one-sided module for native sm_107a The module inherited the global JIT arch flags, which map SM107 to sm_100f while the bundled CUTLASS lacks Sm107. The driver rejects that code on GR100 (Mercury Validator: Validation failed for SM 107), so the first launch failed with cudaErrorUnknown. The kernels use no CUTLASS, so build every target natively. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- flashinfer/jit/comm.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index e60bdeddbcb..c23a549a7ea 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -385,6 +385,9 @@ def gen_moe_ep_nvlink_one_sided_module() -> JitSpec: """ nv_internal = jit_env.FLASHINFER_CSRC_DIR / "nv_internal" common = nv_internal / "cpp" / "common" + # The kernels use no CUTLASS, so SM107 gets native sm_107a code instead of the + # sm100f fallback; the driver rejects that fallback for these kernels on SM107. + nvcc_flags = current_compilation_context.get_nvcc_flags_list() return gen_jit_spec( "moe_ep_nvlink_one_sided_alltoall", [ @@ -404,7 +407,7 @@ def gen_moe_ep_nvlink_one_sided_module() -> JitSpec: str(nv_internal), str(nv_internal / "include"), ], - extra_cuda_cflags=["-DENABLE_BF16"], + extra_cuda_cflags=nvcc_flags + ["-DENABLE_BF16"], extra_ldflags=["-lcuda"], ) From 84bfa6ae022426d844800917bac04fce3e0e5295 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Sat, 17 Oct 2026 06:25:01 +0000 Subject: [PATCH 02/17] refactor(bench)!: remove the moe_a2a_dispatch_combine routine benchmarks/routines/moe_comm.py timed the functional moe_a2a_* ops with the legacy backend="trtllm" kernels, which no flashinfer.moe_ep backend uses. benchmarks/moe_ep/core/comm/bench_moe_ep_comm.py covers MoE communication and gains its fp8_block_scale input (DeepSeek FP8, one FP32 scale per 128 elements); bench_moe_ep.py covers dispatch + MoE + combine. The routine's per-token LoRA id payload has no counterpart, since MoEEpCommunication.dispatch carries no extra payloads. Also drops the routine's README section, sample test list entries, output columns and the moe_utils helpers only it used. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- benchmarks/README.md | 55 +- benchmarks/flashinfer_benchmark.py | 9 - .../moe_ep/core/comm/bench_moe_ep_comm.py | 14 + .../routines/flashinfer_benchmark_utils.py | 11 - benchmarks/routines/moe_comm.py | 1731 ----------------- benchmarks/routines/moe_utils.py | 199 +- benchmarks/samples/sample_testlist.txt | 14 - 7 files changed, 16 insertions(+), 2017 deletions(-) delete mode 100644 benchmarks/routines/moe_comm.py diff --git a/benchmarks/README.md b/benchmarks/README.md index bd311c8c064..aaf48d1f502 100644 --- a/benchmarks/README.md +++ b/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: not a routine of this script. `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 | diff --git a/benchmarks/flashinfer_benchmark.py b/benchmarks/flashinfer_benchmark.py index 61d8821f032..751c4feb370 100644 --- a/benchmarks/flashinfer_benchmark.py +++ b/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 diff --git a/benchmarks/moe_ep/core/comm/bench_moe_ep_comm.py b/benchmarks/moe_ep/core/comm/bench_moe_ep_comm.py index 23dcd600be1..d22de824f26 100644 --- a/benchmarks/moe_ep/core/comm/bench_moe_ep_comm.py +++ b/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": diff --git a/benchmarks/routines/flashinfer_benchmark_utils.py b/benchmarks/routines/flashinfer_benchmark_utils.py index 4dadfb48507..3791531e22f 100644 --- a/benchmarks/routines/flashinfer_benchmark_utils.py +++ b/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", ], diff --git a/benchmarks/routines/moe_comm.py b/benchmarks/routines/moe_comm.py deleted file mode 100644 index b7c7109ec79..00000000000 --- a/benchmarks/routines/moe_comm.py +++ /dev/null @@ -1,1731 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - -""" -MoE Communication Benchmark Routine - -Deprecated: this routine times the module-level moe_a2a_* ops -(``backend="trtllm"``), which no flashinfer.moe_ep communication backend uses. -Migrate to the flashinfer.moe_ep benchmarks: - -- Dispatch and combine only: ``benchmarks/moe_ep/core/comm/bench_moe_ep_comm.py`` - (nvlink_one_sided, cake and nvlink_two_sided backends). -- Dispatch, MoE computation and combine as one layer (what ``--real_math`` - approximates here): ``benchmarks/moe_ep/bench_moe_ep.py``. - -This module provides benchmarking for MoE All-to-All communication operations -using FlashInfer's functional moe_a2a_* ops. Designed to run with mpirun for -multi-GPU benchmarking. - -Launch examples: - # 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 - 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 - - # With validation (recommended for first run) - 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 - - # With per-phase timing (less accurate but shows dispatch and combine times separately) - 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 \ - --per_phase_timing - -Options: - --quant_dtype fp8 : FP8 (float8_e4m3fn) with float32 per-tensor scale - --quant_dtype nvfp4 : NVFP4 (4-bit) with float8_e4m3fn block scales - --quant_dtype fp8_block_scale : FP8 with block scales (128 elements per block) - --real_math : Run actual MoE kernels (trtllm_fp4/fp8_block_scale_moe) - Supported quant_dtype: nvfp4 and fp8_block_scale only. - --intermediate_size N : FFN intermediate size; must be specified if real_math=True - --validate : Run correctness validation for A2A before benchmarking. - Uses a deterministic fake MoE to verify round-trip - communication. For non-quantized mode, performs exact - comparison. For quantized mode, validates output - shape and numerical validity. - --per_phase_timing : Enable per-phase timing (dispatch/combine). Adds slight - overhead from CUDA events. - This is less accurate than the total timing but shows - dispatch and combine times separately. - --nvtx : Enable NVTX markers for Nsight Systems profiling. -""" - -from collections import defaultdict -from contextlib import contextmanager -from typing import List, Optional, Tuple - -import numpy as np -import torch - -from mpi4py import MPI - -from flashinfer.comm import ( - moe_a2a_combine, - moe_a2a_dispatch, - moe_a2a_get_workspace_size_per_rank, - moe_a2a_initialize, - moe_a2a_wrap_payload_tensor_in_workspace, -) -from flashinfer.comm.mapping import Mapping -from flashinfer.comm.mnnvl import MnnvlMemory -from flashinfer.fused_moe import ( - trtllm_fp4_block_scale_routed_moe, - trtllm_fp8_block_scale_routed_moe, - WeightLayout, -) -from flashinfer.testing.utils import bench_gpu_time - -from .flashinfer_benchmark_utils import ( - dtype_str_to_torch_dtype, - print_perf_metrics, -) -from .moe_utils import ( - add_common_moe_args, - calculate_fp4_global_scale, - quantize_fp4, - dequantize_nvfp4, - quantize_fp8, - quantize_fp8_block_scale, - dequantize_fp8, - dequantize_fp8_block_scale, - pack_topk_ids_triton, - calculate_moe_tflops, - calculate_moe_kernel_bandwidth, - generate_moe_weights, - create_moe_output_scale_scalars, - quantize_and_pack_nvfp4, -) - -# Number of distinct LoRA adapter IDs when --use_lora is enabled. -# Matches issue #3109's "up to 8 concurrent adapter IDs" target. -NUM_LORA_ADAPTERS = 8 - - -class _MoeA2A: - """Dispatch and combine with the functional moe_a2a_* ops on one MNNVL - workspace, keeping the round state that combine needs from dispatch.""" - - def __init__( - self, - mapping: Mapping, - max_num_tokens: int, - top_k: int, - num_experts: int, - workspace_size_per_rank: int, - ): - MnnvlMemory.initialize() - self.ep_rank = mapping.moe_ep_rank - self.ep_size = mapping.moe_ep_size - self.top_k = top_k - self.num_experts = num_experts - self._memory = MnnvlMemory(mapping, workspace_size_per_rank) - self.workspace = self._memory.as_torch_strided_tensor(torch.uint8) - self.metainfo = moe_a2a_initialize( - self.workspace, self.ep_rank, self.ep_size, max_num_tokens - ) - # No peer may publish into this workspace until every rank has cleared - # its own slice. - MnnvlMemory.allocated_map[self._memory.ptr].comm.barrier() - self._recv_view_cache: dict = {} - self._local_num_tokens = 0 - self._combine_offset = 0 - - def dispatch( - self, - token_selected_experts: torch.Tensor, - input_payloads: List[torch.Tensor], - runtime_max_tokens_per_rank: int, - ) -> List[torch.Tensor]: - recv, self._combine_offset, _ = moe_a2a_dispatch( - token_selected_experts, - input_payloads, - self.workspace, - self.metainfo, - runtime_max_tokens_per_rank, - self.ep_rank, - self.ep_size, - self.top_k, - self.num_experts, - recv_view_cache=self._recv_view_cache, - ) - self._local_num_tokens = token_selected_experts.size(0) - return recv - - def get_combine_payload_tensor_in_workspace( - self, runtime_max_tokens_per_rank: int, hidden_size: int, dtype: torch.dtype - ) -> torch.Tensor: - start = self._combine_offset - end = ( - start - + self.ep_size * runtime_max_tokens_per_rank * hidden_size * dtype.itemsize - ) - return moe_a2a_wrap_payload_tensor_in_workspace( - self.workspace[self.ep_rank, :], - [self.ep_size, runtime_max_tokens_per_rank], - start, - end, - dtype, - ) - - def combine( - self, - payload: torch.Tensor, - runtime_max_tokens_per_rank: int, - payload_in_workspace: bool = False, - ) -> torch.Tensor: - return moe_a2a_combine( - payload, - self._local_num_tokens, - self.workspace, - self.metainfo, - runtime_max_tokens_per_rank, - self.ep_rank, - self.ep_size, - self.top_k, - self._combine_offset, - payload_in_workspace, - ) - - -@contextmanager -def cuda_event_timer(events_list: list, enabled: bool = True): - """ - Context manager for deferred CUDA event timing. - - Records start/end events and appends them to events_list. - Does NOT synchronize - caller must sync before reading elapsed times. - - Args: - events_list: List to append (start, end) event tuple to. - enabled: If False, skip event recording (no-op). - - Example: - dispatch_events = [] - for _ in range(num_iters): - with cuda_event_timer(dispatch_events): - moe_a2a.dispatch(...) - - torch.cuda.synchronize() # Single sync at end - times = [s.elapsed_time(e) for s, e in dispatch_events] - """ - if not enabled: - yield - return - - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - try: - yield - finally: - end.record() - events_list.append((start, end)) - - -@contextmanager -def nvtx_range(name: str, enabled: bool = True): - if enabled: - torch.cuda.nvtx.range_push(name) - try: - yield - finally: - if enabled: - torch.cuda.nvtx.range_pop() - - -def print_ordered(comm, rank: int, msg, condition: bool = True): - """ - Print messages one rank at a time in order, synchronized with MPI barriers. - - Args: - comm: MPI communicator - rank: Current rank - msg: string to print. - condition: If False, skip printing but still participate in barriers. - """ - for i in range(comm.Get_size()): - if i == rank and condition: - print(msg, flush=True) - comm.Barrier() - - -def run_moe_comm_test(args): - """ - Run a MoE communication test. - - Args: - args: Parsed command line arguments containing test configuration - - Returns: - list: List of dictionaries containing performance results - """ - if args.routine == "moe_a2a_dispatch_combine": - if MPI.COMM_WORLD.Get_rank() == 0: - print( - "[WARNING] moe_a2a_dispatch_combine is deprecated. Use " - "benchmarks/moe_ep/core/comm/bench_moe_ep_comm.py for dispatch and " - "combine, or benchmarks/moe_ep/bench_moe_ep.py for dispatch, MoE " - "computation and combine." - ) - return test_moe_a2a_dispatch_combine(args) - else: - raise ValueError(f"Unsupported routine: {args.routine}") - - -def parse_moe_comm_args(line, parser): - """ - Parse command line arguments for MoE communication test configuration. - - Args: - line: Command line arguments - parser: ArgumentParser object already populated with shared arguments - - Returns: - Parsed argument namespace - """ - # Parse num_tokens/hidden_size/num_experts/top_k/input_dtype in add_common_moe_args - add_common_moe_args(parser) - parser.add_argument( - "--quant_dtype", - type=str, - required=False, - default=None, - choices=["fp8", "nvfp4", "fp8_block_scale"], - help="Quantization format for hidden states. If set, hidden states are quantized and scale factors are communicated. fp8_block_scale: FP8 with block scales (128 elements per block).", - ) - parser.add_argument( - "--real_math", - action="store_true", - help="Runs actual MoE kernels (trtllm_(fp4|fp8)_block_scale_moe).", - ) - parser.add_argument( - "--intermediate_size", - type=int, - required=False, - default=None, - help="Intermediate size for each expert in MoE. Must be specified if real_math=True.", - ) - parser.add_argument( - "--max_num_tokens", - type=int, - required=False, - default=None, - help="Max tokens per rank for workspace allocation. Defaults to num_tokens.", - ) - parser.add_argument( - "--validate", - action="store_true", - help="Run correctness validation before benchmarking. Uses a deterministic fake MoE to verify round-trip communication.", - ) - parser.add_argument( - "--nvtx", - action="store_true", - help="Enable NVTX markers for Nsight Systems profiling. This also turns on --use_cuda_events.", - ) - parser.add_argument( - "--per_phase_timing", - action="store_true", - help="Enable per-phase timing (dispatch/combine). Adds slight overhead from CUDA events.", - ) - parser.add_argument( - "--use_lora", - action="store_true", - help="Enable per-token LoRA adapter ID payload.", - ) - - args = parser.parse_args(line) - - # Default max_num_tokens to num_tokens if not specified - if args.max_num_tokens is None: - args.max_num_tokens = args.num_tokens - - if args.real_math: - # Must specify intermediate_size if real_math=True - assert args.intermediate_size is not None, ( - "intermediate_size must be specified if real_math=True" - ) - # Must specify quant_dtype as one of the following: nvfp4, fp8_block_scale - # Other quant_dtype support is TBD - assert args.quant_dtype in [ - "nvfp4", - "fp8_block_scale", - ], ( - f"real_math=True requires quant_dtype 'nvfp4' or 'fp8_block_scale', got '{args.quant_dtype}'" - ) - - # Derive scale_dtype from quant_dtype - if args.quant_dtype == "nvfp4": - args.scale_dtype = torch.float8_e4m3fn - elif args.quant_dtype == "fp8": - args.scale_dtype = torch.float32 - elif args.quant_dtype == "fp8_block_scale": - args.scale_dtype = torch.float32 # Block scales are float32 - else: - args.scale_dtype = None - - if args.verbose >= 1: - print(f"[INFO] {args = }") - - return args - - -def _setup_mpi_and_device() -> Tuple[MPI.Comm, int, int, int]: - """ - Setup MPI communicator and CUDA device based on local rank. - - Returns: - Tuple of (comm, rank, world_size, local_rank) - """ - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - world_size = comm.Get_size() - - # Get local rank within the node for device assignment - node_comm = comm.Split_type(MPI.COMM_TYPE_SHARED) - local_rank = node_comm.Get_rank() - torch.cuda.set_device(local_rank) - - return comm, rank, world_size, local_rank - - -def _init_moe_weights( - num_experts_local: int, - hidden_size: int, - intermediate_size: int, - quant_dtype: str, - device: torch.device, -) -> dict: - """ - Initialize MoE weights for MoE. - The function does not do weight shuffling required for TRTLLM-Gen MoE as it serves benchmark purposes only. - However, A2A with fake MoE kernel can be validated through flags (--validate). - - Args: - num_experts_local: Number of local experts on this rank - hidden_size: Hidden dimension size - intermediate_size: Intermediate FFN size - quant_dtype: "nvfp4" or "fp8_block_scale" - device: CUDA device - - Returns: - Dictionary containing weights and scales for MoE computation - """ - weights = {} - - # Create quantized weights - # Note: Generate and quantize one expert at a time to avoid OOM with large expert counts - # gemm1: [num_experts, 2*intermediate_size, hidden_size] - # gemm2: [num_experts, hidden_size, intermediate_size] - if quant_dtype == "nvfp4": - # Create FP4 quantized weights - - # Quantize to FP4 using swizzled layout for weights - sf_vec_size = 16 - use_ue8m0 = False - is_sf_swizzled_layout = True - - gemm1_fp4_list = [] - gemm1_sf_list = [] - gemm1_global_sf_list = [] - gemm2_fp4_list = [] - gemm2_sf_list = [] - gemm2_global_sf_list = [] - for _ in range(num_experts_local): - # Generate bf16 weights for this expert using shared utility - w1_batch, w2_batch = generate_moe_weights( - 1, hidden_size, intermediate_size, device, dtype=torch.bfloat16 - ) - expert_w1_bf16 = w1_batch.squeeze(0) - expert_w2_bf16 = w2_batch.squeeze(0) - del w1_batch, w2_batch - - # Quantize gemm1 weights using moe_utils.quantize_fp4 - quantized, sf, global_sf = quantize_fp4( - expert_w1_bf16, - global_scale=None, - use_ue8m0=use_ue8m0, - is_sf_swizzled_layout=is_sf_swizzled_layout, - ) - gemm1_fp4_list.append(quantized.view(torch.uint8)) - gemm1_sf_list.append(sf.view(torch.float8_e4m3fn)) - gemm1_global_sf_list.append(global_sf) - del expert_w1_bf16 - - # Quantize gemm2 weights using moe_utils.quantize_fp4 - quantized, sf, global_sf = quantize_fp4( - expert_w2_bf16, - global_scale=None, - use_ue8m0=use_ue8m0, - is_sf_swizzled_layout=is_sf_swizzled_layout, - ) - - # NOTE: the script chooses not to do weight shuffling as it is intended for benchmarks; - # only A2A with fake MoE kernel is validated - - gemm2_fp4_list.append(quantized.view(torch.uint8)) - gemm2_sf_list.append(sf.view(torch.float8_e4m3fn)) - gemm2_global_sf_list.append(global_sf) - del expert_w2_bf16 - - # Stack and reshape - weights["gemm1_weights"] = torch.stack(gemm1_fp4_list).reshape( - num_experts_local, 2 * intermediate_size, hidden_size // 2 - ) - weights["gemm1_weights_scale"] = torch.stack(gemm1_sf_list).reshape( - num_experts_local, 2 * intermediate_size, hidden_size // sf_vec_size - ) - weights["gemm2_weights"] = torch.stack(gemm2_fp4_list).reshape( - num_experts_local, hidden_size, intermediate_size // 2 - ) - weights["gemm2_weights_scale"] = torch.stack(gemm2_sf_list).reshape( - num_experts_local, hidden_size, intermediate_size // sf_vec_size - ) - - # Scale scalars for output using shared utility - ( - weights["output1_scale_scalar"], - weights["output1_scale_gate_scalar"], - weights["output2_scale_scalar"], - ) = create_moe_output_scale_scalars(num_experts_local, device) - - elif quant_dtype == "fp8_block_scale": - # Create FP8 block-scaled weights - - # Optionally shuffle weights using shared utility - gemm1_weights = [] - gemm2_weights = [] - for _ in range(num_experts_local): - # Generate bf16 weights for this expert using shared utility - w1_batch, w2_batch = generate_moe_weights( - 1, hidden_size, intermediate_size, device, dtype=torch.bfloat16 - ) - expert_w1_bf16 = w1_batch.squeeze(0) - expert_w2_bf16 = w2_batch.squeeze(0) - del w1_batch, w2_batch - - expert_w1_fp8 = expert_w1_bf16.to(torch.float8_e4m3fn) - expert_w2_fp8 = expert_w2_bf16.to(torch.float8_e4m3fn) - del expert_w1_bf16, expert_w2_bf16 # Free memory immediately - - # NOTE: the script chooses not to do weight shuffling as it is intended for benchmarks; - # only A2A with fake MoE kernel is validated - - gemm1_weights.append(expert_w1_fp8) - gemm2_weights.append(expert_w2_fp8) - del expert_w1_fp8, expert_w2_fp8 # Free memory immediately - weights["gemm1_weights"] = torch.stack(gemm1_weights) - weights["gemm2_weights"] = torch.stack(gemm2_weights) - - # Block scales: [num_experts, out_dim // 128, in_dim // 128] - weights["gemm1_weights_scale"] = 2.0 * torch.ones( - (num_experts_local, 2 * intermediate_size // 128, hidden_size // 128), - device=device, - dtype=torch.float32, - ) - weights["gemm2_weights_scale"] = 2.0 * torch.ones( - (num_experts_local, hidden_size // 128, intermediate_size // 128), - device=device, - dtype=torch.float32, - ) - - else: - raise ValueError(f"Unsupported quant_dtype for real computation: {quant_dtype}") - - return weights - - -def _create_moe_inputs( - num_tokens: int, - hidden_size: int, - num_experts: int, - top_k: int, - input_dtype: torch.dtype, - quant_dtype: Optional[str], - device: torch.device, - comm: MPI.Comm, - use_lora: bool = False, -) -> Tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - List[torch.Tensor], -]: - """ - Create input tensors for MoE A2A benchmark. - - Args: - num_tokens: Number of tokens - hidden_size: Hidden dimension size - num_experts: Total number of experts - top_k: Number of experts per token - input_dtype: Data type for hidden states (before quantization) - quant_dtype: None, "fp8", "nvfp4", or "fp8_block_scale" - device: CUDA device - comm: MPI communicator for syncing global scale - use_lora: If True, append per-token LoRA adapter ID as an extra dispatch payload - - Returns: - Tuple of (hidden_states, hidden_states_original, token_selected_experts, token_final_scales, scale_factor, global_scale, lora_ids, input_payloads) - - hidden_states: The tensor to communicate (may be quantized) - - hidden_states_original: The original unquantized hidden states for validation purpose - - token_selected_experts: Expert indices - - token_final_scales: Routing weights - - scale_factor: Block scale factor tensor for NVFP4 (in A2A payloads), None otherwise - - global_scale: Global/per-tensor scale factor for quantized dtype (synced via MPI max, same across ranks), None otherwise - - lora_ids: [num_tokens] int32 per-token LoRA adapter IDs when use_lora=True, None otherwise - - input_payloads: List of all payloads for dispatch - """ - # Generate original hidden states in input_dtype - hidden_states_original = torch.randn( - num_tokens, hidden_size, dtype=input_dtype, device=device - ) - - # Expert selection: random experts for each token - token_selected_experts = torch.stack( - [ - torch.randperm(num_experts, dtype=torch.int32, device=device)[:top_k] - for _ in range(num_tokens) - ], - dim=0, - ).contiguous() - - # Routing weights - token_final_scales = torch.rand( - num_tokens, top_k, dtype=torch.float32, device=device - ) - - # Handle quantization - # Global scale is synced via MPI max reduction to mimic the fact that it is part of model ckpt - scale_factor = None - global_scale = None - if quant_dtype == "nvfp4": - # Compute local global scale, sync via max, then quantize - local_global_scale = calculate_fp4_global_scale(hidden_states_original) - synced_global_scale = comm.allreduce( - local_global_scale.cpu().item(), op=MPI.MAX - ) - global_scale = torch.tensor( - synced_global_scale, dtype=torch.float32, device=device - ) - hidden_states, scale_factor, global_scale = quantize_and_pack_nvfp4( - hidden_states_original, - global_scale, - use_ue8m0=False, - is_sf_swizzled_layout=False, - ) - elif quant_dtype == "fp8": - # Compute local amax, sync via max, then quantize with synced scale - fp8_max = torch.finfo(torch.float8_e4m3fn).max - local_amax = hidden_states_original.abs().max().float().item() - synced_amax = comm.allreduce(local_amax, op=MPI.MAX) - synced_scale = torch.tensor( - synced_amax / fp8_max, dtype=torch.float32, device=device - ) - hidden_states, global_scale = quantize_fp8(hidden_states_original, synced_scale) - elif quant_dtype == "fp8_block_scale": - # FP8 with block scales (128 elements per block) - # Block scales shape: [hidden_size // 128, num_tokens] - hidden_states, scale_factor = quantize_fp8_block_scale( - hidden_states_original, block_size=128 - ) - # Transpose scale_factor to [num_tokens, hidden_size // 128] for A2A payload - # A2A expects [num_tokens, *] shape for payloads - scale_factor = scale_factor.transpose(0, 1).contiguous() - else: - # No quantization - hidden_states = hidden_states_original - - # Build payload list for dispatch - # Base payloads: hidden states, expert IDs, routing weights - input_payloads = [hidden_states, token_selected_experts, token_final_scales] - - # For post-quant communication: include block scale factors so experts can dequantize - if scale_factor is not None: - input_payloads.append(scale_factor) - - # Per-token LoRA adapter IDs - lora_ids = None - if use_lora: - lora_ids = torch.randint( - 0, NUM_LORA_ADAPTERS, (num_tokens,), dtype=torch.int32, device=device - ) - # Dispatch kernel requires 2D payloads [num_tokens, elements_per_token] - input_payloads.append(lora_ids.unsqueeze(-1)) - - return ( - hidden_states, - hidden_states_original, - token_selected_experts, - token_final_scales, - scale_factor, - global_scale, - lora_ids, - input_payloads, - ) - - -def _calculate_exact_comm_traffic( - all_token_selected_experts: List[np.ndarray], - num_experts: int, - ep_size: int, - hidden_size: int, - input_dtype: torch.dtype, - quant_dtype: Optional[str] = None, - use_lora: bool = False, -) -> Tuple[int, int]: - """ - Calculate exact inter-rank traffic from actual expert assignments. - - For MoE A2A, each token is sent to the rank that owns the selected expert. - Assuming that expert `e` is owned by rank `e // num_experts_per_rank`. - - If a token selects multiple experts on the same remote rank, - the data should be sent only once and counted accordingly. - - Args: - all_token_selected_experts: List of expert assignments from all ranks [ep_size][num_tokens, top_k] - num_experts: Total number of experts - ep_size: Expert parallel size (number of ranks) - hidden_size: Hidden dimension size - input_dtype: Data type of hidden states (before quantization) - quant_dtype: None, "fp8", or "nvfp4" - use_lora: If True, include per-token int32 LoRA ID bytes in dispatch traffic - - Returns: - Tuple of (dispatch_bytes, combine_bytes) for actual inter-rank traffic - """ - num_experts_per_rank = num_experts // ep_size - - # Calculate per-element sizes based on quant_dtype - if quant_dtype == "nvfp4": - # NVFP4: 0.5 bytes per element + block scales - hidden_bytes_per_token = hidden_size // 2 - scale_bytes_per_token = (hidden_size // 16) * 1 # float8_e4m3fn - elif quant_dtype == "fp8": - # FP8: 1 byte per element - hidden_bytes_per_token = hidden_size * 1 - scale_bytes_per_token = 0 - else: - # No quantization - element_size = torch.tensor([], dtype=input_dtype).element_size() - hidden_bytes_per_token = hidden_size * element_size - scale_bytes_per_token = 0 - - # Per-token LoRA adapter ID (int32) travels alongside hidden states (once per src→dst pair) - lora_bytes_per_token = 4 if use_lora else 0 - - # Activation dtype element size for combine phase - combine_element_bytes = torch.tensor([], dtype=input_dtype).element_size() - # Expert IDs/scales are sent for each selection - token_topk_id_and_weight_bytes = 4 + 4 - - # Count unique inter-rank transfers - total_dispatch_bytes = 0 - total_combine_bytes = 0 - - for src_rank, experts_on_rank in enumerate(all_token_selected_experts): - # experts_on_rank: [num_tokens, top_k] - num_tokens = experts_on_rank.shape[0] - top_k = experts_on_rank.shape[1] - - for token_idx in range(num_tokens): - # Find unique destination ranks for this token - dst_ranks_for_token = set() - num_expert_ids_per_dst = {} # Count expert selections per dst rank - - for k in range(top_k): - expert_id = experts_on_rank[token_idx, k] - dst_rank = expert_id // num_experts_per_rank - - if dst_rank != src_rank: - dst_ranks_for_token.add(dst_rank) - num_expert_ids_per_dst[dst_rank] = ( - num_expert_ids_per_dst.get(dst_rank, 0) + 1 - ) - - # For each unique dst_rank, token is sent once - for dst_rank in dst_ranks_for_token: - num_experts_to_dst = num_expert_ids_per_dst[dst_rank] - - # Dispatch: hidden states sent once - total_dispatch_bytes += ( - hidden_bytes_per_token - + scale_bytes_per_token - + lora_bytes_per_token - + num_experts_to_dst * token_topk_id_and_weight_bytes - ) - # Combine: one output per token per dst_rank - total_combine_bytes += hidden_size * combine_element_bytes - - return total_dispatch_bytes, total_combine_bytes - - -def _calculate_comm_bandwidth( - num_tokens: int, - hidden_size: int, - top_k: int, - ep_size: int, - time_ms: float, - input_dtype: torch.dtype, - quant_dtype: Optional[str] = None, - phase: str = "dispatch_combine", - actual_traffic: Optional[Tuple[int, int]] = None, - use_lora: bool = False, -) -> float: - """ - Calculate memory bandwidth for MoE A2A communication in TB/sec. - - Args: - num_tokens: Number of tokens per rank - hidden_size: Hidden dimension size - top_k: Number of experts per token - ep_size: Expert parallel size (number of ranks) - time_ms: Time in milliseconds - input_dtype: Data type of hidden states (before quantization) - quant_dtype: None, "fp8", "nvfp4", or "fp8_block_scale" - phase: "dispatch", "combine", or "dispatch_combine" - actual_traffic: Optional tuple of (dispatch_bytes, combine_bytes) from actual routing. - If provided, uses exact traffic instead of uniform distribution estimate. - - Returns: - Bandwidth in TB/sec - """ - if actual_traffic is not None: - # Use actual traffic from expert routing analysis - dispatch_bytes, combine_bytes = actual_traffic - if phase == "dispatch": - total_bytes = dispatch_bytes - elif phase == "combine": - total_bytes = combine_bytes - else: # dispatch_combine - total_bytes = dispatch_bytes + combine_bytes - else: - # Estimate assuming uniform distribution (fallback) - if quant_dtype == "nvfp4": - # NVFP4: 0.5 bytes per element (packed uint8) + block scales - hidden_states_bytes = num_tokens * hidden_size // 2 # packed FP4 - # Block scales: float8_e4m3fn, one per 16 elements - scale_bytes = num_tokens * (hidden_size // 16) * 1 # float8_e4m3fn = 1 byte - elif quant_dtype == "fp8": - # FP8: 1 byte per element, no scale payload in A2A - hidden_states_bytes = num_tokens * hidden_size * 1 # float8_e4m3fn = 1 byte - scale_bytes = 0 - elif quant_dtype == "fp8_block_scale": - # FP8 with block scales: 1 byte per element + block scales (128 elements per block) - hidden_states_bytes = num_tokens * hidden_size * 1 # float8_e4m3fn = 1 byte - # Block scales: float32, one per 128 elements - scale_bytes = num_tokens * (hidden_size // 128) * 4 # float32 = 4 bytes - else: - # No quantization - element_size = torch.tensor([], dtype=input_dtype).element_size() - hidden_states_bytes = num_tokens * hidden_size * element_size - scale_bytes = 0 - - # Dispatch phase: send hidden_states, expert_ids, scales, and quant scales - dispatch_bytes = ( - hidden_states_bytes - + num_tokens * top_k * 4 # token_selected_experts (int32) - + num_tokens * top_k * 4 # token_final_scales (float32) - + scale_bytes - + (num_tokens * 4 if use_lora else 0) # LoRA adapter ID payload - ) - - # Combine phase: receive processed hidden_states back - # Combine ALWAYS uses activation dtype (bfloat16/float16), not quantized format - element_size = torch.tensor([], dtype=input_dtype).element_size() - combine_bytes = num_tokens * hidden_size * element_size - - if phase == "dispatch": - total_bytes = dispatch_bytes - elif phase == "combine": - total_bytes = combine_bytes - else: # dispatch_combine - total_bytes = dispatch_bytes + combine_bytes - - # Account for multi-rank communication: assume uniform distribution - total_bytes *= (ep_size - 1) / ep_size if ep_size > 1 else 1 - - tb_per_sec = total_bytes / (time_ms * 1e-3) / 1e12 - return tb_per_sec - - -# Copied from tests/comm/test_trtllm_moe_alltoall.py -def fake_moe( - hidden_states: torch.Tensor, - token_selected_experts: torch.Tensor, - num_experts: int, - is_ep: bool = False, - ep_rank: Optional[int] = None, - num_experts_per_rank: Optional[int] = None, - lora_ids: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Apply a deterministic fake MoE transformation for validation. - - Each expert applies a predictable scale: (expert_id + 1.0) / num_experts + 0.5 - When lora_ids is provided, the LoRA adapter ID adds an integer step: scale += lora_id + 1.0 - This allows verifying that communication correctly routes tokens to experts - and combines results. - - - Args: - hidden_states: Input tensor [num_tokens, hidden_size] or [world_size, num_tokens, hidden_size] - token_selected_experts: Expert assignments [num_tokens, top_k] or [world_size, num_tokens, top_k] - num_experts: Total number of experts - is_ep: If True, only process experts assigned to this rank - ep_rank: Rank for expert parallel filtering - num_experts_per_rank: Number of experts per rank - lora_ids: Per-token LoRA adapter IDs [num_tokens] int32, or None - - Returns: - Processed tensor with same shape as hidden_states - """ - target_shape = hidden_states.shape - hidden_states = hidden_states.flatten(end_dim=-2) - token_selected_experts = token_selected_experts.flatten(end_dim=-2) - num_tokens, _ = hidden_states.shape - _, top_k = token_selected_experts.shape - - if is_ep: - assert ep_rank is not None and num_experts_per_rank is not None - - # Initialize output - processed_states = torch.zeros_like(hidden_states) - - # Process each token - for token_idx in range(num_tokens): - results = [] - for k in range(top_k): - expert_id = token_selected_experts[token_idx, k].item() - if is_ep and not ( - ep_rank * num_experts_per_rank - <= expert_id - < (ep_rank + 1) * num_experts_per_rank - ): - continue - - # Deterministic scale based on expert_id - scale = (expert_id + 1.0) / num_experts + 0.5 - if lora_ids is not None: - scale += lora_ids[token_idx].item() + 1.0 - results.append(hidden_states[token_idx].to(torch.float32) * scale) - - # Sum results with higher precision to match actual implementation - if results: - processed_states[token_idx] = torch.sum( - torch.stack(results, dim=0), dim=0, dtype=torch.float32 - ).to(processed_states.dtype) - - return processed_states.view(target_shape) - - -def _validate_moe_a2a( - moe_a2a: "_MoeA2A", - hidden_states: torch.Tensor, - hidden_states_original: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: torch.Tensor, - input_payloads: List[torch.Tensor], - runtime_max_tokens_per_rank: int, - hidden_size: int, - input_dtype: torch.dtype, - quant_dtype: Optional[str], - global_scale: Optional[torch.Tensor], - num_experts: int, - ep_size: int, - rank: int, - comm, - verbose: int = 0, - lora_ids: Optional[torch.Tensor] = None, -) -> bool: - """ - Validate MoE A2A communication correctness with a round-trip test. - - Runs dispatch -> fake_moe -> combine and compares with local reference. - - Args: - moe_a2a: all-to-all over the benchmark's workspace - hidden_states: Original hidden states (before quantization) - token_selected_experts: Expert assignments - token_final_scales: Routing weights - input_payloads: Payloads for dispatch - runtime_max_tokens_per_rank: Max tokens per rank - hidden_size: Hidden dimension - input_dtype: Data type for hidden states - quant_dtype: Quantization format - global_scale: Per-tensor scale factor for quantized dtype (e.g., fp8 or nvfp4), None otherwise - num_experts: Total number of experts - ep_size: Expert parallel size - rank: Current rank - comm: MPI communicator - verbose: Verbosity level - - Returns: - True if validation passes, False otherwise - """ - num_experts_per_rank = num_experts // ep_size - - # Dispatch phase - recv_tensors = moe_a2a.dispatch( - token_selected_experts, - input_payloads, - runtime_max_tokens_per_rank, - ) - - # Tuck away comm and rank for the print_ordered function - def _invoke_print_ordered(msg, condition=True): - print_ordered(comm, rank, msg, condition) - - if verbose >= 2: - _invoke_print_ordered( - f"[VVERBOSE][VALIDATE][Rank {rank}] hidden_states shape: {hidden_states.shape} [num_tokens, hidden_size]:\n{hidden_states[:8, :5]}\n" - f"[VVERBOSE][VALIDATE][Rank {rank}] token_selected_experts shape: {token_selected_experts.shape} [num_tokens, top_k]:\n{token_selected_experts[:8, :]}" - ) - - # Unpack recv_tensors - recv_hidden = recv_tensors[0] - recv_experts = recv_tensors[1] - _ = recv_tensors[2] # recv_token_final_scales - # When use_lora, the LoRA ID payload is appended after the scale_factor (or directly - # after the 3 base payloads when there is no quant scale factor). - use_lora = lora_ids is not None - has_scale_factor = quant_dtype in ("nvfp4", "fp8_block_scale") - recv_scale_factor = recv_tensors[3] if has_scale_factor else None - lora_id_payload_index = (4 if has_scale_factor else 3) if use_lora else None - recv_lora_ids = recv_tensors[lora_id_payload_index].flatten() if use_lora else None - - # Note: For quantized dispatch, recv_tensors[0] is quantized. - # Per-tensor scale factor is part of model ckpts, not in A2A payloads. - recv_hidden_dequant = torch.zeros( - (ep_size, runtime_max_tokens_per_rank, hidden_size), - dtype=input_dtype, - device=recv_hidden.device, - ) - if quant_dtype == "nvfp4": - for i in range(recv_hidden.shape[0]): - recv_hidden_dequant[i] = dequantize_nvfp4( - recv_hidden[i], - recv_scale_factor[i], - global_scale, - block_size=16, - dtype=input_dtype, - ) - elif quant_dtype == "fp8": - for i in range(recv_hidden.shape[0]): - recv_hidden_dequant[i] = dequantize_fp8( - recv_hidden[i], - global_scale, - dtype=input_dtype, - ) - elif quant_dtype == "fp8_block_scale": - for i in range(recv_hidden.shape[0]): - # Transpose scales from [max_tokens, hidden_size // block_size] to - # [hidden_size // block_size, max_tokens] - scales_transposed = recv_scale_factor[i].transpose(0, 1).contiguous() - recv_hidden_dequant[i] = dequantize_fp8_block_scale( - recv_hidden[i], - scales_transposed, - block_size=128, - dtype=input_dtype, - ) - else: - recv_hidden_dequant = recv_hidden - - if verbose >= 2: - _invoke_print_ordered( - f"[VVERBOSE][VALIDATE][Rank {rank}] recv_hidden shape: {recv_hidden.shape} [ep_size, max_tokens, hidden_size]:\n{recv_hidden[:, :8, :5]}\n" - f"[VVERBOSE][VALIDATE][Rank {rank}] recv_experts shape: {recv_experts.shape} [ep_size, max_tokens, top_k]:\n{recv_experts[:, :8, :]}", - ) - - # Apply fake MoE (each expert scales by (expert_id + 1) / num_experts + 0.5) - processed = fake_moe( - recv_hidden_dequant, - recv_experts, - num_experts, - is_ep=True, - ep_rank=rank, - num_experts_per_rank=num_experts_per_rank, - lora_ids=recv_lora_ids, - ) - - # Get combine payload workspace - combine_payload = moe_a2a.get_combine_payload_tensor_in_workspace( - runtime_max_tokens_per_rank, - hidden_size, - input_dtype, - ) - - combine_payload.copy_(processed) - - # Combine phase - output_tensor = moe_a2a.combine( - combine_payload, - runtime_max_tokens_per_rank, - payload_in_workspace=True, - ) - - # Verify output dtype - assert output_tensor.dtype == input_dtype, "Output dtype mismatch" - - num_tokens = hidden_states_original.shape[0] - - # Compute exact reference using original (unquantized) hidden states - # Gather all hidden_states_original and token_selected_experts across ranks - # Note: numpy doesn't support bfloat16, so convert to float32 for allgather - all_hidden_states = comm.allgather(hidden_states_original.cpu().float().numpy()) - all_token_selected_experts = comm.allgather(token_selected_experts.cpu().numpy()) - - # Stack into global tensors - global_hidden_states = ( - torch.from_numpy(np.concatenate(all_hidden_states, axis=0)) - .to(hidden_states.device) - .to(input_dtype) - ) - global_token_selected_experts = torch.from_numpy( - np.concatenate(all_token_selected_experts, axis=0) - ).to(token_selected_experts.device) - - global_lora_ids = None - if use_lora: - all_lora_ids = comm.allgather(lora_ids.cpu().numpy()) - global_lora_ids = torch.from_numpy(np.concatenate(all_lora_ids, axis=0)).to( - lora_ids.device - ) - - # Compute expected result locally - expected = fake_moe( - global_hidden_states, - global_token_selected_experts, - num_experts, - is_ep=False, - lora_ids=global_lora_ids, - ) - - # Extract this rank's portion - expected_local = expected[rank * num_tokens : (rank + 1) * num_tokens] - - # Print verbose debug info - if verbose >= 2: - _invoke_print_ordered( - f"[VVERBOSE][VALIDATE][Rank {rank}] output_tensor shape: {output_tensor.shape} [num_tokens, hidden_size]:\n{output_tensor[:8, :5]}\n" - f"[VVERBOSE][VALIDATE][Rank {rank}] expected_local shape: {expected_local.shape} [num_tokens, hidden_size]:\n{expected_local[:8, :5]}", - ) - - # Compare tensors with tolerance based on quantization - # FP4 has very limited precision, so we need larger tolerances - if quant_dtype == "nvfp4": - atol, rtol = 2.0, 0.5 # FP4: very loose tolerance due to 4-bit precision - elif quant_dtype == "fp8": - atol, rtol = 0.1, 0.1 # FP8: moderate tolerance - elif quant_dtype == "fp8_block_scale": - atol, rtol = 0.5, 0.1 # FP8 block scale: slightly worse than per-tensor FP8 - else: - atol, rtol = 1e-2, 1e-2 # Non-quantized: tight tolerance - - error_msg = None - try: - torch.testing.assert_close( - output_tensor, - expected_local, - atol=atol, - rtol=rtol, - ) - passed = True - except AssertionError as e: - passed = False - error_msg = str(e) - - # Print errors rank-by-rank; all ranks must participate to avoid deadlock - _invoke_print_ordered( - f"[VALIDATE][Rank {rank}] ERROR: {error_msg}", - condition=(error_msg is not None), - ) - - # All-reduce pass/fail status - all_passed = comm.allreduce(passed, op=MPI.LAND) - - if rank == 0: - if all_passed: - print("[VALIDATE] PASSED: All ranks validated successfully") - else: - print("[VALIDATE] FAILED: Validation errors detected") - - return all_passed - - -def test_moe_a2a_dispatch_combine(args): - """ - Benchmark MoE A2A dispatch + combine cycle. - - This benchmarks the full round-trip communication pattern: - 1. Dispatch: send tokens to expert ranks - 2. Combine: gather processed tokens back - - Args: - args: Parsed command line arguments - - Returns: - list: List of result dictionaries - """ - # Setup MPI - comm, rank, world_size, local_rank = _setup_mpi_and_device() - ep_size = world_size - device = torch.device("cuda") - - if rank == 0 and args.verbose >= 1: - print("[INFO] Running test_moe_a2a_dispatch_combine") - print(f"[INFO] ep_size={ep_size}, rank={rank}") - - # Initialize MNNVL - try: - MnnvlMemory.initialize() - except Exception as e: - if rank == 0: - print(f"[ERROR] MNNVL initialization failed: {e}") - return [] - - # Parse parameters - num_tokens = args.num_tokens - hidden_size = args.hidden_size - num_experts = args.num_experts - top_k = args.top_k - max_num_tokens = args.max_num_tokens - input_dtype = dtype_str_to_torch_dtype(args.input_dtype) - quant_dtype = args.quant_dtype - use_lora = getattr(args, "use_lora", False) - - res = [] - - # Validate parameters - if num_experts % ep_size != 0: - if rank == 0: - print( - f"[ERROR] num_experts ({num_experts}) must be divisible by ep_size ({ep_size})" - ) - return res - - # Create mapping - mapping = Mapping( - rank=rank, - tp_size=ep_size, - moe_ep_size=ep_size, - world_size=world_size, - ) - - # Size the workspace for 16-bit hidden states, int32 expert ids and FP32 - # weights; when use_lora, add an int32 per token for the LoRA ID payload. - extra_payload_bytes = 4 if use_lora else 0 - workspace_size_per_rank = moe_a2a_get_workspace_size_per_rank( - ep_size, - max_num_tokens, - hidden_size * 2 + top_k * 4 + top_k * 4 + extra_payload_bytes, - hidden_size * 2, - ) - moe_a2a = _MoeA2A( - mapping, max_num_tokens, top_k, num_experts, workspace_size_per_rank - ) - - # Synchronize all_num_tokens across ranks - all_num_tokens = comm.allgather(num_tokens) - runtime_max_tokens_per_rank = max(all_num_tokens) - sum_all_num_tokens = sum(all_num_tokens) - - # Create input data - torch.manual_seed(args.random_seed + rank) - ( - hidden_states, - hidden_states_original, - token_selected_experts, - token_final_scales, - scale_factor, - global_scale, - lora_ids, - input_payloads, - ) = _create_moe_inputs( - num_tokens, - hidden_size, - num_experts, - top_k, - input_dtype, - quant_dtype, - device, - comm, - use_lora=use_lora, - ) - - # Run validation if requested - if getattr(args, "validate", False): - if rank == 0: - print("[INFO] Running validation before benchmarking...") - - validation_passed = _validate_moe_a2a( - moe_a2a=moe_a2a, - hidden_states=hidden_states, - hidden_states_original=hidden_states_original, - token_selected_experts=token_selected_experts, - token_final_scales=token_final_scales, - input_payloads=input_payloads, - runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, - hidden_size=hidden_size, - input_dtype=input_dtype, - quant_dtype=quant_dtype, - global_scale=global_scale, - num_experts=num_experts, - ep_size=ep_size, - rank=rank, - comm=comm, - verbose=args.verbose, - lora_ids=lora_ids, - ) - if not validation_passed: - if rank == 0: - print("[ERROR] Validation failed. Aborting benchmark.") - return res - - # Gather expert assignments from all ranks to calculate actual traffic - all_token_selected_experts = comm.allgather(token_selected_experts.cpu().numpy()) - dispatch_bytes, combine_bytes = _calculate_exact_comm_traffic( - all_token_selected_experts, - num_experts, - ep_size, - hidden_size, - input_dtype, - quant_dtype, - use_lora=use_lora, - ) - # Compute total active experts across all ranks - total_active_experts = int( - np.unique(np.concatenate(all_token_selected_experts).flatten()).size - ) - if rank == 0 and args.verbose >= 1: - print( - f"[INFO] Inter-rank traffic: dispatch={dispatch_bytes / 1024**2:.3f} MiB, combine={combine_bytes / 1024**2:.3f} MiB" - ) - - # Storage for per-phase CUDA events to be populated later during benchmark - # Deferred timing: collect events during iterations, compute times after single sync - dispatch_events = [] - combine_events = [] - moe_events = [] # For MoE kernel timing (excluding packing) - enable_nvtx = getattr(args, "nvtx", False) - if enable_nvtx: - # CUPTI complains subscribers when using CUPTI for timing and nsys profiling at the same time - args.use_cuda_events = True - enable_per_phase_timing = getattr(args, "per_phase_timing", False) - enable_real_math = getattr(args, "real_math", False) - intermediate_size = getattr(args, "intermediate_size", None) - num_experts_local = num_experts // ep_size - - if enable_real_math: - assert intermediate_size is not None, ( - "intermediate_size must be specified if -real_math=True" - ) - - # Initialize MoE weights for real computation mode - moe_weights = None - if enable_real_math: - if quant_dtype not in ["nvfp4", "fp8_block_scale"]: - if rank == 0: - print( - f"[ERROR] Real MoE math requires quant_dtype 'nvfp4' or 'fp8_block_scale', got '{quant_dtype}'" - ) - return res - moe_weights = _init_moe_weights( - num_experts_local=num_experts_local, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - quant_dtype=quant_dtype, - device=device, - ) - - # Define benchmark function that accepts tensors as arguments - # This enables automatic buffer rotation by bench_gpu_time - def run_dispatch_combine(sel_experts, *payloads): - # Dispatch phase: send (possibly quantized) hidden states to experts - with ( - nvtx_range("moe_a2a_dispatch", enable_nvtx), - cuda_event_timer(dispatch_events, enable_per_phase_timing), - ): - recv_tensors = moe_a2a.dispatch( - sel_experts, - list(payloads), - runtime_max_tokens_per_rank, - ) - - # Expert processing in benchmark runs either no-op or real MoE kernel depending on --real_math flag - with nvtx_range("moe_compute", enable_nvtx): - combine_payload = moe_a2a.get_combine_payload_tensor_in_workspace( - runtime_max_tokens_per_rank, - hidden_size, - input_dtype, - ) - - if enable_real_math and moe_weights is not None: - # Real computation using actual MoE kernels - # recv_tensors[0]: the received hidden states [ep_size, max_tokens, hidden_size] - # recv_tensors[1]: the received expert IDs [ep_size, max_tokens, top_k] - # recv_tensors[2]: the received token final scales [ep_size, max_tokens] - # recv_tensors[3]: the received scale factor [ep_size, max_tokens, hidden_size // block_size] - recv_hidden = recv_tensors[0] - recv_experts = recv_tensors[1] - recv_token_final_scales = recv_tensors[2] - recv_scale_factor = recv_tensors[3] if len(recv_tensors) > 3 else None - - # Flatten for MoE kernel: [ep_size * max_tokens, hidden_size] - total_tokens = ep_size * runtime_max_tokens_per_rank - - if quant_dtype == "nvfp4": - # Reshape hidden states for FP4 kernel - hidden_flat = recv_hidden.reshape(total_tokens, -1) - # Reshape scale factors - scale_flat = ( - recv_scale_factor.reshape(total_tokens, -1) - if recv_scale_factor is not None - else None - ) - - # Pack expert IDs with actual routing weights using fused Triton kernel - local_expert_offset = rank * num_experts_local - recv_experts_flat = recv_experts.reshape(total_tokens, top_k) - recv_weights_flat = recv_token_final_scales.reshape( - total_tokens, top_k - ) - packed_topk_ids = pack_topk_ids_triton( - recv_experts_flat, - recv_weights_flat, - local_expert_offset, - ) - - # Run block scale routed MoE - with cuda_event_timer(moe_events, enable_per_phase_timing): - trtllm_fp4_block_scale_routed_moe( - topk_ids=packed_topk_ids, - routing_bias=None, - hidden_states=hidden_flat, - hidden_states_scale=scale_flat, - gemm1_weights=moe_weights["gemm1_weights"], - gemm1_weights_scale=moe_weights["gemm1_weights_scale"], - gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, - gemm2_weights=moe_weights["gemm2_weights"], - gemm2_weights_scale=moe_weights["gemm2_weights_scale"], - gemm2_bias=None, - output1_scale_scalar=moe_weights["output1_scale_scalar"], - output1_scale_gate_scalar=moe_weights[ - "output1_scale_gate_scalar" - ], - output2_scale_scalar=moe_weights["output2_scale_scalar"], - num_experts=num_experts_local, - top_k=top_k, - n_group=None, - topk_group=None, - intermediate_size=intermediate_size, - local_expert_offset=0, - local_num_experts=num_experts_local, - routed_scaling_factor=None, - routing_method_type=1, # Renormalize: TopK -> Softmax - output=combine_payload.view(total_tokens, hidden_size), - ) - - elif quant_dtype == "fp8_block_scale": - # Reshape for FP8 block scale kernel - hidden_flat = recv_hidden.reshape(total_tokens, hidden_size) - # Transpose scale for kernel: [hidden_size // 128, total_tokens] - if recv_scale_factor is not None: - scale_flat = ( - recv_scale_factor.reshape(total_tokens, -1) - .transpose(0, 1) - .contiguous() - ) - - # Pack expert IDs with actual routing weights using fused Triton kernel - local_expert_offset = rank * num_experts_local - recv_experts_flat = recv_experts.reshape(total_tokens, top_k) - recv_weights_flat = recv_token_final_scales.reshape( - total_tokens, top_k - ) - packed_topk_ids = pack_topk_ids_triton( - recv_experts_flat, - recv_weights_flat, - local_expert_offset, - ) - - # Convert hidden states to FP8 - hidden_fp8 = hidden_flat.to(torch.float8_e4m3fn) - - # Run block scale routed MoE - with cuda_event_timer(moe_events, enable_per_phase_timing): - trtllm_fp8_block_scale_routed_moe( - topk_ids=packed_topk_ids, - routing_bias=None, - hidden_states=hidden_fp8, - hidden_states_scale=scale_flat, - gemm1_weights=moe_weights["gemm1_weights"], - gemm1_weights_scale=moe_weights["gemm1_weights_scale"], - gemm2_weights=moe_weights["gemm2_weights"], - gemm2_weights_scale=moe_weights["gemm2_weights_scale"], - num_experts=num_experts_local, - top_k=top_k, - n_group=0, - topk_group=0, - intermediate_size=intermediate_size, - local_expert_offset=0, - local_num_experts=num_experts_local, - routed_scaling_factor=None, - routing_method_type=1, # Renormalize: TopK -> Softmax - use_shuffled_weight=False, - weight_layout=int(WeightLayout.MajorK), - enable_pdl=args.enable_pdl, - output=combine_payload.view(total_tokens, hidden_size), - ) - - # Combine phase: gather processed outputs from all ranks - with ( - nvtx_range("moe_a2a_combine", enable_nvtx), - cuda_event_timer(combine_events, enable_per_phase_timing), - ): - output = moe_a2a.combine( - combine_payload, - runtime_max_tokens_per_rank, - payload_in_workspace=True, - ) - - return output - - # Synchronize before benchmarking - comm.Barrier() - torch.cuda.synchronize() - - # Use bench_gpu_time with cold L2 cache - total_times = bench_gpu_time( - fn=run_dispatch_combine, - input_args=(token_selected_experts, *input_payloads), - dry_run_iters=args.dry_run_iters, - repeat_iters=args.num_iters, - sleep_after_run=False, - enable_cupti=args.use_cupti, - # Note: disable use_cuda_graph when per_phase_timing=True, which inserts CUDA events in the middle - use_cuda_graph=(not args.no_cuda_graph and not enable_per_phase_timing), - cold_l2_cache=True, - ) - - num_measure_iters = len(total_times) - - # Compute per-phase times if enabled - if enable_per_phase_timing: - # Per-phase events include dry runs; only use the last num_measure_iters entries - dispatch_events_measure = dispatch_events[-num_measure_iters:] - combine_events_measure = combine_events[-num_measure_iters:] - moe_events_measure = moe_events[-num_measure_iters:] if moe_events else [] - - # Convert events to times (no additional sync needed - bench_gpu_time already synced) - dispatch_times = [s.elapsed_time(e) for s, e in dispatch_events_measure] - combine_times = [s.elapsed_time(e) for s, e in combine_events_measure] - moe_times = ( - [s.elapsed_time(e) for s, e in moe_events_measure] - if moe_events_measure - else [] - ) - else: - dispatch_times = [] - combine_times = [] - moe_times = [] - - # Gather times from all ranks - all_total_times = comm.allgather(total_times) - all_dispatch_times = comm.allgather(dispatch_times) - all_combine_times = comm.allgather(combine_times) - all_moe_times = comm.allgather(moe_times) - - # Compute statistics from rank 0 - if rank == 0: - # Use max time across ranks as the benchmark result - # since communication is synchronous - total_per_iter_max = [ - max(t[i] for t in all_total_times) for i in range(num_measure_iters) - ] - median_time = np.median(total_per_iter_max) - std_time = np.std(total_per_iter_max) - - # Calculate total bandwidth using actual traffic from expert routing - tb_per_sec_total = _calculate_comm_bandwidth( - num_tokens, - hidden_size, - top_k, - ep_size, - median_time, - input_dtype, - quant_dtype, - phase="dispatch_combine", - actual_traffic=(dispatch_bytes, combine_bytes), - ) - - # Per-phase statistics if enabled --per_phase_timing flag - median_time_dispatch, std_time_dispatch = np.nan, np.nan - median_time_combine, std_time_combine = np.nan, np.nan - median_time_moe, std_time_moe = np.nan, np.nan - tb_per_sec_dispatch, tb_per_sec_combine = np.nan, np.nan - tflops_moe, tb_per_sec_moe = np.nan, np.nan - if enable_per_phase_timing: - dispatch_per_iter_max = [ - max(t[i] for t in all_dispatch_times) for i in range(num_measure_iters) - ] - combine_per_iter_max = [ - max(t[i] for t in all_combine_times) for i in range(num_measure_iters) - ] - median_time_dispatch = np.median(dispatch_per_iter_max) - std_time_dispatch = np.std(dispatch_per_iter_max) - median_time_combine = np.median(combine_per_iter_max) - std_time_combine = np.std(combine_per_iter_max) - - # MoE timing is only available when real_math is enabled - if all_moe_times and all_moe_times[0]: - moe_per_iter_max = [ - max(t[i] for t in all_moe_times) for i in range(num_measure_iters) - ] - median_time_moe = np.median(moe_per_iter_max) - std_time_moe = np.std(moe_per_iter_max) - - tb_per_sec_dispatch = _calculate_comm_bandwidth( - num_tokens, - hidden_size, - top_k, - ep_size, - median_time_dispatch, - input_dtype, - quant_dtype, - phase="dispatch", - actual_traffic=(dispatch_bytes, combine_bytes), - ) - tb_per_sec_combine = _calculate_comm_bandwidth( - num_tokens, - hidden_size, - top_k, - ep_size, - median_time_combine, - input_dtype, - quant_dtype, - phase="combine", - actual_traffic=(dispatch_bytes, combine_bytes), - ) - - # Print per-phase metrics - print_perf_metrics( - "a2a_dispatch", - median_time_dispatch, - std_time_dispatch, - torch.nan, - tb_per_sec_dispatch, - ) - # Only print MoE timing when real_math is enabled - if args.real_math: - # This is the total FLOPS of all ranks, not per rank - tflops_moe = calculate_moe_tflops( - sum_all_num_tokens, - hidden_size, - intermediate_size, - num_experts, # Actually not used - top_k, - median_time_moe, - ) - # This is the total bandwidth of all ranks, not per rank - tb_per_sec_moe = calculate_moe_kernel_bandwidth( - sum_all_num_tokens, - hidden_size, - intermediate_size, - num_experts, - top_k, - median_time_moe, - input_dtype, - input_dtype, - input_format=quant_dtype, - weight_format=quant_dtype, - routing_logits_dtype=None, # No routing logits in routed MoE - active_experts=total_active_experts, - ) - print_perf_metrics( - "moe_kernel", - median_time_moe, - std_time_moe, - tflops_moe, - tb_per_sec_moe, - ) - print_perf_metrics( - "a2a_combine", - median_time_combine, - std_time_combine, - torch.nan, - tb_per_sec_combine, - ) - - # Always print total - print_perf_metrics( - "a2a_total", median_time, std_time, torch.nan, tb_per_sec_total - ) - print( - "[INFO] The reported achieved tflops/tb_per_sec is the aggregate FLOPS/bandwidth of all participating ranks based on timing results of rank 0. Could observe rank-to-rank variations." - ) - - if args.output_path is not None: - cur_res = defaultdict(str) - cur_res["routine"] = args.routine - cur_res["median_time"] = median_time - cur_res["std_time"] = std_time - cur_res["dispatch_time"] = median_time_dispatch - cur_res["dispatch_std"] = std_time_dispatch - cur_res["moe_time"] = median_time_moe - cur_res["moe_std"] = std_time_moe - cur_res["moe_tflops"] = tflops_moe - cur_res["moe_tb_per_sec"] = tb_per_sec_moe - cur_res["combine_time"] = median_time_combine - cur_res["combine_std"] = std_time_combine - cur_res["tflops"] = "N/A" - cur_res["tb_per_sec"] = tb_per_sec_total - cur_res["dispatch_tb_sec"] = tb_per_sec_dispatch - cur_res["combine_tb_sec"] = tb_per_sec_combine - cur_res["backend"] = "moe_a2a" - cur_res["num_tokens"] = num_tokens - cur_res["hidden_size"] = hidden_size - cur_res["num_experts"] = num_experts - cur_res["top_k"] = top_k - cur_res["ep_size"] = ep_size - cur_res["input_dtype"] = str(input_dtype) - cur_res["quant_dtype"] = quant_dtype if quant_dtype else "" - cur_res["max_num_tokens"] = max_num_tokens - res.append(cur_res) - - return res diff --git a/benchmarks/routines/moe_utils.py b/benchmarks/routines/moe_utils.py index 082e16366d7..bfbbe32048a 100644 --- a/benchmarks/routines/moe_utils.py +++ b/benchmarks/routines/moe_utils.py @@ -30,8 +30,6 @@ import torch -import triton -import triton.language as tl from flashinfer import fp4_quantize, shuffle_matrix_a from flashinfer.fused_moe import WeightLayout, convert_to_block_layout @@ -317,148 +315,6 @@ def dequantize_fp8( return (tensor_fp8.float() * scale.float()).to(dtype) -def dequantize_fp8_block_scale( - tensor_fp8: torch.Tensor, - block_scales: torch.Tensor, - block_size: int = 128, - dtype: torch.dtype = torch.bfloat16, -) -> torch.Tensor: - """ - Dequantize block-scaled FP8 tensor back to high precision. - - Args: - tensor_fp8: FP8 quantized tensor [num_tokens, hidden_size] - block_scales: Block scales [hidden_size // block_size, num_tokens] - block_size: Number of elements per scale block - dtype: Output dtype - - Returns: - Dequantized tensor in specified dtype - """ - num_tokens, hidden_size = tensor_fp8.shape - num_blocks = hidden_size // block_size - - # Reshape tensor for block-wise dequantization - reshaped = tensor_fp8.float().reshape(num_tokens, num_blocks, block_size) - - # Transpose scales from [num_blocks, num_tokens] to [num_tokens, num_blocks] - block_scales_t = block_scales.transpose(0, 1).contiguous() - - # Apply scales - scales_expanded = block_scales_t.unsqueeze(-1) # [num_tokens, num_blocks, 1] - dequantized = reshaped * scales_expanded - - return dequantized.reshape(num_tokens, hidden_size).to(dtype) - - -@triton.jit -def _pack_topk_ids_kernel( - expert_ids_ptr, # [total_tokens, top_k] int32/int64 - expert_weights_ptr, # [total_tokens, top_k] float32 - output_ptr, # [total_tokens, top_k] int32 - local_expert_offset, # scalar int - stride_ids_row, # stride for expert_ids row dimension - stride_ids_col, # stride for expert_ids col dimension - stride_weights_row, # stride for weights row dimension - stride_weights_col, # stride for weights col dimension - n_rows, - n_cols, - BLOCK_SIZE: tl.constexpr, -): - """ - Fused kernel to pack expert IDs with actual weights into packed format: - packed = ((expert_id - local_offset) << 16) | (weight_as_bf16 bits) - - This eliminates: - - dtype conversion kernel (float32 -> bf16) - - subtraction, shift, view, cast, bitwise_or kernels - All fused into a single kernel. - """ - pid = tl.program_id(0) - - # Calculate row and column from linear index - linear_idx = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - row_idx = linear_idx // n_cols - col_idx = linear_idx % n_cols - - mask = linear_idx < (n_rows * n_cols) - - # Compute actual memory offset using strides for expert_ids - ids_offset = row_idx * stride_ids_row + col_idx * stride_ids_col - weights_offset = row_idx * stride_weights_row + col_idx * stride_weights_col - - # Load expert IDs and compute local IDs - expert_ids = tl.load(expert_ids_ptr + ids_offset, mask=mask) - local_ids = expert_ids - local_expert_offset - - # Load weights as float32 and convert to bf16 bits - weights_f32 = tl.load(expert_weights_ptr + weights_offset, mask=mask) - # Convert to bf16, then reinterpret as int16 - weights_bf16 = weights_f32.to(tl.bfloat16) - weights_int16 = weights_bf16.to(tl.int16, bitcast=True) - weights_int32 = weights_int16.to(tl.int32) & 0xFFFF - - # Pack: (local_id << 16) | weight_bits - packed = (local_ids.to(tl.int32) << 16) | weights_int32 - - # Output is always contiguous - tl.store(output_ptr + linear_idx, packed, mask=mask) - - -def pack_topk_ids_triton( - expert_ids: torch.Tensor, - expert_weights: torch.Tensor, - local_expert_offset: int, - output: torch.Tensor = None, -) -> torch.Tensor: - """ - Pack expert IDs with actual weights into packed format using a fused Triton kernel. - - This fused kernel handles: - - Non-contiguous input tensors via strides - - float32 -> bf16 conversion for weights - - Packing: (expert_id - offset) << 16 | weight_bf16_bits - - Args: - expert_ids: [total_tokens, top_k] expert indices (int32 or int64), can be non-contiguous - expert_weights: [total_tokens, top_k] routing weights (float32), can be non-contiguous - local_expert_offset: offset to subtract from global expert IDs - output: optional pre-allocated output tensor [total_tokens, top_k] int32 - - Returns: - packed_topk_ids: [total_tokens, top_k] int32 where each element is - ((expert_id - offset) << 16) | (weight_bf16 as int16) - """ - assert expert_ids.ndim == 2 - assert expert_weights.ndim == 2 - assert expert_ids.shape == expert_weights.shape - n_rows, n_cols = expert_ids.shape - - if output is None: - output = torch.empty( - n_rows, n_cols, dtype=torch.int32, device=expert_ids.device - ) - - n_elements = n_rows * n_cols - BLOCK_SIZE = 1024 - grid = (triton.cdiv(n_elements, BLOCK_SIZE),) - - _pack_topk_ids_kernel[grid]( - expert_ids, - expert_weights, - output, - local_expert_offset, - expert_ids.stride(0), - expert_ids.stride(1), - expert_weights.stride(0), - expert_weights.stride(1), - n_rows, - n_cols, - BLOCK_SIZE=BLOCK_SIZE, - ) - return output - - def calculate_moe_tflops( num_tokens: int, hidden_size: int, @@ -627,11 +483,7 @@ def add_common_moe_args(parser: argparse.ArgumentParser) -> None: """ Add common MoE CLI arguments to a parser. - This adds arguments shared between moe.py and moe_comm.py: - - num_tokens, hidden_size, num_experts, top_k, input_dtype - - In constrast to moe.py, intermediate_size is only optional for moe_comm.py, - hence not counted as a common argument. + Adds num_tokens, hidden_size, num_experts, top_k and input_dtype. Args: parser: ArgumentParser to add arguments to @@ -727,52 +579,3 @@ def create_moe_output_scale_scalars( ) output2_scale_scalar = torch.ones(num_experts, device=device, dtype=torch.float32) return output1_scale_scalar, output1_scale_gate_scalar, output2_scale_scalar - - -def quantize_and_pack_nvfp4( - tensor: torch.Tensor, - global_scale: Optional[torch.Tensor] = None, - use_ue8m0: bool = False, - is_sf_swizzled_layout: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Quantize tensor to NVFP4 and pack into communication-ready format. - - This wraps quantize_fp4 and performs the reshaping needed for MoE kernels: - - Packs 2 FP4 values into 1 byte (uint8) - - Reshapes block scales appropriately - - Args: - tensor: Input tensor [num_tokens, hidden_size] - global_scale: Optional pre-computed global scale. If None, computed from tensor. - use_ue8m0: Whether to use UE8M0 format - is_sf_swizzled_layout: Whether to use swizzled layout for scale factors. - Use False for activations, True for weights. - - Returns: - Tuple of (quantized_packed, block_scales, global_scale) - - quantized_packed: uint8 tensor [num_tokens, hidden_size // 2] - - block_scales: float8_e4m3fn tensor [num_tokens, hidden_size // 16] - - global_scale: float32 scalar tensor - """ - sf_vec_size = 16 - num_tokens, hidden_size = tensor.shape - - # Quantize using the standard FP4 quantization - quantized, block_scales, global_scale = quantize_fp4( - tensor, global_scale, use_ue8m0, is_sf_swizzled_layout - ) - - # Pack 2 FP4 values into 1 byte - quantized_packed = quantized.view(torch.uint8).reshape(num_tokens, hidden_size // 2) - - # Reshape block scales - block_scales_reshaped = block_scales.view(torch.float8_e4m3fn).reshape( - num_tokens, hidden_size // sf_vec_size - ) - - # Validate scale shape - expected_scale_elems = (num_tokens * hidden_size) // sf_vec_size - assert block_scales_reshaped.numel() == expected_scale_elems, "Invalid scale shape" - - return quantized_packed, block_scales_reshaped, global_scale diff --git a/benchmarks/samples/sample_testlist.txt b/benchmarks/samples/sample_testlist.txt index dd07833c8bb..9f0fd07e9b6 100644 --- a/benchmarks/samples/sample_testlist.txt +++ b/benchmarks/samples/sample_testlist.txt @@ -90,20 +90,6 @@ --routine b12x_fused_moe --num_tokens 1024 --hidden_size 1024 --intermediate_size 2688 --num_experts 512 --top_k 22 -vv --generate_repro_command --case_tag "b12x_moe_swiglu" --routine b12x_fused_moe --num_tokens 1 --hidden_size 1024 --intermediate_size 2688 --num_experts 512 --top_k 22 --activation-type Relu2 -vv --generate_repro_command --case_tag "b12x_moe_relu2_decode" -## MoE Communication (requires mpirun, e.g.: mpirun -np 8 python benchmarks/flashinfer_benchmark.py ...) -# Basic A2A dispatch+combine without quantization -#--routine moe_a2a_dispatch_combine --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 -vv --generate_repro_command --case_tag "moe_a2a_basic" -# With FP8 per-tensor quantization -#--routine moe_a2a_dispatch_combine --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 --quant_dtype fp8 -vv --generate_repro_command --case_tag "moe_a2a_fp8" -# With NVFP4 block-scale quantization -#--routine moe_a2a_dispatch_combine --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 --quant_dtype nvfp4 -vv --generate_repro_command --case_tag "moe_a2a_nvfp4" -# With FP8 block-scale quantization -#--routine moe_a2a_dispatch_combine --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 --quant_dtype fp8_block_scale -vv --generate_repro_command --case_tag "moe_a2a_fp8_block_scale" -# With real MoE kernel (NVFP4) -#--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 --per_phase_timing -vv --generate_repro_command --case_tag "moe_a2a_nvfp4_real_math" -# With real MoE kernel (FP8 block-scale) -#--routine moe_a2a_dispatch_combine --num_tokens 1024 --hidden_size 7168 --num_experts 256 --top_k 8 --quant_dtype fp8_block_scale --real_math --intermediate_size 18432 --per_phase_timing -vv --generate_repro_command --case_tag "moe_a2a_fp8_bs_real_math" - ## RMSNorm # Basic RMSNorm with 2D input shape --routine rmsnorm --batch_size 32 --hidden_size 4096 --input_dtype bfloat16 --refcheck -vv --generate_repro_command --case_tag "rmsnorm_llama_hidden" From 6bf31c02aa3a2d6afbed7fd872f0fa04b9e92e98 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Sat, 17 Oct 2026 06:36:45 +0000 Subject: [PATCH 03/17] refactor(comm)!: remove the backend="trtllm" MoE all-to-all ops; move Cake sources The functional moe_a2a_* ops kept the TensorRT-LLM one-sided kernels as their default backend after flashinfer.moe_ep.NVLinkOneSidedAlltoAll moved to the current TensorRT-LLM kernels. Nothing in FlashInfer uses that backend now: - Remove backend="trtllm" and its kernels (communicationKernels/moeAlltoAllKernels.cu). The ops keep the backend argument with "cake" as its only and default value; backend="trtllm" raises and points to NVLinkOneSidedAlltoAll. The "cake" custom ops keep their per-architecture names. - Move the Cake sources out of nv_internal, following the other Cake ops: csrc/trtllm_moe_alltoall.cu -> csrc/cake_moe_alltoall.cu, moeAlltoAllFusedKernels.cu / moeAlltoAllKernels.h -> csrc/cake_moe_alltoall_dispatch.{cu,cuh}, thop/moeAlltoAllMeta.h -> csrc/cake_moe_alltoall_meta.h, and the generated kernels to csrc/generated/moe_alltoall/. - Delete the seven per-phase sources (moeAlltoAll{PrepareDispatch,Dispatch, StageCombine,PublishCombine,Combine,QuantizeCombine,Sanitize}.cu), which no module compiled. The functional-op tests now run against the Cake kernels and are unchanged apart from the default backend; the adapter tests drop their legacy cases. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- ...m_moe_alltoall.cu => cake_moe_alltoall.cu} | 4 +- ...rnels.cu => cake_moe_alltoall_dispatch.cu} | 2 +- ...rnels.h => cake_moe_alltoall_dispatch.cuh} | 2 +- ...lltoAllMeta.h => cake_moe_alltoall_meta.h} | 0 .../mnnvl_moe_alltoall_sm100.cu | 0 .../mnnvl_moe_alltoall_sm103.cu | 0 .../moeAlltoAllCombine.cu | 270 --- .../moeAlltoAllDispatch.cu | 377 ---- .../moeAlltoAllKernels.cu | 1690 ----------------- .../moeAlltoAllPrepareDispatch.cu | 55 - .../moeAlltoAllPublishCombine.cu | 114 -- .../moeAlltoAllQuantizeCombine.cu | 226 --- .../moeAlltoAllSanitize.cu | 61 - .../moeAlltoAllStageCombine.cu | 65 - flashinfer/comm/trtllm_moe_alltoall.py | 84 +- flashinfer/jit/comm.py | 38 +- tests/comm/moe_a2a_driver.py | 6 +- .../test_mnnvl_moe_alltoall_fused_adapter.py | 157 +- 18 files changed, 104 insertions(+), 3047 deletions(-) rename csrc/{trtllm_moe_alltoall.cu => cake_moe_alltoall.cu} (99%) rename csrc/{nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllFusedKernels.cu => cake_moe_alltoall_dispatch.cu} (99%) rename csrc/{nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h => cake_moe_alltoall_dispatch.cuh} (99%) rename csrc/{nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h => cake_moe_alltoall_meta.h} (100%) rename csrc/generated/{ => moe_alltoall}/mnnvl_moe_alltoall_sm100.cu (100%) rename csrc/generated/{ => moe_alltoall}/mnnvl_moe_alltoall_sm103.cu (100%) delete mode 100644 csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllCombine.cu delete mode 100644 csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllDispatch.cu delete mode 100644 csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu delete mode 100644 csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllPrepareDispatch.cu delete mode 100644 csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllPublishCombine.cu delete mode 100644 csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllQuantizeCombine.cu delete mode 100644 csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllSanitize.cu delete mode 100644 csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllStageCombine.cu diff --git a/csrc/trtllm_moe_alltoall.cu b/csrc/cake_moe_alltoall.cu similarity index 99% rename from csrc/trtllm_moe_alltoall.cu rename to csrc/cake_moe_alltoall.cu index e4e27314310..2d66f46a967 100644 --- a/csrc/trtllm_moe_alltoall.cu +++ b/csrc/cake_moe_alltoall.cu @@ -24,12 +24,12 @@ #include #include +#include "cake_moe_alltoall_dispatch.cuh" +#include "cake_moe_alltoall_meta.h" #include "flashinfer/utils.cuh" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/dataType.h" #include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h" -#include "tensorrt_llm/thop/moeAlltoAllMeta.h" #include "tvm_ffi_utils.h" using tvm::ffi::Array; diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllFusedKernels.cu b/csrc/cake_moe_alltoall_dispatch.cu similarity index 99% rename from csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllFusedKernels.cu rename to csrc/cake_moe_alltoall_dispatch.cu index 87aab4e98cf..791a0c93bd5 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllFusedKernels.cu +++ b/csrc/cake_moe_alltoall_dispatch.cu @@ -9,10 +9,10 @@ #include #include +#include "cake_moe_alltoall_dispatch.cuh" #include "flashinfer/exception.h" #include "tensorrt_llm/common/dataType.h" #include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h" // These entry points are frozen CUDA sources compiled as part of the same JIT // module. Keep this launch-only translation unit free of framework tensor diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/csrc/cake_moe_alltoall_dispatch.cuh similarity index 99% rename from csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h rename to csrc/cake_moe_alltoall_dispatch.cuh index 6ce641ade12..b06304ce6d2 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/csrc/cake_moe_alltoall_dispatch.cuh @@ -254,7 +254,7 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params); void moe_a2a_prepare_combine_launch(MoeA2ACombineParams const& params); // True when combine and output block quantization run as one launch, so the caller must not -// allocate the BF16/FP16 accumulation_data intermediate. The trtllm backend always returns false. +// allocate the BF16/FP16 accumulation_data intermediate. bool moe_a2a_combine_fuses_quantization(MoeA2ACombineParams const& params); // Sanitize expert IDs for invalid tokens diff --git a/csrc/nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h b/csrc/cake_moe_alltoall_meta.h similarity index 100% rename from csrc/nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h rename to csrc/cake_moe_alltoall_meta.h diff --git a/csrc/generated/mnnvl_moe_alltoall_sm100.cu b/csrc/generated/moe_alltoall/mnnvl_moe_alltoall_sm100.cu similarity index 100% rename from csrc/generated/mnnvl_moe_alltoall_sm100.cu rename to csrc/generated/moe_alltoall/mnnvl_moe_alltoall_sm100.cu diff --git a/csrc/generated/mnnvl_moe_alltoall_sm103.cu b/csrc/generated/moe_alltoall/mnnvl_moe_alltoall_sm103.cu similarity index 100% rename from csrc/generated/mnnvl_moe_alltoall_sm103.cu rename to csrc/generated/moe_alltoall/mnnvl_moe_alltoall_sm103.cu diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllCombine.cu b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllCombine.cu deleted file mode 100644 index 847f0464e62..00000000000 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllCombine.cu +++ /dev/null @@ -1,270 +0,0 @@ -/* - * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef signed int int32_t; -typedef short int int16_t; -typedef struct __align__(64) { - uint64_t opaque[16]; -} CUtensorMap; - -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x)); - return result; -} - -#define NUM_MAIN_STAGES 1 -#define THREADS 256 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - -template -__device__ __forceinline__ float reduce_contributions(float (&contributions)[TOP_K]) { - if constexpr (TOP_K == 22) { - contributions[0] += contributions[1]; - contributions[2] += contributions[3]; - contributions[4] += contributions[5]; - contributions[6] += contributions[7]; - contributions[8] += contributions[9]; - contributions[10] += contributions[11]; - contributions[12] += contributions[13]; - contributions[14] += contributions[15]; - contributions[16] += contributions[17]; - contributions[18] += contributions[19]; - contributions[20] += contributions[21]; - contributions[0] += contributions[2]; - contributions[4] += contributions[6]; - contributions[8] += contributions[10]; - contributions[12] += contributions[14]; - contributions[16] += contributions[18]; - contributions[0] += contributions[4]; - contributions[8] += contributions[12]; - contributions[16] += contributions[20]; - contributions[0] += contributions[8]; - contributions[0] += contributions[16]; - } else if constexpr (TOP_K == 16) { - contributions[0] += contributions[1]; - contributions[2] += contributions[3]; - contributions[4] += contributions[5]; - contributions[6] += contributions[7]; - contributions[8] += contributions[9]; - contributions[10] += contributions[11]; - contributions[12] += contributions[13]; - contributions[14] += contributions[15]; - contributions[0] += contributions[2]; - contributions[4] += contributions[6]; - contributions[8] += contributions[10]; - contributions[12] += contributions[14]; - contributions[0] += contributions[4]; - contributions[8] += contributions[12]; - contributions[0] += contributions[8]; - } else if constexpr (TOP_K == 10) { - contributions[0] += contributions[1]; - contributions[2] += contributions[3]; - contributions[4] += contributions[5]; - contributions[6] += contributions[7]; - contributions[8] += contributions[9]; - contributions[0] += contributions[2]; - contributions[4] += contributions[6]; - contributions[0] += contributions[4]; - contributions[0] += contributions[8]; - } else if constexpr (TOP_K == 8) { - contributions[0] += contributions[1]; - contributions[2] += contributions[3]; - contributions[4] += contributions[5]; - contributions[6] += contributions[7]; - contributions[0] += contributions[2]; - contributions[4] += contributions[6]; - contributions[0] += contributions[4]; - } else if constexpr (TOP_K == 6) { - contributions[0] += contributions[1]; - contributions[2] += contributions[3]; - contributions[4] += contributions[5]; - contributions[0] += contributions[2]; - contributions[0] += contributions[4]; - } else if constexpr (TOP_K == 4) { - contributions[0] += contributions[1]; - contributions[2] += contributions[3]; - contributions[0] += contributions[2]; - } else if constexpr (TOP_K == 2) { - contributions[0] += contributions[1]; - } else if constexpr (TOP_K != 1) { -#pragma unroll - for (int route = 1; route < TOP_K; ++route) { - contributions[0] += contributions[route]; - } - } - return contributions[0]; -} - -template -__device__ __forceinline__ void combine_top_k( - uint8_t* __restrict__ workspace, uint8_t* __restrict__ output, - unsigned long long workspace_stride_bytes, unsigned long long topk_target_ranks_offset, - unsigned long long topk_send_indices_offset, unsigned long long combine_payload_offset, - int max_tokens_per_rank, int local_num_tokens, int elements_per_token, - int payload_element_bytes, int payload_dtype_code, int output_dtype_code, int ep_rank, - bool use_low_precision, bool enable_pdl) { - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int bid = blockIdx.x; - - if (enable_pdl) { - asm volatile("griddepcontrol.wait;" ::: "memory"); - } - int* workspace_i32 = reinterpret_cast(workspace); - unsigned long long local_workspace_base = (unsigned long long)ep_rank * workspace_stride_bytes; - int token = bid; - if (token < local_num_tokens) { - unsigned long long target_base = (local_workspace_base + topk_target_ranks_offset) / 4 + - (unsigned long long)token * (unsigned long long)TOP_K; - unsigned long long send_base = (local_workspace_base + topk_send_indices_offset) / 4 + - (unsigned long long)token * (unsigned long long)TOP_K; -#pragma unroll 1 - for (int column = tid; column < elements_per_token; column += 256) { - float contributions[TOP_K]; -#pragma unroll - for (int route = 0; route < TOP_K; ++route) { - contributions[route] = 0.0f; - } -#pragma unroll - for (int route = 0; route < TOP_K; ++route) { - int target_rank = workspace_i32[target_base + (unsigned long long)route]; - int send_index = workspace_i32[send_base + (unsigned long long)route]; - if (target_rank >= 0 && send_index >= 0) { - unsigned long long payload_item = - ((unsigned long long)ep_rank * (unsigned long long)max_tokens_per_rank + - (unsigned long long)send_index) * - (unsigned long long)elements_per_token + - (unsigned long long)column; - unsigned long long byte_index = (unsigned long long)target_rank * workspace_stride_bytes + - combine_payload_offset + - payload_item * (unsigned long long)payload_element_bytes; - float value = 0.0f; - if (payload_dtype_code == 0) { - __nv_bfloat16* source_bf16 = reinterpret_cast<__nv_bfloat16*>(workspace); - value = (float)source_bf16[byte_index / 2]; - } - if (payload_dtype_code == 1) { - __half* source_f16 = reinterpret_cast<__half*>(workspace); - value = (float)source_f16[byte_index / 2]; - } - if (payload_dtype_code == 2) { - uint8_t* source_fp8 = reinterpret_cast(workspace); - float _vec_load_0[1]; - { - uint8_t _fp8_byte_0 = *reinterpret_cast(source_fp8 + byte_index); - uint16_t _e4m3x2_0 = (uint16_t)_fp8_byte_0; - uint32_t _f16x2_0; - asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(_f16x2_0) : "h"(_e4m3x2_0)); - uint16_t _h0_0 = (uint16_t)(_f16x2_0 & 0xFFFFu); - asm volatile("cvt.f32.f16 %0, %1;" : "=f"(_vec_load_0[0]) : "h"(_h0_0)); - } - value = _vec_load_0[0]; - } - if (payload_dtype_code == 3) { - float* source_f32 = reinterpret_cast(workspace); - value = source_f32[byte_index / 4]; - } - float contribution = value; - if (use_low_precision) { - float _fp8_rt_0; - uint16_t _e4m3x2_1; - uint32_t _f16x2_1; - asm volatile("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" - : "=h"(_e4m3x2_1) - : "f"(0.0f), "f"(contribution)); - asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(_f16x2_1) : "h"(_e4m3x2_1)); - uint16_t _fp8_h0_1 = (uint16_t)(_f16x2_1 & 0xFFFFu); - asm volatile("cvt.f32.f16 %0, %1;" : "=f"(_fp8_rt_0) : "h"(_fp8_h0_1)); - contribution = _fp8_rt_0; - } - contributions[route] = contribution; - } - } - float result = reduce_contributions(contributions); - unsigned long long output_item = - (unsigned long long)token * (unsigned long long)elements_per_token + - (unsigned long long)column; - if (output_dtype_code == 0) { - __nv_bfloat16* output_bf16 = reinterpret_cast<__nv_bfloat16*>(output); - output_bf16[output_item] = result; - } - if (output_dtype_code == 1) { - __half* output_f16 = reinterpret_cast<__half*>(output); - output_f16[output_item] = result; - } - if (output_dtype_code == 2) { - uint8_t* output_fp8 = reinterpret_cast(output); - { - unsigned short _fp8_pair; - asm("cvt.rn.satfinite.e4m3x2.f32 %0, 0f00000000, %1;" : "=h"(_fp8_pair) : "f"(result)); - *(reinterpret_cast(output_fp8 + output_item) + (0)) = - (unsigned char)(_fp8_pair & 0xFF); - } - } - if (output_dtype_code == 3) { - float* output_f32 = reinterpret_cast(output); - output_f32[output_item] = result; - } - } - } - if (enable_pdl) { - if (warp == 0) { - if (elect_sync()) { - asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); - } - } - } -} - -#define DEFINE_COMBINE_TOP_K(TOP_K) \ - extern "C" __global__ __launch_bounds__(256) void \ - kernel_flashinfer_mnnvl_moe_alltoall_combine_top_k_##TOP_K( \ - uint8_t* __restrict__ workspace, uint8_t* __restrict__ output, \ - unsigned long long workspace_stride_bytes, unsigned long long topk_target_ranks_offset, \ - unsigned long long topk_send_indices_offset, unsigned long long combine_payload_offset, \ - int max_tokens_per_rank, int local_num_tokens, int elements_per_token, \ - int payload_element_bytes, int payload_dtype_code, int output_dtype_code, int ep_rank, \ - bool use_low_precision, bool enable_pdl) { \ - combine_top_k(workspace, output, workspace_stride_bytes, topk_target_ranks_offset, \ - topk_send_indices_offset, combine_payload_offset, max_tokens_per_rank, \ - local_num_tokens, elements_per_token, payload_element_bytes, \ - payload_dtype_code, output_dtype_code, ep_rank, use_low_precision, \ - enable_pdl); \ - } - -DEFINE_COMBINE_TOP_K(1) -DEFINE_COMBINE_TOP_K(2) -DEFINE_COMBINE_TOP_K(4) -DEFINE_COMBINE_TOP_K(6) -DEFINE_COMBINE_TOP_K(8) -DEFINE_COMBINE_TOP_K(10) -DEFINE_COMBINE_TOP_K(12) -DEFINE_COMBINE_TOP_K(14) -DEFINE_COMBINE_TOP_K(16) -DEFINE_COMBINE_TOP_K(18) -DEFINE_COMBINE_TOP_K(22) - -#undef DEFINE_COMBINE_TOP_K diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllDispatch.cu b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllDispatch.cu deleted file mode 100644 index ccdc8072a72..00000000000 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllDispatch.cu +++ /dev/null @@ -1,377 +0,0 @@ -/* - * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef signed int int32_t; -typedef short int int16_t; -typedef struct __align__(64) { - uint64_t opaque[16]; -} CUtensorMap; - -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x)); - return result; -} - -#define NUM_MAIN_STAGES 1 -#define SMEM_SMEM_TARGET_RANKS_OFF 0 -#define SMEM_SMEM_TARGET_RANKS_STAGE_BYTES 88 -#define SMEM_SMEM_TARGET_RANKS_STRIDE 88 -#define SMEM_SMEM_SEND_INDICES_OFF 88 -#define SMEM_SMEM_SEND_INDICES_STAGE_BYTES 88 -#define SMEM_SMEM_SEND_INDICES_STRIDE 88 -#define SMEM_TOTAL 256 -#define THREADS 256 - -#include - -extern "C" { - -__global__ __launch_bounds__(256) void kernel_flashinfer_mnnvl_moe_alltoall_dispatch( - int* __restrict__ token_selected_experts, uint8_t* __restrict__ payload_0, - uint8_t* __restrict__ payload_1, uint8_t* __restrict__ payload_2, - uint8_t* __restrict__ payload_3, uint8_t* __restrict__ payload_4, - uint8_t* __restrict__ payload_5, uint8_t* __restrict__ workspace, - int* __restrict__ eplb_local_stats, unsigned long long workspace_stride_bytes, - unsigned long long flag_val_offset, unsigned long long local_token_counter_offset, - unsigned long long send_counters_offset, unsigned long long recv_counters_offset, - unsigned long long completion_flags_offset, unsigned long long topk_target_ranks_offset, - unsigned long long topk_send_indices_offset, unsigned long long eplb_gathered_stats_offset, - unsigned long long payload_0_offset, unsigned long long payload_1_offset, - unsigned long long payload_2_offset, unsigned long long payload_3_offset, - unsigned long long payload_4_offset, unsigned long long payload_5_offset, int payload_0_bytes, - int payload_1_bytes, int payload_2_bytes, int payload_3_bytes, int payload_4_bytes, - int payload_5_bytes, int num_payloads, int max_tokens_per_rank, int local_num_tokens, - int ep_rank, int ep_size, int num_experts, int top_k, int eplb_stats_num_experts, - bool enable_pdl, bool enable_eplb, bool enable_rank_mask, unsigned long long active_rank_mask) { - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - // Kernel setup ops - int* smem_target_ranks = reinterpret_cast(smem_raw + 0); - const int smem_target_ranks_addr = smem + 0; - int* smem_send_indices = reinterpret_cast(smem_raw + 88); - const int smem_send_indices_addr = smem + 88; - - // === Task calls (dependency order) === - int* workspace_i32 = reinterpret_cast(workspace); - unsigned int* workspace_u32 = reinterpret_cast(workspace); - unsigned long long local_workspace_base = (unsigned long long)ep_rank * workspace_stride_bytes; - int local_token_idx = bid; - if (enable_pdl) { - asm volatile("griddepcontrol.wait;" ::: "memory"); - } - if (local_num_tokens > 0) { - if (tid == 0) { - int ep_base = num_experts / ep_size; - int ep_remainder = num_experts - ep_base * ep_size; - int split = ep_remainder * (ep_base + 1); - unsigned long long seen_ranks = 0; - int route_base = local_token_idx * top_k; - unsigned long long topk_workspace_base = - (local_workspace_base + topk_target_ranks_offset) / 4 + - (unsigned long long)local_token_idx * (unsigned long long)top_k; - unsigned long long send_index_workspace_base = - (local_workspace_base + topk_send_indices_offset) / 4 + - (unsigned long long)local_token_idx * (unsigned long long)top_k; -#pragma unroll 1 - for (int k = 0; k < top_k; k++) { - int expert_id = token_selected_experts[route_base + k]; - int target_rank = 0; - if (ep_remainder == 0) { - target_rank = expert_id / ep_base; - } else if (expert_id < split) { - target_rank = expert_id / (ep_base + 1); - } else { - target_rank = ep_remainder + (expert_id - split) / ep_base; - } - unsigned long long target_bit = 1ULL << (unsigned long long)target_rank; - int first_for_rank = (((seen_ranks & target_bit) == 0) ? 1 : 0); - seen_ranks = seen_ranks | target_bit; - int rank_is_active = 1; - if (enable_rank_mask) { - rank_is_active = (((active_rank_mask & target_bit) != 0) ? 1 : 0); - } - int stored_rank = -1; - int send_index = -1; - if (first_for_rank != 0 && rank_is_active != 0) { - unsigned long long send_counter_index = - (local_workspace_base + send_counters_offset) / 4 + (unsigned long long)target_rank; - int _atomic_old_0 = atomicAdd(&workspace_i32[send_counter_index], 1); - send_index = _atomic_old_0; - stored_rank = target_rank; - } - smem_target_ranks[k] = stored_rank; - smem_send_indices[k] = send_index; - workspace_i32[topk_workspace_base + (unsigned long long)k] = stored_rank; - workspace_i32[send_index_workspace_base + (unsigned long long)k] = send_index; - } - } - __syncthreads(); - if (num_payloads > 0) { - unsigned long long source_base_0 = - (unsigned long long)local_token_idx * (unsigned long long)payload_0_bytes; -#pragma unroll 1 - for (int byte_0 = tid; byte_0 < payload_0_bytes; byte_0 += 256) { - uint8_t byte_value_0 = payload_0[source_base_0 + (unsigned long long)byte_0]; -#pragma unroll 1 - for (int k_0 = 0; k_0 < top_k; k_0++) { - int target_rank_0 = smem_target_ranks[k_0]; - int send_index_0 = smem_send_indices[k_0]; - if (send_index_0 >= 0) { - unsigned long long destination_0 = - (unsigned long long)target_rank_0 * workspace_stride_bytes + payload_0_offset + - ((unsigned long long)ep_rank * (unsigned long long)max_tokens_per_rank + - (unsigned long long)send_index_0) * - (unsigned long long)payload_0_bytes + - (unsigned long long)byte_0; - workspace[destination_0] = byte_value_0; - } - } - } - } - if (num_payloads > 1) { - unsigned long long source_base_1 = - (unsigned long long)local_token_idx * (unsigned long long)payload_1_bytes; -#pragma unroll 1 - for (int byte_1 = tid; byte_1 < payload_1_bytes; byte_1 += 256) { - uint8_t byte_value_1 = payload_1[source_base_1 + (unsigned long long)byte_1]; -#pragma unroll 1 - for (int k_1 = 0; k_1 < top_k; k_1++) { - int target_rank_1 = smem_target_ranks[k_1]; - int send_index_1 = smem_send_indices[k_1]; - if (send_index_1 >= 0) { - unsigned long long destination_1 = - (unsigned long long)target_rank_1 * workspace_stride_bytes + payload_1_offset + - ((unsigned long long)ep_rank * (unsigned long long)max_tokens_per_rank + - (unsigned long long)send_index_1) * - (unsigned long long)payload_1_bytes + - (unsigned long long)byte_1; - workspace[destination_1] = byte_value_1; - } - } - } - } - if (num_payloads > 2) { - unsigned long long source_base_2 = - (unsigned long long)local_token_idx * (unsigned long long)payload_2_bytes; -#pragma unroll 1 - for (int byte_2 = tid; byte_2 < payload_2_bytes; byte_2 += 256) { - uint8_t byte_value_2 = payload_2[source_base_2 + (unsigned long long)byte_2]; -#pragma unroll 1 - for (int k_2 = 0; k_2 < top_k; k_2++) { - int target_rank_2 = smem_target_ranks[k_2]; - int send_index_2 = smem_send_indices[k_2]; - if (send_index_2 >= 0) { - unsigned long long destination_2 = - (unsigned long long)target_rank_2 * workspace_stride_bytes + payload_2_offset + - ((unsigned long long)ep_rank * (unsigned long long)max_tokens_per_rank + - (unsigned long long)send_index_2) * - (unsigned long long)payload_2_bytes + - (unsigned long long)byte_2; - workspace[destination_2] = byte_value_2; - } - } - } - } - if (num_payloads > 3) { - unsigned long long source_base_3 = - (unsigned long long)local_token_idx * (unsigned long long)payload_3_bytes; -#pragma unroll 1 - for (int byte_3 = tid; byte_3 < payload_3_bytes; byte_3 += 256) { - uint8_t byte_value_3 = payload_3[source_base_3 + (unsigned long long)byte_3]; -#pragma unroll 1 - for (int k_3 = 0; k_3 < top_k; k_3++) { - int target_rank_3 = smem_target_ranks[k_3]; - int send_index_3 = smem_send_indices[k_3]; - if (send_index_3 >= 0) { - unsigned long long destination_3 = - (unsigned long long)target_rank_3 * workspace_stride_bytes + payload_3_offset + - ((unsigned long long)ep_rank * (unsigned long long)max_tokens_per_rank + - (unsigned long long)send_index_3) * - (unsigned long long)payload_3_bytes + - (unsigned long long)byte_3; - workspace[destination_3] = byte_value_3; - } - } - } - } - if (num_payloads > 4) { - unsigned long long source_base_4 = - (unsigned long long)local_token_idx * (unsigned long long)payload_4_bytes; -#pragma unroll 1 - for (int byte_4 = tid; byte_4 < payload_4_bytes; byte_4 += 256) { - uint8_t byte_value_4 = payload_4[source_base_4 + (unsigned long long)byte_4]; -#pragma unroll 1 - for (int k_4 = 0; k_4 < top_k; k_4++) { - int target_rank_4 = smem_target_ranks[k_4]; - int send_index_4 = smem_send_indices[k_4]; - if (send_index_4 >= 0) { - unsigned long long destination_4 = - (unsigned long long)target_rank_4 * workspace_stride_bytes + payload_4_offset + - ((unsigned long long)ep_rank * (unsigned long long)max_tokens_per_rank + - (unsigned long long)send_index_4) * - (unsigned long long)payload_4_bytes + - (unsigned long long)byte_4; - workspace[destination_4] = byte_value_4; - } - } - } - } - if (num_payloads > 5) { - unsigned long long source_base_5 = - (unsigned long long)local_token_idx * (unsigned long long)payload_5_bytes; -#pragma unroll 1 - for (int byte_5 = tid; byte_5 < payload_5_bytes; byte_5 += 256) { - uint8_t byte_value_5 = payload_5[source_base_5 + (unsigned long long)byte_5]; -#pragma unroll 1 - for (int k_5 = 0; k_5 < top_k; k_5++) { - int target_rank_5 = smem_target_ranks[k_5]; - int send_index_5 = smem_send_indices[k_5]; - if (send_index_5 >= 0) { - unsigned long long destination_5 = - (unsigned long long)target_rank_5 * workspace_stride_bytes + payload_5_offset + - ((unsigned long long)ep_rank * (unsigned long long)max_tokens_per_rank + - (unsigned long long)send_index_5) * - (unsigned long long)payload_5_bytes + - (unsigned long long)byte_5; - workspace[destination_5] = byte_value_5; - } - } - } - } - __syncthreads(); - } - asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); - if (warp == 0) { - int is_last_token = 0; - if (lane == 0) { - if (local_num_tokens == 0) { - is_last_token = 1; - } else { - unsigned long long local_counter_index = - (local_workspace_base + local_token_counter_offset) / 4; - int _atomic_old_1 = atomicAdd(&workspace_i32[local_counter_index], 1); - int completed_before = _atomic_old_1; - if (completed_before + 1 == local_num_tokens) { - is_last_token = 1; - } - } - } - int _shfl_0 = __shfl_sync(0xFFFFFFFF, is_last_token, 0); - is_last_token = _shfl_0; - if (is_last_token != 0) { -#pragma unroll 1 - for (int target_rank_6 = lane; target_rank_6 < ep_size; target_rank_6 += 32) { - int rank_is_active_1 = 1; - if (enable_rank_mask) { - unsigned long long target_bit_1 = 1ULL << (unsigned long long)target_rank_6; - rank_is_active_1 = (((active_rank_mask & target_bit_1) != 0) ? 1 : 0); - } - if (rank_is_active_1 != 0) { - unsigned long long send_counter_index_1 = - (local_workspace_base + send_counters_offset) / 4 + (unsigned long long)target_rank_6; - unsigned long long recv_counter_index = - ((unsigned long long)target_rank_6 * workspace_stride_bytes + recv_counters_offset) / - 4 + - (unsigned long long)ep_rank; - workspace_i32[recv_counter_index] = workspace_i32[send_counter_index_1]; - } - } - if (enable_eplb) { -#pragma unroll 1 - for (int target_rank_eplb = 0; target_rank_eplb < ep_size; target_rank_eplb++) { - int rank_is_active_eplb = 1; - if (enable_rank_mask) { - unsigned long long target_bit_eplb = 1ULL << (unsigned long long)target_rank_eplb; - rank_is_active_eplb = (((active_rank_mask & target_bit_eplb) != 0) ? 1 : 0); - } - if (rank_is_active_eplb != 0) { - unsigned long long gathered_base = - ((unsigned long long)target_rank_eplb * workspace_stride_bytes + - eplb_gathered_stats_offset) / - 4 + - (unsigned long long)ep_rank * (unsigned long long)eplb_stats_num_experts; -#pragma unroll 1 - for (int expert_stat = lane; expert_stat < eplb_stats_num_experts; expert_stat += 32) { - workspace_i32[gathered_base + (unsigned long long)expert_stat] = - eplb_local_stats[expert_stat]; - } - } - } - } - unsigned long long expected_value_index = (local_workspace_base + flag_val_offset) / 4; - unsigned int expected_value = workspace_u32[expected_value_index]; - asm volatile("fence.release.sys;" ::: "memory"); -#pragma unroll 1 - for (int target_rank_flag = lane; target_rank_flag < ep_size; target_rank_flag += 32) { - int rank_is_active_flag = 1; - if (enable_rank_mask) { - unsigned long long target_bit_flag = 1ULL << (unsigned long long)target_rank_flag; - rank_is_active_flag = (((active_rank_mask & target_bit_flag) != 0) ? 1 : 0); - } - if (rank_is_active_flag != 0) { - unsigned long long remote_flag_index = - ((unsigned long long)target_rank_flag * workspace_stride_bytes + - completion_flags_offset) / - 4 + - (unsigned long long)ep_rank; - asm volatile("st.relaxed.sys.u32 [%0], %1;" ::"l"( - (reinterpret_cast(workspace_u32) + (remote_flag_index))), - "r"(static_cast(expected_value)) - : "memory"); - } - } -#pragma unroll 1 - for (int peer_rank = lane; peer_rank < ep_size; peer_rank += 32) { - int peer_is_active = 1; - if (enable_rank_mask) { - unsigned long long peer_bit = 1ULL << (unsigned long long)peer_rank; - peer_is_active = (((active_rank_mask & peer_bit) != 0) ? 1 : 0); - } - if (peer_is_active != 0) { - unsigned long long local_flag_index = - (local_workspace_base + completion_flags_offset) / 4 + (unsigned long long)peer_rank; - { - unsigned int* _sre_ptr_0 = - (reinterpret_cast(workspace_u32) + (local_flag_index)); - const unsigned int _sre_expected_0 = static_cast(expected_value); - const unsigned long long _sre_start_0 = clock64(); - bool _sre_matched_0 = false; - do { - unsigned int _sre_value_0; - asm volatile("ld.relaxed.sys.u32 %0, [%1];" - : "=r"(_sre_value_0) - : "l"(_sre_ptr_0) - : "memory"); - _sre_matched_0 = (_sre_value_0 == _sre_expected_0); - } while (!_sre_matched_0 && - ((clock64() - _sre_start_0) <= static_cast(600000000000))); - if (__builtin_expect(!_sre_matched_0, 0)) { - asm volatile("trap;"); - return; - } - } - } - } - } - } -} - -} // extern "C" diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu deleted file mode 100644 index 2fcf9def6cc..00000000000 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ /dev/null @@ -1,1690 +0,0 @@ -/* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. - * - * 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 - -#include -#include - -#include "flashinfer/exception.h" -#include "flashinfer/utils.cuh" -#include "flashinfer/vec_dtypes.cuh" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h" -#include "tensorrt_llm/kernels/quantization_utils.cuh" - -namespace tensorrt_llm::kernels::moe_alltoall { - -using tensorrt_llm::common::launchWithPdlWhenEnabled; - -#define ENABLE_DEBUG_PRINT 0 -#define DISABLE_SYNC_FOR_PROFILING 0 - -constexpr int kEp4Size = 4; -// Two 32-thread warps can cover at most 64 16-byte activation vectors. -constexpr int kCompactDispatchMaxPayloadBytes = 1024; -constexpr int kCompactDispatchBlockSize = 128; -constexpr int kDispatchWarpSize = 32; -constexpr int kDispatchActivationWarps = 2; -constexpr int kDispatchWideVectorBytes = 16; -constexpr int kDispatchMetadataVectorBytes = 8; -constexpr int kDispatchHalfWarpSize = 16; -constexpr int kNvfp4ActivationToScaleRatio = 8; -constexpr int kNvfp4H2048ActivationBytes = 2048 / 2; -constexpr int kNvfp4H2048ScaleBytes = 2048 / 16; -constexpr int kTopK22MetadataBytes = 22 * sizeof(int32_t); -constexpr int kNvfp4H2048ActivationWorkers = kNvfp4H2048ActivationBytes / kDispatchWideVectorBytes; -constexpr int kNvfp4H2048ScaleWorkers = kNvfp4H2048ScaleBytes / kDispatchWideVectorBytes; -constexpr int kTopK22MetadataWorkers = kTopK22MetadataBytes / kDispatchMetadataVectorBytes; -static_assert(kCompactDispatchMaxPayloadBytes % kDispatchWideVectorBytes == 0); -static_assert(kCompactDispatchMaxPayloadBytes / kDispatchWideVectorBytes <= - kDispatchActivationWarps * kDispatchWarpSize); -static_assert(kCompactDispatchBlockSize == (kDispatchActivationWarps + 2) * kDispatchWarpSize); - -#ifndef DISABLE_TIMEOUT -#define DISABLE_TIMEOUT 0 -#endif - -// Helper function for ceiling division -template -__host__ __device__ inline T ceilDiv(T m, T n) { - return (m + n - 1) / n; -} - -// The phased schedule reserves two warps for packed activations, one for scales, -// and splits the last warp between the two routing-metadata payloads. The -// activation cap intentionally limits this path to hidden sizes up to 2048. -static bool isPhasedNvfp4PayloadLayout(DispatchKernelPointers const& ptrs, int num_payloads, - int top_k) { - if (num_payloads != 4 || top_k <= kEp4Size) { - return false; - } - - int const activation_bytes = ptrs.payload_bytes_per_token[0]; - int const scale_bytes = ptrs.payload_bytes_per_token[1]; - int const metadata_bytes = top_k * static_cast(sizeof(int32_t)); - // Smaller top-k layouts support runtime activation widths (including H1024). - // Top-k 22 is restricted to the measured H2048 configuration. - bool const validated_top_k_layout = top_k < 22 || activation_bytes == kNvfp4H2048ActivationBytes; - return validated_top_k_layout && activation_bytes > 0 && - activation_bytes <= kCompactDispatchMaxPayloadBytes && - activation_bytes % kDispatchWideVectorBytes == 0 && scale_bytes > 0 && - scale_bytes % kDispatchWideVectorBytes == 0 && - activation_bytes == kNvfp4ActivationToScaleRatio * scale_bytes && - metadata_bytes % kDispatchMetadataVectorBytes == 0 && - metadata_bytes / kDispatchMetadataVectorBytes <= kDispatchHalfWarpSize && - ptrs.payload_bytes_per_token[2] == metadata_bytes && - ptrs.payload_bytes_per_token[3] == metadata_bytes; -} - -// Macros for concise launch-time specialization -#define SWITCH_BOOL(flag, NAME, ...) \ - if (flag) { \ - constexpr bool NAME = true; \ - __VA_ARGS__ \ - } else { \ - constexpr bool NAME = false; \ - __VA_ARGS__ \ - } - -#define SWITCH_TOP_K(top_k, TOP_K, ...) \ - switch (top_k) { \ - case 22: { \ - constexpr int TOP_K = 22; \ - __VA_ARGS__; \ - break; \ - } \ - case 18: { \ - constexpr int TOP_K = 18; \ - __VA_ARGS__; \ - break; \ - } \ - case 16: { \ - constexpr int TOP_K = 16; \ - __VA_ARGS__; \ - break; \ - } \ - case 14: { \ - constexpr int TOP_K = 14; \ - __VA_ARGS__; \ - break; \ - } \ - case 12: { \ - constexpr int TOP_K = 12; \ - __VA_ARGS__; \ - break; \ - } \ - case 10: { \ - constexpr int TOP_K = 10; \ - __VA_ARGS__; \ - break; \ - } \ - case 8: { \ - constexpr int TOP_K = 8; \ - __VA_ARGS__; \ - break; \ - } \ - case 6: { \ - constexpr int TOP_K = 6; \ - __VA_ARGS__; \ - break; \ - } \ - case 4: { \ - constexpr int TOP_K = 4; \ - __VA_ARGS__; \ - break; \ - } \ - case 2: { \ - constexpr int TOP_K = 2; \ - __VA_ARGS__; \ - break; \ - } \ - case 1: { \ - constexpr int TOP_K = 1; \ - __VA_ARGS__; \ - break; \ - } \ - default: { \ - FLASHINFER_CHECK(false, "Unsupported top_k"); \ - } \ - } - -#define SWITCH_DTYPE(dtype, TYPE, ...) \ - switch (dtype) { \ - case nvinfer1::DataType::kHALF: { \ - using TYPE = half; \ - __VA_ARGS__; \ - break; \ - } \ - case nvinfer1::DataType::kBF16: { \ - using TYPE = __nv_bfloat16; \ - __VA_ARGS__; \ - break; \ - } \ - case nvinfer1::DataType::kFLOAT: { \ - using TYPE = float; \ - __VA_ARGS__; \ - break; \ - } \ - case nvinfer1::DataType::kFP8: { \ - using TYPE = __nv_fp8_e4m3; \ - __VA_ARGS__; \ - break; \ - } \ - default: { \ - FLASHINFER_CHECK(false, "Unsupported dtype for moe_a2a_combine"); \ - } \ - } - -// Quantized combine currently only supports FP16/BF16 inputs. -#define SWITCH_QUANT_MODE(dtype, quant_mode, QUANT_MODE, ...) \ - if constexpr (std::is_same_v || std::is_same_v) { \ - switch (quant_mode) { \ - case MoeA2ACombineQuantMode::NONE: { \ - constexpr auto QUANT_MODE = MoeA2ACombineQuantMode::NONE; \ - __VA_ARGS__; \ - break; \ - } \ - case MoeA2ACombineQuantMode::MXFP8: { \ - constexpr auto QUANT_MODE = MoeA2ACombineQuantMode::MXFP8; \ - __VA_ARGS__; \ - break; \ - } \ - case MoeA2ACombineQuantMode::NVFP4: { \ - constexpr auto QUANT_MODE = MoeA2ACombineQuantMode::NVFP4; \ - __VA_ARGS__; \ - break; \ - } \ - case MoeA2ACombineQuantMode::MXFP4: { \ - constexpr auto QUANT_MODE = MoeA2ACombineQuantMode::MXFP4; \ - __VA_ARGS__; \ - break; \ - } \ - default: { \ - FLASHINFER_CHECK(false, "Unsupported quant_mode for moe_a2a_combine"); \ - } \ - } \ - } else { \ - FLASHINFER_CHECK((quant_mode) == MoeA2ACombineQuantMode::NONE, \ - "All to All Combine Quantization currently only supports " \ - "FP16/BF16"); \ - constexpr auto QUANT_MODE = MoeA2ACombineQuantMode::NONE; \ - __VA_ARGS__; \ - } - -#define SWITCH_SWIZZLE_MODE(swizzle_mode, SWIZZLE_MODE, ...) \ - switch (swizzle_mode) { \ - case MoeA2ACombineSwizzleSFMode::LINEAR: { \ - constexpr auto SWIZZLE_MODE = MoeA2ACombineSwizzleSFMode::LINEAR; \ - __VA_ARGS__; \ - break; \ - } \ - case MoeA2ACombineSwizzleSFMode::SWIZZLE_128x4: { \ - constexpr auto SWIZZLE_MODE = MoeA2ACombineSwizzleSFMode::SWIZZLE_128x4; \ - __VA_ARGS__; \ - break; \ - } \ - case MoeA2ACombineSwizzleSFMode::SWIZZLE_8x4: { \ - constexpr auto SWIZZLE_MODE = MoeA2ACombineSwizzleSFMode::SWIZZLE_8x4; \ - __VA_ARGS__; \ - break; \ - } \ - default: { \ - FLASHINFER_CHECK(false, "Unsupported swizzle_mode for moe_a2a_combine"); \ - } \ - } - -#if DISABLE_TIMEOUT -#define check_timeout(s) false -#else -// 300 * 2000 MHz - should be high enough on any GPU but will prevent a hang -#define check_timeout(s) ((clock64() - (s)) > (300ll * 2000ll * 1000ll * 1000ll)) -#endif - -// ============================================================================ -// Helper Functions for Expert-to-Rank Mapping -// ============================================================================ - -// Compute which rank owns a given expert using contiguous ceil/floor partitioning. -// Supports non-divisible distribution when num_experts % ep_size != 0: -// base = num_experts / ep_size -// remainder = num_experts % ep_size -// - Ranks [0, remainder) each own (base + 1) experts. -// - Ranks [remainder, ep_size) each own base experts. -// -// Example A (uniform): 32 experts, 4 ranks -> base=8, remainder=0 -// - Rank 0: experts 0-7, Rank 1: 8-15, Rank 2: 16-23, Rank 3: 24-31 -// Example B (non-divisible): 384 experts, 5 ranks -> base=76, remainder=4 -// - Ranks 0-3: 77 experts each, Rank 4: 76 experts -// -// base and remainder are precomputed by the caller once outside the per-token TOP_K loop -// so the hot path performs at most one integer divide. -__device__ __forceinline__ int compute_target_rank_id(int expert_id, int base, int remainder) { - // Fast path for the uniform (num_experts % ep_size == 0) case: identical to the - // pre-ceil/floor implementation, so existing divisible deployments incur no overhead. - if (remainder == 0) { - return expert_id / base; - } - int const split = remainder * (base + 1); // boundary expert id - if (expert_id < split) { - // Falls inside the (base + 1)-sized prefix block. - return expert_id / (base + 1); - } - // Falls inside the base-sized suffix block. - return remainder + (expert_id - split) / base; -} - -// Test bit `rank` in a kRankMaskWords-wide little-endian uint64 bitmask. -// Word 0 covers ranks 0..63, word 1 covers ranks 64..127, etc. -// `rank >> 6` and `rank & 63` divide / modulo by 64. -__device__ __forceinline__ bool is_rank_active(uint64_t const* mask, int rank) { - return (mask[rank >> 6] >> (rank & 63)) & 1ULL; -} - -// ============================================================================ -// Helper Functions for Vectorized Memory Operations -// ============================================================================ - -template -__device__ void vectorized_copy_impl(void* dst, void const* src, int size) { - using flashinfer::vec_t; - - uint8_t* dst_ptr = static_cast(dst); - uint8_t const* src_ptr = static_cast(src); - - int const stride = blockDim.x * VEC_SIZE; - - for (int offset = threadIdx.x * VEC_SIZE; offset < size; offset += stride) { - vec_t v; - v.load(src_ptr + offset); - v.store(dst_ptr + offset); - } -} - -__device__ void vectorized_copy(void* dst, void const* src, int size) { - if (size % 16 == 0) { - vectorized_copy_impl<16>(dst, src, size); - } else if (size % 8 == 0) { - vectorized_copy_impl<8>(dst, src, size); - } else if (size % 4 == 0) { - vectorized_copy_impl<4>(dst, src, size); - } else if (size % 2 == 0) { - vectorized_copy_impl<2>(dst, src, size); - } else { - vectorized_copy_impl<1>(dst, src, size); - } -} - -// Vectorized dispatch: load one vec from source and write to up to TOP_K destinations -template -__device__ void vectorized_dispatch_impl(uint8_t const* src_ptr, int bytes_per_token, int rank_id, - int max_tokens_per_rank, int payload_idx, - DispatchKernelPointers const& ptrs, - int const* topk_target_ranks, int const* topk_send_indices, - int worker_idx, int num_workers) { - using flashinfer::vec_t; - - // Precompute destination base pointers per k - uint8_t* dst_base_k[TOP_K]; -#pragma unroll - for (int k = 0; k < TOP_K; ++k) { - int dst_idx_k = topk_send_indices[k]; - int target_rank_k = topk_target_ranks[k]; - if (dst_idx_k < 0) { - dst_base_k[k] = nullptr; - continue; - } - uint8_t* dst_data = static_cast(ptrs.recv_buffers[target_rank_k][payload_idx]); - size_t base_source_rank = - static_cast(rank_id) * static_cast(max_tokens_per_rank) + - static_cast(dst_idx_k); - size_t base_token = base_source_rank * static_cast(bytes_per_token); - dst_base_k[k] = dst_data + base_token; - } - - // The caller assigns a worker subset to this payload. Only those workers pay - // the per-destination pointer setup and copy cost. - int const stride = num_workers * VEC_SIZE; - for (int offset = worker_idx * VEC_SIZE; offset < bytes_per_token; offset += stride) { - vec_t v; - v.load(src_ptr + offset); - -#pragma unroll - for (int k = 0; k < TOP_K; ++k) { - uint8_t* dst_base = dst_base_k[k]; - if (dst_base == nullptr) { - continue; - } - v.store(dst_base + offset); - } - } -} - -template -__device__ void vectorized_dispatch(uint8_t const* src_ptr, int bytes_per_token, int rank_id, - int max_tokens_per_rank, int payload_idx, - DispatchKernelPointers const& ptrs, - int const* topk_target_ranks, int const* topk_send_indices) { - // Receive planes are packed without vector-width padding. A small preceding - // payload can misalign the next plane even when its row size is divisible by - // 16. Token offsets are multiples of bytes_per_token, so checking every base - // together with the row size preserves wide copies exactly where safe. - uintptr_t copy_alignment = - reinterpret_cast(src_ptr) | static_cast(bytes_per_token); -#pragma unroll - for (int k = 0; k < TOP_K; ++k) { - if (topk_send_indices[k] >= 0) { - copy_alignment |= - reinterpret_cast(ptrs.recv_buffers[topk_target_ranks[k]][payload_idx]); - } - } - if (copy_alignment % 16 == 0) { - vectorized_dispatch_impl<16, TOP_K>(src_ptr, bytes_per_token, rank_id, max_tokens_per_rank, - payload_idx, ptrs, topk_target_ranks, topk_send_indices, - threadIdx.x, blockDim.x); - } else if (copy_alignment % 8 == 0) { - vectorized_dispatch_impl<8, TOP_K>(src_ptr, bytes_per_token, rank_id, max_tokens_per_rank, - payload_idx, ptrs, topk_target_ranks, topk_send_indices, - threadIdx.x, blockDim.x); - } else if (copy_alignment % 4 == 0) { - vectorized_dispatch_impl<4, TOP_K>(src_ptr, bytes_per_token, rank_id, max_tokens_per_rank, - payload_idx, ptrs, topk_target_ranks, topk_send_indices, - threadIdx.x, blockDim.x); - } else if (copy_alignment % 2 == 0) { - vectorized_dispatch_impl<2, TOP_K>(src_ptr, bytes_per_token, rank_id, max_tokens_per_rank, - payload_idx, ptrs, topk_target_ranks, topk_send_indices, - threadIdx.x, blockDim.x); - } else { - vectorized_dispatch_impl<1, TOP_K>(src_ptr, bytes_per_token, rank_id, max_tokens_per_rank, - payload_idx, ptrs, topk_target_ranks, topk_send_indices, - threadIdx.x, blockDim.x); - } -} - -__global__ void moeA2APrepareDispatchKernel(int* send_counters, int* local_token_counter, - int ep_size, uint32_t* flag_val_ptr, bool enable_pdl) { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (enable_pdl) cudaGridDependencySynchronize(); - cudaTriggerProgrammaticLaunchCompletion(); -#endif - int idx = blockIdx.x * blockDim.x + threadIdx.x; - // Zero send_counters - if (idx < ep_size) { - send_counters[idx] = 0; - } - // Zero local_token_counter and increment flag_val - if (idx == 0) { - *local_token_counter = 0; - // Increment flag_val for this dispatch round - *flag_val_ptr = *flag_val_ptr + 1; - } -} - -// ============================================================================ -// Generic Dispatch Kernel Implementation -// One CTA processes one token and all its payloads. -// ============================================================================ - -template -__global__ void moeA2ADispatchKernel( - int32_t const* token_selected_experts, // [local_num_tokens, TOP_K] - const DispatchKernelPointers ptrs, // Struct containing all kernel pointers - int num_payloads, // Number of payloads - int max_tokens_per_rank, // Maximum tokens per rank - int local_num_tokens, int rank_id, int ep_size, int num_experts, int eplb_stats_num_experts, - bool enable_pdl) { - static_assert(!COMPACT_EP4 || TOP_K > kEp4Size); - static_assert(!PHASE_H2048_NVFP4 || TOP_K == 22); - static_assert(!PHASE_NVFP4_PAYLOADS || - (TOP_K > kEp4Size && TOP_K % 2 == 0 && - TOP_K * static_cast(sizeof(int32_t)) / kDispatchMetadataVectorBytes <= - kDispatchHalfWarpSize)); - static_assert(!PHASE_H2048_NVFP4 || !PHASE_NVFP4_PAYLOADS); - int thread_idx = threadIdx.x; - int local_token_idx = blockIdx.x; - - if (local_num_tokens == 0) { - // Special case: If local_num_tokens == 0, - // we need to keep the threads where local_token_idx == 0 alive to participate in the - // synchronization. Other threads should return. - if (local_token_idx > 0) return; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (enable_pdl) cudaGridDependencySynchronize(); -#endif - } else { - // Threads that do not have a token to process should return. - if (local_token_idx >= local_num_tokens) return; - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (enable_pdl) cudaGridDependencySynchronize(); -#endif - - // Prepare per-policy shared-memory tiles for this token - extern __shared__ int smem[]; - int* smem_topk_target_ranks = smem; - int* smem_topk_send_indices = nullptr; - int* smem_compact_send_indices = smem; - if constexpr (!COMPACT_EP4) { - smem_topk_send_indices = smem + TOP_K; - } - - // Precompute the ceil/floor partition parameters once per thread, outside the - // per-token TOP_K loop. The fast path (remainder == 0) then collapses to a single - // integer divide per call, matching the pre-PR uniform-partition cost exactly. - int const ep_base = num_experts / ep_size; - int const ep_remainder = num_experts - ep_base * ep_size; // == num_experts % ep_size - - if (thread_idx < warpSize) { - int lane_id = thread_idx; - if constexpr (COMPACT_EP4) { - if (lane_id < kEp4Size) { - smem_compact_send_indices[lane_id] = -1; - } - __syncwarp(); - } - - unsigned topk_mask = __ballot_sync(0xffffffff, lane_id < TOP_K); - if (lane_id < TOP_K) { - int k = lane_id; - int expert_id = token_selected_experts[local_token_idx * TOP_K + k]; - // Use contiguous ceil/floor partitioning (supports non-divisible num_experts % ep_size). - int target_rank = compute_target_rank_id(expert_id, ep_base, ep_remainder); - - // Elect the first top-k lane for each destination rank; duplicate targets within a - // token collapse to a single send (replaces the old serial already_copied bitmask). - unsigned matching_lanes = __match_any_sync(topk_mask, target_rank); - bool is_first = lane_id == __ffs(matching_lanes) - 1; - - // Dead ranks (ENABLE_RANK_MASK) are dropped the same way as duplicates: no send is - // issued and topk_send_indices[k] stays < 0, so combine's existing skip check on that - // sentinel handles both cases uniformly. When ENABLE_RANK_MASK is false, the dead-rank - // check is compiled out entirely (no active_rank_mask load/branch on this hot path). - bool is_valid = is_first; - if constexpr (ENABLE_RANK_MASK) { - is_valid = is_valid && is_rank_active(ptrs.active_rank_mask, target_rank); - } - - int dst_token_idx = -1; - if (is_valid) { - dst_token_idx = atomicAdd(&ptrs.send_counters[target_rank], 1); - if constexpr (COMPACT_EP4) { - smem_compact_send_indices[target_rank] = dst_token_idx; - } - } else { - target_rank = -1; - } - - ptrs.topk_target_ranks[local_token_idx * TOP_K + k] = target_rank; - ptrs.topk_send_indices[local_token_idx * TOP_K + k] = dst_token_idx; - if constexpr (!COMPACT_EP4) { - smem_topk_target_ranks[k] = target_rank; - smem_topk_send_indices[k] = dst_token_idx; - } - } - } - // Sync before dispatching data - __syncthreads(); - - // EP4 has at most four unique destinations, regardless of TOP_K. The compact - // specialization avoids expanding those destinations back to TOP_K entries. - constexpr int NUM_DESTINATIONS = COMPACT_EP4 ? kEp4Size : TOP_K; - int target_ranks[NUM_DESTINATIONS]; - int send_indices[NUM_DESTINATIONS]; - if constexpr (COMPACT_EP4) { -#pragma unroll - for (int target_rank = 0; target_rank < NUM_DESTINATIONS; ++target_rank) { - target_ranks[target_rank] = target_rank; - send_indices[target_rank] = smem_compact_send_indices[target_rank]; - } - } else { -#pragma unroll - for (int k = 0; k < TOP_K; ++k) { - target_ranks[k] = smem_topk_target_ranks[k]; - send_indices[k] = smem_topk_send_indices[k]; - } - } - - if constexpr (PHASE_H2048_NVFP4) { - // Keep compile-time bounds for the measured Nemotron H2048 target. - int warp_id = thread_idx / warpSize; - int lane_id = thread_idx % warpSize; - if (warp_id < kDispatchActivationWarps) { - auto const* src = static_cast(ptrs.src_data_ptrs[0]) + - static_cast(local_token_idx) * kNvfp4H2048ActivationBytes; - vectorized_dispatch_impl<16, NUM_DESTINATIONS>( - src, kNvfp4H2048ActivationBytes, rank_id, max_tokens_per_rank, 0, ptrs, target_ranks, - send_indices, thread_idx, kNvfp4H2048ActivationWorkers); - } else if (warp_id == kDispatchActivationWarps && lane_id < kNvfp4H2048ScaleWorkers) { - auto const* src = static_cast(ptrs.src_data_ptrs[1]) + - static_cast(local_token_idx) * kNvfp4H2048ScaleBytes; - vectorized_dispatch_impl<16, NUM_DESTINATIONS>( - src, kNvfp4H2048ScaleBytes, rank_id, max_tokens_per_rank, 1, ptrs, target_ranks, - send_indices, lane_id, kNvfp4H2048ScaleWorkers); - } else if (warp_id == kDispatchActivationWarps + 1 && lane_id < 2 * kTopK22MetadataWorkers) { - bool second_tail = lane_id >= kTopK22MetadataWorkers; - int payload_idx = second_tail ? 3 : 2; - int worker_idx = second_tail ? lane_id - kTopK22MetadataWorkers : lane_id; - auto const* src = static_cast(ptrs.src_data_ptrs[payload_idx]) + - static_cast(local_token_idx) * kTopK22MetadataBytes; - vectorized_dispatch_impl<8, NUM_DESTINATIONS>( - src, kTopK22MetadataBytes, rank_id, max_tokens_per_rank, payload_idx, ptrs, - target_ranks, send_indices, worker_idx, kTopK22MetadataWorkers); - } - } else if constexpr (PHASE_NVFP4_PAYLOADS) { - // Match warps to the NVFP4 payload widths instead of making all 128 - // threads materialize destination pointers for every payload. - int const activation_bytes = ptrs.payload_bytes_per_token[0]; - int const scale_bytes = ptrs.payload_bytes_per_token[1]; - constexpr int METADATA_BYTES = TOP_K * sizeof(int32_t); - constexpr int METADATA_WORKERS = METADATA_BYTES / kDispatchMetadataVectorBytes; - int const activation_workers = activation_bytes / kDispatchWideVectorBytes; - int const scale_workers = scale_bytes / kDispatchWideVectorBytes; - int warp_id = thread_idx / warpSize; - int lane_id = thread_idx % warpSize; - if (warp_id < kDispatchActivationWarps && thread_idx < activation_workers) { - auto const* src = static_cast(ptrs.src_data_ptrs[0]) + - static_cast(local_token_idx) * activation_bytes; - vectorized_dispatch_impl( - src, activation_bytes, rank_id, max_tokens_per_rank, 0, ptrs, target_ranks, - send_indices, thread_idx, activation_workers); - } else if (warp_id == kDispatchActivationWarps && lane_id < scale_workers) { - auto const* src = static_cast(ptrs.src_data_ptrs[1]) + - static_cast(local_token_idx) * scale_bytes; - vectorized_dispatch_impl( - src, scale_bytes, rank_id, max_tokens_per_rank, 1, ptrs, target_ranks, send_indices, - lane_id, scale_workers); - } else if (warp_id == kDispatchActivationWarps + 1 && lane_id < 2 * METADATA_WORKERS) { - bool second_tail = lane_id >= METADATA_WORKERS; - int payload_idx = second_tail ? 3 : 2; - int worker_idx = second_tail ? lane_id - METADATA_WORKERS : lane_id; - auto const* src = static_cast(ptrs.src_data_ptrs[payload_idx]) + - static_cast(local_token_idx) * METADATA_BYTES; - vectorized_dispatch_impl( - src, METADATA_BYTES, rank_id, max_tokens_per_rank, payload_idx, ptrs, target_ranks, - send_indices, worker_idx, METADATA_WORKERS); - } - } else { - // Perform a single source load and fan out to each unique destination. - for (int payload_idx = 0; payload_idx < num_payloads; payload_idx++) { - uint8_t const* src_data = static_cast(ptrs.src_data_ptrs[payload_idx]); - int bytes_per_token = ptrs.payload_bytes_per_token[payload_idx]; - uint8_t const* src_ptr = src_data + local_token_idx * bytes_per_token; - - vectorized_dispatch(src_ptr, bytes_per_token, rank_id, - max_tokens_per_rank, payload_idx, ptrs, target_ranks, - send_indices); - } - } - - __syncthreads(); - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif - - bool is_first_warp = threadIdx.x / warpSize == 0; - if (is_first_warp) { - int lane_id = threadIdx.x % warpSize; - - bool is_last_token = false; - if (lane_id == 0) { - if (local_num_tokens != 0) { - int cnt = atomicAdd(ptrs.local_token_counter, 1); - is_last_token = cnt + 1 == local_num_tokens; - } else { - is_last_token = true; - } - } - is_last_token = __shfl_sync(0xffffffff, is_last_token, 0); - - if (is_last_token) { -// Store send_counters to recv_counters. -// Skip masked target ranks: their symmetric memory may be inaccessible. -#pragma unroll 1 // No unroll as one iter is typically enough - for (int target_rank = lane_id; target_rank < ep_size; target_rank += warpSize) { - if constexpr (ENABLE_RANK_MASK) { - if (!is_rank_active(ptrs.active_rank_mask, target_rank)) continue; - } - int send_count = ptrs.send_counters[target_rank]; - ptrs.recv_counters[target_rank][rank_id] = send_count; - } - - if constexpr (ENABLE_EPLB) { - // Write local stats into peer buffers before the release fence below. - // Skip masked target ranks for the same reason as above. -#pragma unroll 1 - for (int target_rank = 0; target_rank < ep_size; ++target_rank) { - if constexpr (ENABLE_RANK_MASK) { - if (!is_rank_active(ptrs.active_rank_mask, target_rank)) continue; - } - int* target_stats = ptrs.eplb_gathered_stats[target_rank]; - for (int expert_id = lane_id; expert_id < eplb_stats_num_experts; expert_id += warpSize) { - int stat_val = ptrs.eplb_local_stats[expert_id]; - target_stats[rank_id * eplb_stats_num_experts + expert_id] = stat_val; - } - } - } - -#if !DISABLE_SYNC_FOR_PROFILING - uint32_t expected_value = *ptrs.flag_val; - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("fence.release.sys;"); -#else - __threadfence_system(); -#endif -// Signal completion to all active peers; skip dead ranks (their symmetric memory is -// unreachable). -#pragma unroll 1 // No unroll as one iter is typically enough - for (int target_rank = lane_id; target_rank < ep_size; target_rank += warpSize) { - if constexpr (ENABLE_RANK_MASK) { - if (!is_rank_active(ptrs.active_rank_mask, target_rank)) continue; - } - uint32_t* flag_addr = &ptrs.completion_flags[target_rank][rank_id]; - asm volatile("st.relaxed.sys.u32 [%0], %1;" ::"l"(flag_addr), "r"(expected_value)); - -#if ENABLE_DEBUG_PRINT - printf("dispatch: +++Rank %d setting completion flag to %d for rank %d\n", rank_id, - expected_value, target_rank); -#endif - } - -// Wait for all active peers to signal; skip dead ranks (otherwise we would spin forever — -// this is the bug the rank-mask is here to prevent). -#pragma unroll 1 // No unroll - for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize) { - if constexpr (ENABLE_RANK_MASK) { - if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) continue; - } - bool flag_set = false; - [[maybe_unused]] auto s = clock64(); - do { - uint32_t* flag_ptr = &ptrs.completion_flags[rank_id][peer_rank]; - uint32_t flag_value; - // Acquire load to ensure visibility of peer's release-store - asm volatile("ld.relaxed.sys.u32 %0, [%1];" : "=r"(flag_value) : "l"(flag_ptr)); -#if ENABLE_DEBUG_PRINT - printf( - "dispatch: ---Rank %d received completion flag from rank %d, flag_value: %d, " - "expected_value: " - "%d, address: %p\n", - rank_id, peer_rank, flag_value, expected_value, flag_ptr); -#endif - flag_set = flag_value == expected_value; - } while (!flag_set && !check_timeout(s)); - - if (__builtin_expect(!flag_set, 0)) { - printf("dispatch: ---Rank %d timed out waiting for completion flag from rank %d\n", - rank_id, peer_rank); - asm volatile("trap;"); - return; - } - } - // asm volatile("fence.acquire.sys;"); -#endif - } - } -} - -void moe_a2a_prepare_dispatch_launch(MoeA2ADispatchParams const& params) { - launchWithPdlWhenEnabled("moeA2APrepareDispatchKernel", params.enable_pdl, - moeA2APrepareDispatchKernel, 1, params.ep_size, 0, params.stream, - params.send_counters, params.local_token_counter, params.ep_size, - params.flag_val, params.enable_pdl); -} - -// ============================================================================ -// Launch Functions -// ============================================================================ - -void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) { - // Validate parameters - TLLM_CHECK(params.top_k > 0 && params.top_k <= kMaxTopK); - TLLM_CHECK(params.ep_size > 0 && params.ep_size <= kMaxRanks); - TLLM_CHECK(params.ep_rank >= 0 && params.ep_rank < params.ep_size); - TLLM_CHECK(params.local_num_tokens >= 0); - TLLM_CHECK(params.num_payloads > 0 && params.num_payloads <= kMaxPayloads); - if (params.enable_rank_mask) { - // The local rank must always be marked active in its own view of the mask; otherwise the - // kernel itself would be running on a "dead" rank. - TLLM_CHECK_WITH_INFO( - (params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL, - "active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank); - } - - // Prepare kernel pointers struct - DispatchKernelPointers kernel_ptrs = {}; - - // Fill source data pointers and payload sizes - for (int i = 0; i < params.num_payloads; i++) { - kernel_ptrs.src_data_ptrs[i] = params.payloads[i].src_data; - kernel_ptrs.payload_bytes_per_token[i] = - params.payloads[i].element_size * params.payloads[i].elements_per_token; - } - - // Fill receive buffer pointers - for (int target_rank = 0; target_rank < params.ep_size; target_rank++) { - kernel_ptrs.recv_counters[target_rank] = params.recv_counters[target_rank]; - kernel_ptrs.eplb_gathered_stats[target_rank] = params.eplb_gathered_stats[target_rank]; - for (int payload = 0; payload < params.num_payloads; payload++) { - kernel_ptrs.recv_buffers[target_rank][payload] = params.recv_buffers[target_rank][payload]; - } - } - - // Copy completion flag pointers - for (int i = 0; i < params.ep_size; i++) { - kernel_ptrs.completion_flags[i] = params.completion_flags[i]; - } - kernel_ptrs.flag_val = params.flag_val; - - // Copy communication tracking pointers - kernel_ptrs.send_counters = params.send_counters; - kernel_ptrs.local_token_counter = params.local_token_counter; - kernel_ptrs.topk_target_ranks = params.topk_target_ranks; - kernel_ptrs.topk_send_indices = params.topk_send_indices; - kernel_ptrs.eplb_local_stats = params.eplb_local_stats; - - // Copy active-rank bitmask into the kernel pointers struct - for (int w = 0; w < kRankMaskWords; ++w) { - kernel_ptrs.active_rank_mask[w] = params.active_rank_mask[w]; - } - - int block_size = tensorrt_llm::common::getEnvMoeA2ADispatchBlockSize(); - - bool const use_compact_ep4 = params.ep_size == kEp4Size && params.top_k > kEp4Size; - - int max_payload_bytes_per_token = 0; - for (int i = 0; i < params.num_payloads; ++i) { - if (kernel_ptrs.payload_bytes_per_token[i] > max_payload_bytes_per_token) { - max_payload_bytes_per_token = kernel_ptrs.payload_bytes_per_token[i]; - } - } - bool const supports_phased_nvfp4 = - params.ep_size == kEp4Size || params.ep_size == 8 || params.ep_size == 16; - bool const phased_nvfp4_layout = - supports_phased_nvfp4 && - isPhasedNvfp4PayloadLayout(kernel_ptrs, params.num_payloads, params.top_k); - - // Small payloads do not have enough 16-byte vectors to use larger CTAs. A GB200 - // sweep found 128 threads faster than both 64 and 256 at prefill token counts. - if ((use_compact_ep4 || phased_nvfp4_layout) && - max_payload_bytes_per_token <= kCompactDispatchMaxPayloadBytes && - block_size > kCompactDispatchBlockSize) { - block_size = kCompactDispatchBlockSize; - } - bool const phase_nvfp4_payloads = phased_nvfp4_layout && block_size == kCompactDispatchBlockSize; - - // Configure kernel launch: one block per token - int grid_size = params.local_num_tokens; - // If local_num_tokens is 0, we still need to launch a minimal kernel to participate in the - // synchronization. - if (grid_size == 0) { - grid_size = 1; - } - SWITCH_BOOL(params.enable_rank_mask, ENABLE_RANK_MASK, { - SWITCH_BOOL(params.enable_eplb, EPLB_STATS, { - if (phase_nvfp4_payloads && params.top_k == 22 && - kernel_ptrs.payload_bytes_per_token[0] == kNvfp4H2048ActivationBytes) { - FLASHINFER_CHECK(kernel_ptrs.payload_bytes_per_token[1] == kNvfp4H2048ScaleBytes && - kernel_ptrs.payload_bytes_per_token[2] == kTopK22MetadataBytes && - kernel_ptrs.payload_bytes_per_token[3] == kTopK22MetadataBytes, - "Invalid H2048/top-k-22 NVFP4 dispatch payload layout"); - if (use_compact_ep4) { - int shared_bytes = kEp4Size * (int)sizeof(int); - auto kernel_fn = moeA2ADispatchKernel<22, EPLB_STATS, ENABLE_RANK_MASK, true, true>; - launchWithPdlWhenEnabled("moeA2ADispatchKernel", params.enable_pdl, kernel_fn, grid_size, - block_size, shared_bytes, params.stream, - params.token_selected_experts, kernel_ptrs, params.num_payloads, - params.max_tokens_per_rank, params.local_num_tokens, - params.ep_rank, params.ep_size, params.num_experts, - params.eplb_stats_num_experts, params.enable_pdl); - } else { - int shared_bytes = 2 * params.top_k * (int)sizeof(int); - auto kernel_fn = moeA2ADispatchKernel<22, EPLB_STATS, ENABLE_RANK_MASK, false, true>; - launchWithPdlWhenEnabled("moeA2ADispatchKernel", params.enable_pdl, kernel_fn, grid_size, - block_size, shared_bytes, params.stream, - params.token_selected_experts, kernel_ptrs, params.num_payloads, - params.max_tokens_per_rank, params.local_num_tokens, - params.ep_rank, params.ep_size, params.num_experts, - params.eplb_stats_num_experts, params.enable_pdl); - } - } else if (phase_nvfp4_payloads) { - if (use_compact_ep4) { - int shared_bytes = kEp4Size * (int)sizeof(int); - SWITCH_TOP_K( - params.top_k, TOP_K, if constexpr (TOP_K > kEp4Size && TOP_K % 2 == 0) { - auto kernel_fn = - moeA2ADispatchKernel; - launchWithPdlWhenEnabled( - "moeA2ADispatchKernel", params.enable_pdl, kernel_fn, grid_size, block_size, - shared_bytes, params.stream, params.token_selected_experts, kernel_ptrs, - params.num_payloads, params.max_tokens_per_rank, params.local_num_tokens, - params.ep_rank, params.ep_size, params.num_experts, - params.eplb_stats_num_experts, params.enable_pdl); - }) - } else { - int shared_bytes = 2 * params.top_k * (int)sizeof(int); - SWITCH_TOP_K( - params.top_k, TOP_K, if constexpr (TOP_K > kEp4Size && TOP_K % 2 == 0) { - auto kernel_fn = - moeA2ADispatchKernel; - launchWithPdlWhenEnabled( - "moeA2ADispatchKernel", params.enable_pdl, kernel_fn, grid_size, block_size, - shared_bytes, params.stream, params.token_selected_experts, kernel_ptrs, - params.num_payloads, params.max_tokens_per_rank, params.local_num_tokens, - params.ep_rank, params.ep_size, params.num_experts, - params.eplb_stats_num_experts, params.enable_pdl); - }) - } - } else if (use_compact_ep4) { - int shared_bytes = kEp4Size * (int)sizeof(int); - SWITCH_TOP_K( - params.top_k, TOP_K, if constexpr (TOP_K > kEp4Size) { - auto kernel_fn = moeA2ADispatchKernel; - launchWithPdlWhenEnabled( - "moeA2ADispatchKernel", params.enable_pdl, kernel_fn, grid_size, block_size, - shared_bytes, params.stream, params.token_selected_experts, kernel_ptrs, - params.num_payloads, params.max_tokens_per_rank, params.local_num_tokens, - params.ep_rank, params.ep_size, params.num_experts, params.eplb_stats_num_experts, - params.enable_pdl); - }) - } else { - int shared_bytes = 2 * params.top_k * (int)sizeof(int); - SWITCH_TOP_K(params.top_k, TOP_K, { - auto kernel_fn = moeA2ADispatchKernel; - launchWithPdlWhenEnabled("moeA2ADispatchKernel", params.enable_pdl, kernel_fn, grid_size, - block_size, shared_bytes, params.stream, - params.token_selected_experts, kernel_ptrs, params.num_payloads, - params.max_tokens_per_rank, params.local_num_tokens, - params.ep_rank, params.ep_size, params.num_experts, - params.eplb_stats_num_experts, params.enable_pdl); - }); - } - }); - }) -} - -// ============================================================================ -// Combine kernels -// ============================================================================ - -// Accumulate across all valid ranks into float32 registers, then store as T (QuantMode==NONE) -// or quantize the result to MXFP8 with a per-block scale factor (QuantMode==MXFP8). -// -// InT: input element type in the recv buffer (defaults to T for same-type accumulation; -// __nv_fp8_e4m3 for the low-precision combine path). sizeof(InT) must divide VEC_SIZE_BYTES. -// T: output / accumulation-cast element type (and the input type when InT == T). -// -// Unified path: load VEC_SIZE_BYTES bytes, reinterpret as InT[elems_per_vec], accumulate as -// float32, then either cast_store as T (NONE) or convert to MXFP8 + emit the block scale factor -// (MXFP8, which always uses InT == T). Block-per-token indexing -// (threadIdx.x/blockDim.x/blockIdx.x). -template -__device__ void vectorized_combine_impl(void* output_buffer, void* sf_output, int row_idx, - int row_size, int stride_per_token, int rank_id, - int max_tokens_per_rank, CombineKernelPointers const& ptrs, - float OutputScalarScale = 1.0f) { - using flashinfer::vec_t; - - // elems_per_vec: number of InT elements per VEC_SIZE_BYTES-byte load (constexpr). - constexpr int elems_per_vec = VEC_SIZE_BYTES / static_cast(sizeof(InT)); - // size_per_token: byte span of one token in the recv buffer (InT-typed). - const int size_per_token = row_size * static_cast(sizeof(InT)); - - // Output base pointers. NONE writes T elements; MXFP8 writes packed fp8 bytes. - T* dst_typed_base = nullptr; - uint8_t* dst_bytes = nullptr; - if constexpr (QuantMode == MoeA2ACombineQuantMode::NONE) { - dst_typed_base = static_cast(output_buffer) + static_cast(row_idx) * row_size; - } else { - // MXFP8 stores one byte per logical element; FP4 packs two e2m1 values per byte, so its - // rows are half as wide. The accumulation still reads InT-sized inputs either way. - size_t const bytes_per_row = QuantMode == MoeA2ACombineQuantMode::MXFP8 - ? static_cast(row_size) - : static_cast(row_size) / 2; - dst_bytes = static_cast(output_buffer) + static_cast(row_idx) * bytes_per_row; - } - - int const stride = blockDim.x * VEC_SIZE_BYTES; - int const local_token_idx = blockIdx.x; - - // offset is a byte offset into the recv buffer, stepping by VEC_SIZE_BYTES bytes. - for (int offset = threadIdx.x * VEC_SIZE_BYTES; offset < size_per_token; offset += stride) { - int logical_offset = offset / static_cast(sizeof(InT)); - // Per-k vec_t accumulators, zero-initialised via fill(). - vec_t acc[TOP_K]; - - // Pass 1: issue all TOP_K loads back-to-back without any type conversion. - // Raw InT bytes are loaded directly into acc[k]'s register storage, reinterpreted as - // vec_t (VEC_SIZE_BYTES bytes, fitting in the low end of acc[k]'s - // sizeof(float)*elems_per_vec allocation). Separating load from cast lets the compiler - // schedule all VEC_SIZE_BYTES-byte global loads consecutively, hiding memory latency across k. -#pragma unroll - for (int k = 0; k < TOP_K; ++k) { - int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k]; - int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k]; - // dst_idx < 0: duplicate/dead-target sentinel already set by dispatch (which consults - // active_rank_mask once per token). Rechecking is_rank_active here is redundant and, - // being on the per-k vectorized hot path, measurably regresses combine throughput. - if (dst_idx < 0) { - acc[k].fill(0.0f); - continue; - } - - uint8_t const* recv_buffer = static_cast(ptrs.recv_buffers[target_rank][0]); - size_t base_source_rank = - static_cast(rank_id) * static_cast(max_tokens_per_rank) + - static_cast(dst_idx); - // stride_per_token: byte distance between tokens in the recv buffer. - // Equals size_per_token for normal cases; may differ for FP8 in-place - // (BF16-stride workspace but FP8-sized payload). - size_t base_token = base_source_rank * static_cast(stride_per_token); - - reinterpret_cast&>(acc[k]).load( - reinterpret_cast(recv_buffer + base_token + offset)); - } - - // Pass 2: in-place cast InT → float, iterating j in descending order. - // float[j] occupies bytes [j*4, j*4+3]; InT[j] occupies [j*sizeof(InT), ...). - // For sizeof(InT) < sizeof(float), high-j float writes land above all remaining - // InT bytes, so descending order is always write-after-read safe. -#pragma unroll - for (int k = 0; k < TOP_K; ++k) { - if (ptrs.topk_send_indices[local_token_idx * TOP_K + k] < 0) - continue; // acc[k] already holds 0.0f from fill() above -#pragma unroll - for (int j = elems_per_vec - 1; j >= 0; --j) - acc[k][j] = static_cast(reinterpret_cast(&acc[k])[j]); - } - // Reduce acc[TOP_K] into acc[0] via unrolled tree-reduction. - // acc[k][j] uses vec_t::operator[] which returns float& — no indirection overhead. - if constexpr (TOP_K == 22) { -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[1][j]; - acc[2][j] += acc[3][j]; - acc[4][j] += acc[5][j]; - acc[6][j] += acc[7][j]; - acc[8][j] += acc[9][j]; - acc[10][j] += acc[11][j]; - acc[12][j] += acc[13][j]; - acc[14][j] += acc[15][j]; - acc[16][j] += acc[17][j]; - acc[18][j] += acc[19][j]; - acc[20][j] += acc[21][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[2][j]; - acc[4][j] += acc[6][j]; - acc[8][j] += acc[10][j]; - acc[12][j] += acc[14][j]; - acc[16][j] += acc[18][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[4][j]; - acc[8][j] += acc[12][j]; - acc[16][j] += acc[20][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[8][j]; - acc[0][j] += acc[16][j]; - } - } else if constexpr (TOP_K == 16) { -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[1][j]; - acc[2][j] += acc[3][j]; - acc[4][j] += acc[5][j]; - acc[6][j] += acc[7][j]; - acc[8][j] += acc[9][j]; - acc[10][j] += acc[11][j]; - acc[12][j] += acc[13][j]; - acc[14][j] += acc[15][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[2][j]; - acc[4][j] += acc[6][j]; - acc[8][j] += acc[10][j]; - acc[12][j] += acc[14][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[4][j]; - acc[8][j] += acc[12][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[8][j]; - } - } else if constexpr (TOP_K == 10) { -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[1][j]; - acc[2][j] += acc[3][j]; - acc[4][j] += acc[5][j]; - acc[6][j] += acc[7][j]; - acc[8][j] += acc[9][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[2][j]; - acc[4][j] += acc[6][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[4][j]; - acc[0][j] += acc[8][j]; - } - } else if constexpr (TOP_K == 8) { -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[1][j]; - acc[2][j] += acc[3][j]; - acc[4][j] += acc[5][j]; - acc[6][j] += acc[7][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[2][j]; - acc[4][j] += acc[6][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[4][j]; - } - } else if constexpr (TOP_K == 6) { -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[1][j]; - acc[2][j] += acc[3][j]; - acc[4][j] += acc[5][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[2][j]; - acc[0][j] += acc[4][j]; - } - } else if constexpr (TOP_K == 4) { -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[1][j]; - acc[2][j] += acc[3][j]; - } -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[2][j]; - } - } else if constexpr (TOP_K == 2) { -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[1][j]; - } - } else if constexpr (TOP_K == 1) { - // nothing to do - } else { - // Generic fallback: accumulate all into acc[0] -#pragma unroll - for (int k = 1; k < TOP_K; ++k) { -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - acc[0][j] += acc[k][j]; - } - } - } - // Epilogue: store the float accumulator as T (NONE) or quantize to MXFP8 + scale factor. - if constexpr (QuantMode == MoeA2ACombineQuantMode::NONE) { - // cast_store: converts float->T element-by-element then writes via vectorized int4 store. - acc[0].cast_store(dst_typed_base + offset / static_cast(sizeof(InT))); - } else { - constexpr uint32_t sf_vec_size = QuantMode == MoeA2ACombineQuantMode::NVFP4 ? 16 : 32; - constexpr uint32_t threads_per_sf = sf_vec_size / elems_per_vec; - uint8_t scale; - auto store_sf = [&]() { - int64_t sf_offset; - const int64_t num_vecs_per_row = (row_size + sf_vec_size - 1) / sf_vec_size; - const int64_t sf_idx = logical_offset / sf_vec_size; - if constexpr (SwizzleMode == MoeA2ACombineSwizzleSFMode::LINEAR) { - sf_offset = row_idx * num_vecs_per_row + sf_idx; - } else if constexpr (SwizzleMode == MoeA2ACombineSwizzleSFMode::SWIZZLE_128x4) { - sf_offset = get_sf_out_offset_128x4(row_idx, sf_idx, num_vecs_per_row); - } else { // SWIZZLE_8x4 - sf_offset = get_sf_out_offset_8x4(row_idx, sf_idx, num_vecs_per_row); - } - if (threadIdx.x % threads_per_sf == 0) { - reinterpret_cast(sf_output)[sf_offset] = scale; - } - }; - // Cast the float32 accumulator down to T to feed the MXFP8/FP4 converter's PackedVec input. - tensorrt_llm::kernels::PackedVec packed_vec; - T* packed_elems = reinterpret_cast(&packed_vec); -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - packed_elems[j] = static_cast(acc[0][j]); - } - if constexpr (QuantMode == MoeA2ACombineQuantMode::MXFP8) { - static_assert(elems_per_vec == 8, "MXFP8 quantization requires 8 elements per vector"); - uint64_t fp8x8 = - tensorrt_llm::kernels::cvt_warp_fp16_to_mxfp8(packed_vec, &scale); - reinterpret_cast(dst_bytes)[logical_offset / elems_per_vec] = fp8x8; - } else if constexpr (QuantMode == MoeA2ACombineQuantMode::NVFP4 || - QuantMode == MoeA2ACombineQuantMode::MXFP4) { - static_assert(elems_per_vec == 8 || elems_per_vec == 16, - "FP4 quantization requires 8 or 16 elements per vector"); - constexpr int SF_VEC_SIZE = QuantMode == MoeA2ACombineQuantMode::MXFP4 ? 32 : 16; - auto fp4_packed = tensorrt_llm::kernels::cvt_warp_fp16_to_fp4 < T, SF_VEC_SIZE, - elems_per_vec, - QuantMode == MoeA2ACombineQuantMode::MXFP4 > (packed_vec, OutputScalarScale, &scale); - // cvt_warp_fp16_to_fp4 returns uint32_t for 8 elems (4 packed bytes) and uint64_t for 16 - // (8 packed bytes); store at the matching width so packed rows stay contiguous and fit the - // dim/2-byte output buffer (a wider store would overflow into the next token's row). - reinterpret_cast(dst_bytes)[logical_offset / elems_per_vec] = - fp4_packed; - } - store_sf(); - } - } -} - -// Wrapper that selects vector width based on size_per_token alignment. -// QuantMode==NONE writes T (optionally upcast from an fp8 InT); QuantMode==MXFP8 writes packed -// fp8 + scale factors (always InT == T, fixed 16-byte / 8-element vectors). -// stride_per_token: byte distance between tokens in the recv buffer (may differ from -// size_per_token when FP8 in-place uses a BF16-stride workspace with an FP8-sized payload). -// InT: input element type in recv buffer (defaults to T for same-type accumulation). -template -__device__ void vectorized_combine(void* output_buffer, void* sf_output, int row_idx, int row_size, - int stride_per_token, int rank_id, int max_tokens_per_rank, - CombineKernelPointers const& ptrs, - float OutputScalarScale = 1.0f) { - const int size_per_token = row_size * static_cast(sizeof(InT)); - if constexpr (QuantMode != MoeA2ACombineQuantMode::NONE) { - // MXFP8/FP4 require 16-byte vectors (8 or 16 elements per vec), always InT == T. - vectorized_combine_impl<16, TOP_K, T, InT, QuantMode, SwizzleMode>( - output_buffer, sf_output, row_idx, row_size, stride_per_token, rank_id, max_tokens_per_rank, - ptrs, OutputScalarScale); - } else { - // Each branch is guarded by if constexpr (sizeof(InT) <= VEC_SIZE_BYTES) so the compiler - // never instantiates vectorized_combine_impl with elems_per_vec == 0. - if (size_per_token % 16 == 0) { - if constexpr (static_cast(sizeof(InT)) <= 16) - vectorized_combine_impl<16, TOP_K, T, InT>(output_buffer, nullptr, row_idx, row_size, - stride_per_token, rank_id, max_tokens_per_rank, - ptrs); - } else if (size_per_token % 8 == 0) { - if constexpr (static_cast(sizeof(InT)) <= 8) - vectorized_combine_impl<8, TOP_K, T, InT>(output_buffer, nullptr, row_idx, row_size, - stride_per_token, rank_id, max_tokens_per_rank, - ptrs); - } else if (size_per_token % 4 == 0) { - if constexpr (static_cast(sizeof(InT)) <= 4) - vectorized_combine_impl<4, TOP_K, T, InT>(output_buffer, nullptr, row_idx, row_size, - stride_per_token, rank_id, max_tokens_per_rank, - ptrs); - } else if (size_per_token % 2 == 0) { - if constexpr (static_cast(sizeof(InT)) <= 2) - vectorized_combine_impl<2, TOP_K, T, InT>(output_buffer, nullptr, row_idx, row_size, - stride_per_token, rank_id, max_tokens_per_rank, - ptrs); - } else { - if constexpr (static_cast(sizeof(InT)) <= 1) - vectorized_combine_impl<1, TOP_K, T, InT>(output_buffer, nullptr, row_idx, row_size, - stride_per_token, rank_id, max_tokens_per_rank, - ptrs); - } - } -} - -// ---- vec_convert: per-vector type conversion, specialized by PTX where available ---- -// Generic: SrcT -> float -> DstT (all architectures, all type combinations). -template -__device__ __forceinline__ void vec_convert(flashinfer::vec_t& out, - flashinfer::vec_t const& in) { -#pragma unroll - for (int j = 0; j < VEC_SIZE; ++j) out[j] = DstT(static_cast(in[j])); -} - -// BF16 -> FP8 e4m3: use CUDA intrinsic (SM100+, Blackwell). -// Inline PTX "cvt.rn.satfinite.e4m3x2.bf16x2 %h, %r" is rejected by SM100a ptxas -// ("Unexpected instruction types for cvt") because SM100a requires a 32-bit output -// register for this instruction. __nv_fp8x2_e4m3(bfloat162) emits the correct form. -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -template = 0> -__device__ __forceinline__ void vec_convert(flashinfer::vec_t<__nv_fp8_e4m3, VEC_SIZE>& out, - flashinfer::vec_t<__nv_bfloat16, VEC_SIZE> const& in) { - __nv_fp8x2_e4m3* out_fp8x2 = reinterpret_cast<__nv_fp8x2_e4m3*>(&out); - __nv_bfloat162 const* in_bf16x2 = reinterpret_cast<__nv_bfloat162 const*>(&in); -#pragma unroll - for (int p = 0; p < static_cast(VEC_SIZE) / 2; ++p) - out_fp8x2[p] = __nv_fp8x2_e4m3(in_bf16x2[p]); -} -#endif - -// FP16 -> FP8 e4m3: paired PTX cvt.rn.satfinite.e4m3x2.f16x2 (SM89+, Hopper). -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) -template = 0> -__device__ __forceinline__ void vec_convert(flashinfer::vec_t<__nv_fp8_e4m3, VEC_SIZE>& out, - flashinfer::vec_t const& in) { - uint32_t const* src_u32 = reinterpret_cast(&in); - uint16_t* dst_u16 = reinterpret_cast(&out); -#pragma unroll - for (int p = 0; p < VEC_SIZE / 2; ++p) { - uint16_t d; - asm volatile("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(d) : "r"(src_u32[p])); - dst_u16[p] = d; - } -} -#endif - -// ---- vectorized_quant_impl: load -> sync -> convert -> store ---- -// VEC_SIZE is in elements (not bytes), so both SrcT and DstT vectors hold VEC_SIZE values. -// Block-per-token: threadIdx.x indexes vectors within the token, blockDim.x is the stride. -template -__device__ void vectorized_quant_impl(DstT* dst, SrcT const* src, int num_elements) { - using flashinfer::vec_t; - - // num_elements is a multiple of VEC_SIZE here (selected by the wrapper). - int const num_vecs = num_elements / VEC_SIZE; - // Round the iteration bound up to a multiple of blockDim.x so every thread reaches the - // __syncthreads() the same number of times (a divergent barrier would deadlock when - // num_vecs is not a multiple of blockDim.x, e.g. hidden=7168 with 256 threads). - int const vec_iters = ((num_vecs + blockDim.x - 1) / blockDim.x) * blockDim.x; - - for (int iter = threadIdx.x; iter < vec_iters; iter += blockDim.x) { - bool const active = iter < num_vecs; - int const e = iter * VEC_SIZE; - vec_t in_vec; - if (active) in_vec.load(src + e); - - // Sync so all threads finish loading inputs before any thread writes output. This avoids - // write-after-read hazards in the FP8 in-place case where this kernel's output is read as - // input on the next iteration. All threads participate (active guards the work, not the sync). - __syncthreads(); - - if (active) { - vec_t out_vec; - vec_convert(out_vec, in_vec); - out_vec.store(dst + e); - } - } -} - -template -__device__ void vectorized_quant(DstT* dst, SrcT const* src, int num_elements) { - if (num_elements % 16 == 0) - vectorized_quant_impl<16, SrcT, DstT>(dst, src, num_elements); - else if (num_elements % 8 == 0) - vectorized_quant_impl<8, SrcT, DstT>(dst, src, num_elements); - else if (num_elements % 4 == 0) - vectorized_quant_impl<4, SrcT, DstT>(dst, src, num_elements); - else if (num_elements % 2 == 0) - vectorized_quant_impl<2, SrcT, DstT>(dst, src, num_elements); - else - vectorized_quant_impl<1, SrcT, DstT>(dst, src, num_elements); -} - -// LOW_PRECISION=false: vectorized byte-copy (SrcT = payload dtype). -// LOW_PRECISION=true: vectorized SrcT->FP8 quantization via vectorized_quant. -// stride_per_token: byte distance between tokens in recv_buffer_bytes (host-computed, avoids -// per-thread recomputation): -// - FP8 external payload: elements_per_token x 1 (compact FP8 layout) -// - FP8 in-place / byte-copy: elements_per_token x sizeof(SrcT) (payload-dtype stride) -// Copy payload to recv buffer; one block per token. -template -__global__ void moeA2APrepareCombineKernel(uint8_t* recv_buffer_bytes, void const* payload, - int elements_per_token, int ep_size, - int max_tokens_per_rank, uint32_t* flag_val_ptr, - int const* recv_counters, int stride_per_token, - bool enable_pdl) { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (enable_pdl) cudaGridDependencySynchronize(); - cudaTriggerProgrammaticLaunchCompletion(); -#endif - - if (blockIdx.x == 0 && threadIdx.x == 0) { - // Increment flag_val for this combine round - *flag_val_ptr = *flag_val_ptr + 1; - } - - // Copy path: null payload means data is already in workspace — nothing to do. - if (!LOW_PRECISION && payload == nullptr) return; - - int global_token_idx = blockIdx.x; - - int global_token_num = ep_size * max_tokens_per_rank; - if (global_token_idx >= global_token_num) return; - - // Map global_token_idx to (rank_idx, local_token_idx) - int rank_idx = global_token_idx / max_tokens_per_rank; - int local_token_idx = global_token_idx % max_tokens_per_rank; - - // Skip invalid tokens beyond per-rank recv count - if (local_token_idx >= recv_counters[rank_idx]) return; - - size_t const token_offset = static_cast(global_token_idx) * stride_per_token; - - if constexpr (LOW_PRECISION) { - // Source pointer: external payload or in-place from workspace. - SrcT const* src_ptr = (payload != nullptr) - ? static_cast(payload) + - static_cast(global_token_idx) * elements_per_token - : reinterpret_cast(recv_buffer_bytes + token_offset); - - // Destination: stride_per_token encodes the correct layout for both paths - // (compact FP8 for external, payload-dtype stride for in-place). - __nv_fp8_e4m3* dst_ptr = reinterpret_cast<__nv_fp8_e4m3*>(recv_buffer_bytes + token_offset); - - vectorized_quant(dst_ptr, src_ptr, elements_per_token); - } else { - // Generic byte copy (payload guaranteed non-null by early return above). - vectorized_copy(recv_buffer_bytes + token_offset, - static_cast(payload) + token_offset, stride_per_token); - } -} - -// ============================================================================ -// Generic Combine Kernel Implementation (Templated by data type) -// ============================================================================ - -template -__global__ void moeA2ACombineKernel( - const CombineKernelPointers ptrs, // Combine-specific struct, src_data_ptrs[0] is output - int max_tokens_per_rank, int elements_per_token, int stride_per_token, int local_num_tokens, - int rank_id, int ep_size, bool enable_pdl, float OutputScalarScale) { - int local_token_idx = blockIdx.x; - - if (local_num_tokens == 0) { - // Special case: If local_num_tokens == 0, - // we need to keep the threads where local_token_idx == 0 alive to participate in the - // synchronization. Other threads should return. - if (local_token_idx > 0) return; - } else { - // Threads that do not have a token to process should return. - if (local_token_idx >= local_num_tokens) return; - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (enable_pdl) cudaGridDependencySynchronize(); -#endif - -#if !DISABLE_SYNC_FOR_PROFILING - // In-kernel readiness synchronization at start of combine: - // - One warp signals readiness to all peers with current flag_val. - // - The first warp of each block waits for all peers' readiness (equality), then __syncthreads. - bool is_first_warp = threadIdx.x / warpSize == 0; - if (is_first_warp) { - int lane_id = threadIdx.x % warpSize; - uint32_t expected_value = *ptrs.flag_val; - - if (blockIdx.x == 0) { - // asm volatile("fence.release.sys;"); - // Signal readiness to all active peers; skip dead ranks (their symmetric memory is - // unreachable). -#pragma unroll 1 // No unroll - for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize) { - if constexpr (ENABLE_RANK_MASK) { - if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) continue; - } - uint32_t* flag_addr = &ptrs.completion_flags[peer_rank][rank_id]; - asm volatile("st.relaxed.sys.u32 [%0], %1;" ::"l"(flag_addr), "r"(expected_value)); -#if ENABLE_DEBUG_PRINT - printf("combine: +++Rank %d setting completion flag to %d for rank %d\n", rank_id, - expected_value, peer_rank); -#endif - } - } - -// Wait for all active peers to signal; skip dead ranks (otherwise we would spin forever — -// this is the bug the rank-mask is here to prevent). -#pragma unroll 1 // No unroll - for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize) { - if constexpr (ENABLE_RANK_MASK) { - if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) continue; - } - bool flag_set = false; - [[maybe_unused]] auto s = clock64(); - do { - uint32_t* flag_ptr = &ptrs.completion_flags[rank_id][peer_rank]; - uint32_t flag_value; - // Acquire load to ensure visibility of peer's release-store - asm volatile("ld.relaxed.sys.u32 %0, [%1];" : "=r"(flag_value) : "l"(flag_ptr)); -#if ENABLE_DEBUG_PRINT - printf( - "combine: ---Rank %d received completion flag from rank %d, flag_value: %d, " - "expected_value: %d, " - "address: %p\n", - rank_id, peer_rank, flag_value, expected_value, flag_ptr); -#endif - flag_set = flag_value == expected_value; - } while (!flag_set && !check_timeout(s)); - - if (__builtin_expect(!flag_set, 0)) { - printf("combine: ---Rank %d timed out waiting for completion flag from rank %d\n", rank_id, - peer_rank); - asm volatile("trap;"); - return; - } - } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("fence.acquire.sys;"); -#else - __threadfence_system(); -#endif - } - __syncthreads(); -#endif - - if (local_num_tokens == 0) return; - - // Dispatch the combine path: - // - T == fp8_e4m3: low-precision combine, FP8 recv buffer -> BF16 output (InT=fp8, NONE quant). - // - otherwise: same-type accumulate, then store as T (NONE) or quantize to MXFP8/FP4 - // (QuantMode). - // vectorized_combine derives the per-token output pointer from row_idx (= local_token_idx). - if constexpr (std::is_same_v) { - // src_data_ptrs[0] points to a BF16 output buffer (set by moeA2ACombineOp for low precision). - vectorized_combine( - ptrs.src_data_ptrs[0], nullptr, local_token_idx, elements_per_token, stride_per_token, - rank_id, max_tokens_per_rank, ptrs); - } else { - vectorized_combine( - ptrs.src_data_ptrs[0], ptrs.output_scales, local_token_idx, elements_per_token, - stride_per_token, rank_id, max_tokens_per_rank, ptrs, OutputScalarScale); - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (enable_pdl) cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -void moe_a2a_prepare_combine_launch(MoeA2ACombineParams const& params) { - constexpr int kBlockSize = 256; - - // FP8 in-place (use_low_precision, prepare_payload==nullptr): each CTA writes FP8 at the - // BF16-stride position, so CTAs never race — all tokens must be processed. Copy path with a - // null payload is a no-op; 1 block suffices for the flag increment only. - int grid = (params.use_low_precision || params.prepare_payload != nullptr) - ? params.ep_size * params.max_tokens_per_rank - : 1; - - uint8_t* recv_buffer_bytes = - static_cast(const_cast(params.recv_buffers[params.ep_rank])); - void const* payload = params.prepare_payload; - - // stride_per_token (host-computed, avoids per-thread recompute): - // FP8 external: EPT x 1 (compact FP8, dst packed tightly) - // FP8 in-place / byte-copy: EPT x sizeof(SrcT) (payload-dtype stride) - SWITCH_BOOL(params.use_low_precision, LOW_PRECISION, { - SWITCH_DTYPE(params.dtype, SrcT, { - bool const low_precision_staged = LOW_PRECISION && (params.prepare_payload != nullptr); - int const stride_per_token = low_precision_staged - ? params.elements_per_token - : params.elements_per_token * static_cast(sizeof(SrcT)); - auto kernel_fn = moeA2APrepareCombineKernel; - launchWithPdlWhenEnabled("moeA2APrepareCombineKernel", params.enable_pdl, kernel_fn, grid, - kBlockSize, 0, params.stream, recv_buffer_bytes, payload, - params.elements_per_token, params.ep_size, - params.max_tokens_per_rank, params.flag_val, params.recv_counters, - stride_per_token, params.enable_pdl); - }); - }); -} - -// ============================================================================ -// Combine Launch Function -// ============================================================================ - -// The trtllm backend always writes the BF16/FP16 intermediate and quantizes it in a second -// launch; only the generated cake backend fuses the two. -bool moe_a2a_combine_fuses_quantization(MoeA2ACombineParams const&) { return false; } - -void moe_a2a_combine_launch(MoeA2ACombineParams const& params) { - // Validate parameters - TLLM_CHECK(params.top_k > 0 && params.top_k <= kMaxTopK); - TLLM_CHECK(params.ep_size > 0 && params.ep_size <= kMaxRanks); - TLLM_CHECK(params.ep_rank >= 0 && params.ep_rank < params.ep_size); - TLLM_CHECK(params.local_num_tokens >= 0); - TLLM_CHECK(params.elements_per_token > 0); - if (params.enable_rank_mask) { - // The local rank must always be marked active in its own view of the mask; otherwise the - // kernel itself would be running on a "dead" rank. - TLLM_CHECK_WITH_INFO( - (params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL, - "active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank); - } - - // Configure kernel launch: one block per token - int const kBlockSize = tensorrt_llm::common::getEnvMoeA2ACombineBlockSize(); - int grid_size_block = params.local_num_tokens; - // If local_num_tokens is 0, we still need to launch a minimal kernel to participate in the - // synchronization. - if (grid_size_block == 0) { - grid_size_block = 1; - } - - // Prepare kernel pointers struct for combine - CombineKernelPointers kernel_ptrs = {}; // Zero-initialize - - // Set output data pointer in src_data_ptrs[0] - kernel_ptrs.src_data_ptrs[0] = params.output_data; - kernel_ptrs.output_scales = params.output_scales; - - // Fill recv buffer pointers - for (int rank = 0; rank < params.ep_size; rank++) { - kernel_ptrs.recv_buffers[rank][0] = params.recv_buffers[rank]; - } - - // Copy completion flag pointers - for (int i = 0; i < params.ep_size; i++) { - kernel_ptrs.completion_flags[i] = params.completion_flags[i]; - } - kernel_ptrs.flag_val = params.flag_val; - - // Copy communication tracking pointers - kernel_ptrs.topk_target_ranks = params.topk_target_ranks; - kernel_ptrs.topk_send_indices = params.topk_send_indices; - - // Copy active-rank bitmask into the kernel pointers struct - for (int w = 0; w < kRankMaskWords; ++w) { - kernel_ptrs.active_rank_mask[w] = params.active_rank_mask[w]; - } - - int grid = grid_size_block; // one block per token - - // stride_per_token: byte distance between tokens in the recv buffer. - // FP8 external payload: EPT × 1 (compact FP8 layout) - // FP8 in-place / non-FP8: EPT × sizeof(PayloadT) (payload-dtype stride) - bool const low_precision_staged = params.use_low_precision && (params.prepare_payload != nullptr); - int stride_per_token; - SWITCH_DTYPE(params.dtype, PayloadT, { - stride_per_token = low_precision_staged - ? params.elements_per_token - : params.elements_per_token * static_cast(sizeof(PayloadT)); - }); - - // When use_low_precision is set the recv buffers contain FP8 data regardless of params.dtype, - // so dispatch the FP8 accumulation kernel in that case. - auto const effective_dtype = params.use_low_precision ? nvinfer1::DataType::kFP8 : params.dtype; - - // Launch appropriate kernel with compact macros - SWITCH_BOOL(params.enable_rank_mask, ENABLE_RANK_MASK, { - SWITCH_DTYPE(effective_dtype, TKernelType, { - SWITCH_TOP_K(params.top_k, TOP_K, { - SWITCH_QUANT_MODE(TKernelType, params.quant_mode, QUANT_MODE, { - SWITCH_SWIZZLE_MODE(params.swizzle_mode, SWIZZLE_MODE, { - auto kernel_fn = - moeA2ACombineKernel; - launchWithPdlWhenEnabled("moeA2ACombineKernel", params.enable_pdl, kernel_fn, grid, - kBlockSize, 0, params.stream, kernel_ptrs, - params.max_tokens_per_rank, params.elements_per_token, - stride_per_token, params.local_num_tokens, params.ep_rank, - params.ep_size, params.enable_pdl, params.output_scalar_scale); - }); - }); - }); - }); - }) -} - -// Kernel to sanitize expert ids for invalid tokens -__global__ void moeA2ASanitizeExpertIdsKernel(int32_t* expert_ids_ptr, - int32_t const* recv_counters_ptr, int ep_size, - int max_tokens_per_rank, int top_k, - int32_t invalid_id, bool enable_pdl) { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (enable_pdl) cudaGridDependencySynchronize(); - cudaTriggerProgrammaticLaunchCompletion(); -#endif - int tid = blockIdx.x * blockDim.x + threadIdx.x; - int total_tokens = ep_size * max_tokens_per_rank; - if (tid >= total_tokens) return; - - int source_rank = tid / max_tokens_per_rank; - int token_idx = tid % max_tokens_per_rank; - - if (token_idx >= recv_counters_ptr[source_rank]) { - int32_t* token_expert_ids = expert_ids_ptr + tid * top_k; - for (int k = 0; k < top_k; ++k) { - token_expert_ids[k] = invalid_id; - } - } -} - -void moe_a2a_sanitize_expert_ids_launch(int32_t* expert_ids, int32_t const* recv_counters, - int32_t invalid_id, int ep_size, int max_tokens_per_rank, - int top_k, cudaStream_t stream, bool enable_pdl) { - constexpr int kBlockSize = 256; - int total_tokens = ep_size * max_tokens_per_rank; - int grid = ceilDiv(total_tokens, kBlockSize); - launchWithPdlWhenEnabled("moeA2ASanitizeExpertIdsKernel", enable_pdl, - moeA2ASanitizeExpertIdsKernel, grid, kBlockSize, 0, stream, expert_ids, - recv_counters, ep_size, max_tokens_per_rank, top_k, invalid_id, - enable_pdl); -} - -} // namespace tensorrt_llm::kernels::moe_alltoall diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllPrepareDispatch.cu b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllPrepareDispatch.cu deleted file mode 100644 index 4fdb7c9e945..00000000000 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllPrepareDispatch.cu +++ /dev/null @@ -1,55 +0,0 @@ -/* - * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef signed int int32_t; -typedef short int int16_t; -typedef struct __align__(64) { - uint64_t opaque[16]; -} CUtensorMap; - -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x)); - return result; -} - -#define NUM_MAIN_STAGES 1 -#define THREADS 64 - -#include - -extern "C" { - -__global__ __launch_bounds__(64) void kernel_flashinfer_mnnvl_moe_alltoall_prepare_dispatch( - int* __restrict__ send_counters, int* __restrict__ local_token_counter, int ep_size, - unsigned int* __restrict__ flag_val, bool enable_pdl) { - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - // === Task calls (dependency order) === - if (enable_pdl) { - asm volatile("griddepcontrol.wait;" ::: "memory"); - } - asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); - if (tid < ep_size) { - send_counters[tid] = 0; - } - if (tid == 0) { - local_token_counter[0] = 0; - flag_val[0] = flag_val[0] + 1; - } -} - -} // extern "C" diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllPublishCombine.cu b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllPublishCombine.cu deleted file mode 100644 index a4caa0e6a6c..00000000000 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllPublishCombine.cu +++ /dev/null @@ -1,114 +0,0 @@ -/* - * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef signed int int32_t; -typedef short int int16_t; -typedef struct __align__(64) { - uint64_t opaque[16]; -} CUtensorMap; - -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x)); - return result; -} - -#define NUM_MAIN_STAGES 1 -#define THREADS 64 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - -extern "C" { - -__global__ __launch_bounds__(64) void kernel_flashinfer_mnnvl_moe_alltoall_publish_combine( - uint8_t* __restrict__ workspace, unsigned long long workspace_stride_bytes, - unsigned long long flag_val_offset, unsigned long long completion_flags_offset, int ep_rank, - int ep_size, bool enable_pdl, bool enable_rank_mask, unsigned long long active_rank_mask) { - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - // === Task calls (dependency order) === - unsigned int* workspace_u32 = reinterpret_cast(workspace); - unsigned long long local_workspace_base = (unsigned long long)ep_rank * workspace_stride_bytes; - if (enable_pdl) { - asm volatile("griddepcontrol.wait;" ::: "memory"); - } - unsigned long long expected_index = (local_workspace_base + flag_val_offset) / 4; - unsigned int expected_value = workspace_u32[expected_index]; - asm volatile("fence.release.sys;" ::: "memory"); - if (tid < ep_size) { - int target_rank = tid; - int target_is_active = 1; - if (enable_rank_mask) { - unsigned long long rank_mask_bit = active_rank_mask >> (unsigned long long)target_rank & 1; - target_is_active = ((rank_mask_bit != 0) ? 1 : 0); - } - if (target_is_active != 0) { - unsigned long long remote_flag_index = - ((unsigned long long)target_rank * workspace_stride_bytes + completion_flags_offset) / 4 + - (unsigned long long)ep_rank; - asm volatile("st.relaxed.sys.u32 [%0], %1;" ::"l"( - (reinterpret_cast(workspace_u32) + (remote_flag_index))), - "r"(static_cast(expected_value)) - : "memory"); - unsigned long long local_flag_index = - (local_workspace_base + completion_flags_offset) / 4 + (unsigned long long)target_rank; - { - unsigned int* _sre_ptr_0 = - (reinterpret_cast(workspace_u32) + (local_flag_index)); - const unsigned int _sre_expected_0 = static_cast(expected_value); - const unsigned long long _sre_start_0 = clock64(); - bool _sre_matched_0 = false; - do { - unsigned int _sre_value_0; - asm volatile("ld.relaxed.sys.u32 %0, [%1];" - : "=r"(_sre_value_0) - : "l"(_sre_ptr_0) - : "memory"); - _sre_matched_0 = (_sre_value_0 == _sre_expected_0); - } while (!_sre_matched_0 && - ((clock64() - _sre_start_0) <= static_cast(600000000000))); - if (__builtin_expect(!_sre_matched_0, 0)) { - asm volatile("trap;"); - return; - } - } - } - } - __syncthreads(); - asm volatile("fence.acquire.sys;" ::: "memory"); - if (enable_pdl) { - if (warp == 0) { - if (elect_sync()) { - asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); - } - } - } -} - -} // extern "C" diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllQuantizeCombine.cu b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllQuantizeCombine.cu deleted file mode 100644 index 57294f600a7..00000000000 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllQuantizeCombine.cu +++ /dev/null @@ -1,226 +0,0 @@ -/* - * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef signed int int32_t; -typedef short int int16_t; -typedef struct __align__(64) { - uint64_t opaque[16]; -} CUtensorMap; - -#include -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x)); - return result; -} - -#define NUM_MAIN_STAGES 1 -#define THREADS 32 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - -__device__ __forceinline__ float reciprocal_approximate_ftz(float value) { - float result; - asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(value)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(32) void kernel_flashinfer_mnnvl_moe_alltoall_quantize_combine( - float* __restrict__ accumulated, uint8_t* __restrict__ quantized_fp8, - uint8_t* __restrict__ quantized_packed, uint8_t* __restrict__ scales_u8, - uint8_t* __restrict__ scales_fp8, int elements_per_token, int payload_dtype_code, - int quant_mode, int scale_layout, float output_scalar_scale, int blocks_per_row, - int padded_scale_cols, bool enable_pdl) { - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - // === Task calls (dependency order) === - if (enable_pdl) { - asm volatile("griddepcontrol.wait;" ::: "memory"); - } - int logical_block = bid; - int token = logical_block / blocks_per_row; - int block_column = logical_block - token * blocks_per_row; - int block_size = 32; - if (quant_mode == 3) { - block_size = 16; - } - int column = block_column * block_size + lane; - int active_lane = ((block_size > lane && elements_per_token > column) ? 1 : 0); - float value = 0.0f; - if (active_lane != 0) { - value = accumulated[token * elements_per_token + column]; - if (payload_dtype_code == 0) { - value = __bfloat162float(__float2bfloat16(value)); - } - if (payload_dtype_code == 1) { - value = __half2float(__float2half_rn(value)); - } - } - float scaled_value = value; - float _fabs_0 = fabsf(scaled_value); - float block_max = _fabs_0; - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, block_max, 16); - float _max_0 = max_noftz(block_max, _shfl_xor_0); - block_max = _max_0; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_max, 8); - float _max_1 = max_noftz(block_max, _shfl_xor_1); - block_max = _max_1; - float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, block_max, 4); - float _max_2 = max_noftz(block_max, _shfl_xor_2); - block_max = _max_2; - float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, block_max, 2); - float _max_3 = max_noftz(block_max, _shfl_xor_3); - block_max = _max_3; - float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, block_max, 1); - float _max_4 = max_noftz(block_max, _shfl_xor_4); - block_max = _max_4; - float _shfl_0 = __shfl_sync(0xFFFFFFFF, block_max, 0); - block_max = _shfl_0; - int scale_byte = 0; - float actual_scale = 0.0f; - float fp8_scale = 0.0f; - if (quant_mode == 3) { - float sf_value = output_scalar_scale * (block_max * reciprocal_approximate_ftz(6.0f)); - float _fp8_rt_0; - uint16_t _e4m3x2_0; - uint32_t _f16x2_0; - asm volatile("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" - : "=h"(_e4m3x2_0) - : "f"(0.0f), "f"(sf_value)); - asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(_f16x2_0) : "h"(_e4m3x2_0)); - uint16_t _fp8_h0_0 = (uint16_t)(_f16x2_0 & 0xFFFFu); - asm volatile("cvt.f32.f16 %0, %1;" : "=f"(_fp8_rt_0) : "h"(_fp8_h0_0)); - fp8_scale = _fp8_rt_0; - actual_scale = block_max != 0.0f - ? reciprocal_approximate_ftz(fp8_scale * - reciprocal_approximate_ftz(output_scalar_scale)) - : 0.0f; - } else { - float denominator = 6.0f; - if (quant_mode == 1) { - denominator = 448.0f; - } - float raw_scale = block_max * reciprocal_approximate_ftz(denominator); - int raw_bits = 0; - raw_bits = reinterpret_cast(&raw_scale)[0]; - int exponent = raw_bits >> 23 & 255; - int mantissa = raw_bits & 8388607; - int round_up = 0; - if (mantissa != 0) { - if (exponent != 0) { - round_up = 1; - } else if (mantissa > 4194304) { - round_up = 1; - } - } - if (raw_scale > 0.0f) { - scale_byte = exponent + round_up; - if (scale_byte > 254) { - scale_byte = 254; - } - } - if (quant_mode == 1) { - int scale_bits = scale_byte << 23; - if (scale_byte == 0) { - scale_bits = 0x00400000; - } - float decoded_scale = reinterpret_cast(&scale_bits)[0]; - actual_scale = reciprocal_approximate_ftz(decoded_scale); - } else if (block_max != 0.0f) { - actual_scale = scale_byte == 0 ? 1.0f : exp2f(127.0f - scale_byte); - } - } - int scale_index = token * blocks_per_row + block_column; - if (scale_layout == 0) { - scale_index = block_column % 4 + block_column / 4 * 512 + token % 32 * 16 + - token % 128 / 32 * 4 + token / 128 * (128 * padded_scale_cols); - } - if (scale_layout == 1) { - int tiles = padded_scale_cols / 4; - scale_index = - token / 8 * (tiles * 32) + block_column / 4 * 32 + token % 8 * 4 + block_column % 4; - } - if (lane == 0) { - if (quant_mode == 3) { - { - unsigned short _fp8_pair; - asm("cvt.rn.satfinite.e4m3x2.f32 %0, 0f00000000, %1;" : "=h"(_fp8_pair) : "f"(fp8_scale)); - *(reinterpret_cast(scales_fp8 + scale_index) + (0)) = - (unsigned char)(_fp8_pair & 0xFF); - } - } else { - scales_u8[scale_index] = scale_byte; - } - } - float normalized = value * actual_scale; - if (quant_mode == 1) { - if (active_lane != 0) { - { - unsigned short _fp8_pair; - asm("cvt.rn.satfinite.e4m3x2.f32 %0, 0f00000000, %1;" : "=h"(_fp8_pair) : "f"(normalized)); - *(reinterpret_cast(quantized_fp8 + (token * elements_per_token + column)) + - (0)) = (unsigned char)(_fp8_pair & 0xFF); - } - } - } else { - float partner = __shfl_xor_sync(0xFFFFFFFF, normalized, 1); - if (active_lane != 0) { - if ((lane & 1) == 0 && column + 1 < elements_per_token) { - uint32_t packed; - asm volatile( - "{\n" - ".reg .b8 byte0;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" - "mov.b32 %0, {byte0, 0, 0, 0};\n" - "}" - : "=r"(packed) - : "f"(normalized), "f"(partner)); - int packed_column = block_column * (block_size / 2) + lane / 2; - quantized_packed[token * (elements_per_token / 2) + packed_column] = packed; - } - } - } - if (enable_pdl) { - if (warp == 0) { - if (elect_sync()) { - asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); - } - } - } -} - -} // extern "C" diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllSanitize.cu b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllSanitize.cu deleted file mode 100644 index 6807ebed35c..00000000000 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllSanitize.cu +++ /dev/null @@ -1,61 +0,0 @@ -/* - * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef signed int int32_t; -typedef short int int16_t; -typedef struct __align__(64) { - uint64_t opaque[16]; -} CUtensorMap; - -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x)); - return result; -} - -#define NUM_MAIN_STAGES 1 -#define THREADS 256 - -#include - -extern "C" { - -__global__ __launch_bounds__(256) void kernel_flashinfer_mnnvl_moe_alltoall_sanitize_expert_ids( - int* __restrict__ expert_ids, int* __restrict__ recv_counters, int ep_size, - int max_tokens_per_rank, int top_k, int invalid_id, int enable_pdl) { - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - // === Task calls (dependency order) === - if (enable_pdl != 0) { - asm volatile("griddepcontrol.wait;" ::: "memory"); - } - asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); - int tid_0 = bid * 256 + tid; - int total_tokens = ep_size * max_tokens_per_rank; - if (tid_0 < total_tokens) { - int source_rank = tid_0 / max_tokens_per_rank; - int token_idx = tid_0 - source_rank * max_tokens_per_rank; - if (token_idx >= recv_counters[source_rank]) { - int token_base = tid_0 * top_k; -#pragma unroll 1 - for (int k = 0; k < top_k; k++) { - expert_ids[token_base + k] = invalid_id; - } - } - } -} - -} // extern "C" diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllStageCombine.cu b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllStageCombine.cu deleted file mode 100644 index 91c4cb4d165..00000000000 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllStageCombine.cu +++ /dev/null @@ -1,65 +0,0 @@ -/* - * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef signed int int32_t; -typedef short int int16_t; -typedef struct __align__(64) { - uint64_t opaque[16]; -} CUtensorMap; - -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x)); - return result; -} - -#define NUM_MAIN_STAGES 1 -#define THREADS 256 - -#include - -extern "C" { - -__global__ __launch_bounds__(256) void kernel_flashinfer_mnnvl_moe_alltoall_stage_combine( - uint8_t* __restrict__ payload, uint8_t* __restrict__ workspace, - unsigned long long workspace_stride_bytes, unsigned long long flag_val_offset, - unsigned long long destination_offset, unsigned long long payload_bytes, int ep_rank, - bool enable_pdl) { - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - // === Task calls (dependency order) === - if (enable_pdl) { - asm volatile("griddepcontrol.wait;" ::: "memory"); - } - asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); - if (bid == 0) { - if (tid == 0) { - unsigned int* workspace_u32 = reinterpret_cast(workspace); - unsigned long long flag_index = - ((unsigned long long)ep_rank * workspace_stride_bytes + flag_val_offset) / 4; - workspace_u32[flag_index] = workspace_u32[flag_index] + 1; - } - } - unsigned long long linear_thread = (unsigned long long)bid * 256 + (unsigned long long)tid; - unsigned long long thread_stride = (unsigned long long)num_bids * 256; -#pragma unroll 1 - for (unsigned long long byte_index = linear_thread; byte_index < payload_bytes; - byte_index += thread_stride) { - workspace[destination_offset + byte_index] = payload[byte_index]; - } -} - -} // extern "C" diff --git a/flashinfer/comm/trtllm_moe_alltoall.py b/flashinfer/comm/trtllm_moe_alltoall.py index 41afee5dc7a..4b417eaa8d3 100644 --- a/flashinfer/comm/trtllm_moe_alltoall.py +++ b/flashinfer/comm/trtllm_moe_alltoall.py @@ -4,12 +4,11 @@ This module provides the throughput-optimized all-to-all backend for MoE expert parallelism, supporting multiple payloads per collective operation. -The TensorRT-LLM implementation remains the default on every architecture. -Pass ``backend="cake"`` explicitly to use the generated Blackwell implementation. -All ranks in a collective must select the same backend, and use it consistently -for workspace sizing, initialization, dispatch, combine, and sanitization. -:class:`flashinfer.moe_ep.CakeAlltoAll` drives these ops with ``backend="cake"``; -:class:`flashinfer.moe_ep.NVLinkOneSidedAlltoAll` runs newer one-sided kernels. +The ops run the generated Cake kernels (``backend="cake"``, compute capability +10.0 or 10.3) and are driven by :class:`flashinfer.moe_ep.CakeAlltoAll`. The +TensorRT-LLM kernels (``backend="trtllm"``) were removed; use +:class:`flashinfer.moe_ep.NVLinkOneSidedAlltoAll`, which runs the current +TensorRT-LLM one-sided kernels. """ from types import SimpleNamespace @@ -25,11 +24,11 @@ from ..tllm_enums import SfLayout # Number of uint64 words in the active_rank_mask ABI (must match kRankMaskWords in -# csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h, which is +# csrc/cake_moe_alltoall_dispatch.cuh, which is # derived from kMaxRanks there). A single word covers up to 64 ranks. MOE_A2A_RANK_MASK_WORDS = 1 -MoeAlltoAllTarget = Literal["legacy", "sm100a", "sm103a"] -MoeAlltoAllBackend = Literal["trtllm", "cake"] +MoeAlltoAllTarget = Literal["sm100a", "sm103a"] +MoeAlltoAllBackend = Literal["cake"] def moe_a2a_active_rank_mask(active_ranks: Sequence[int], ep_size: int) -> torch.Tensor: @@ -71,17 +70,17 @@ def _moe_alltoall_target(device_index: int) -> MoeAlltoAllTarget: if capability == (10, 3): return "sm103a" raise ValueError( - f'backend="cake" requires compute capability 10.0 or 10.3, got {capability}' + "The MoE all-to-all ops require compute capability 10.0 or 10.3, got " + f"{capability}; use flashinfer.moe_ep.NVLinkOneSidedAlltoAll on other GPUs." ) @functools.cache def _get_moe_alltoall_module_for_target(target: MoeAlltoAllTarget): - """Build or load the legacy or exact-architecture all-to-all module.""" + """Build or load the all-to-all module of one exact architecture.""" module = gen_moe_alltoall_module(target).build_and_load() - # Keep legacy names stable, but do not let loading an opt-in module replace - # the default backend's custom-op implementations (or captured graphs). - op_suffix = "" if target == "legacy" else f"_{target}" + # Each architecture registers its own custom ops. + op_suffix = f"_{target}" @register_custom_op( f"flashinfer::moe_a2a_initialize{op_suffix}", @@ -362,14 +361,16 @@ def moe_a2a_get_aux_data_size( ) -def get_moe_alltoall_module(backend: MoeAlltoAllBackend = "trtllm"): - """Use TRT-LLM by default; exact-architecture kernels require explicit opt-in.""" +def get_moe_alltoall_module(backend: MoeAlltoAllBackend = "cake"): + """Return the Cake all-to-all module for the current device.""" if backend == "trtllm": - return _get_moe_alltoall_module_for_target("legacy") - if backend != "cake": raise ValueError( - f"Unknown MoE all-to-all backend {backend!r}; expected 'trtllm' or 'cake'" + 'The backend="trtllm" MoE all-to-all ops were removed; use ' + "flashinfer.moe_ep.NVLinkOneSidedAlltoAll, which runs the current " + "TensorRT-LLM one-sided kernels." ) + if backend != "cake": + raise ValueError(f"Unknown MoE all-to-all backend {backend!r}; expected 'cake'") device_index = torch.cuda.current_device() return _get_moe_alltoall_module_for_target(_moe_alltoall_target(device_index)) @@ -392,7 +393,7 @@ def moe_a2a_initialize( max_num_tokens: int, eplb_stats_num_experts: int = 0, *, - backend: MoeAlltoAllBackend = "trtllm", + backend: MoeAlltoAllBackend = "cake", ): r"""Initialize the MoE all-to-all workspace and return a metainfo tensor. @@ -418,10 +419,8 @@ def moe_a2a_initialize( region. ``0`` (default) disables the EPLB region; when non-zero it must match the ``eplb_local_stats`` length passed to :func:`moe_a2a_dispatch`. - backend : {"trtllm", "cake"} - Defaults to ``"trtllm"`` on all architectures. ``"cake"`` explicitly - opts in to generated Blackwell kernels (CC 10.0/10.3). Use the same - backend for all operations on this workspace and on all ranks. + backend : {"cake"} + The generated Blackwell kernels (CC 10.0/10.3), the only backend. Returns ------- @@ -503,7 +502,7 @@ def moe_a2a_dispatch( enable_rank_mask: bool = False, active_rank_mask: Optional[torch.Tensor] = None, *, - backend: MoeAlltoAllBackend = "trtllm", + backend: MoeAlltoAllBackend = "cake", recv_view_cache: Optional[dict] = None, ): r"""Dispatch tokens and payloads to their target expert ranks. @@ -551,8 +550,8 @@ def moe_a2a_dispatch( Masking a peer does not advance that peer's transport epoch. Before a skipped peer rejoins, all ranks must quiesce and coordinate workspace reinitialization; merely restoring its mask bit is insufficient. - backend : {"trtllm", "cake"} - Defaults to ``"trtllm"``. Pass ``"cake"`` to opt in on CC 10.0/10.3; + backend : {"cake"} + The generated Blackwell kernels (CC 10.0/10.3), the only backend; must match workspace initialization, combine, and all peer ranks. recv_view_cache : dict, optional Opaque cache of receive views. Pass an empty dictionary to enable @@ -675,17 +674,13 @@ def moe_a2a_combine( enable_pdl: Optional[bool] = None, enable_rank_mask: bool = False, active_rank_mask: Optional[torch.Tensor] = None, - backend: MoeAlltoAllBackend = "trtllm", + backend: MoeAlltoAllBackend = "cake", ) -> torch.Tensor: r"""Combine per-expert outputs back to the originating ranks. Inverse of :func:`moe_a2a_dispatch`: scatters the rank-local expert output rows back to the ranks that supplied the original tokens. - ``backend="trtllm"`` is the default, including on Blackwell. Opt in with - ``backend="cake"`` only for workspaces initialized and dispatched with - that backend; all peer ranks must use the same selection. - Parameters ---------- payload : torch.Tensor @@ -748,9 +743,8 @@ def moe_a2a_combine( :func:`moe_a2a_active_rank_mask`). Should match the mask passed to the corresponding :func:`moe_a2a_dispatch` call (or be omitted from both). Requires ``enable_rank_mask=True``. - backend : MoeAlltoAllBackend - Communication backend. Defaults to ``"trtllm"``; use ``"cake"`` only - with a workspace initialized by the Cake backend. + backend : {"cake"} + The generated Blackwell kernels (CC 10.0/10.3), the only backend. Returns ------- @@ -803,14 +797,11 @@ def moe_a2a_sanitize_expert_ids( invalid_expert_id: int, enable_pdl: Optional[bool] = None, *, - backend: MoeAlltoAllBackend = "trtllm", + backend: MoeAlltoAllBackend = "cake", ): r"""Sanitize invalid slots that contain no token routed to this rank by setting their expert IDs to ``invalid_expert_id``. - ``backend="trtllm"`` is the default. For a workspace initialized with - ``backend="cake"``, explicitly pass ``backend="cake"`` here as well. - Parameters ---------- expert_ids : torch.Tensor @@ -828,9 +819,8 @@ def moe_a2a_sanitize_expert_ids( enable_pdl : Optional[bool] Whether to use programmatic dependent launch. ``None`` auto-detects from the device. - backend : MoeAlltoAllBackend - Communication backend. Defaults to ``"trtllm"``; use ``"cake"`` only - with a workspace initialized by the Cake backend. + backend : {"cake"} + The generated Blackwell kernels (CC 10.0/10.3), the only backend. """ if enable_pdl is None: enable_pdl = device_support_pdl(expert_ids.device) @@ -847,13 +837,10 @@ def moe_a2a_get_workspace_size_per_rank( combine_payload_size_per_token: int, eplb_stats_num_experts: int = 0, *, - backend: MoeAlltoAllBackend = "trtllm", + backend: MoeAlltoAllBackend = "cake", ): r"""Compute the per-rank workspace size for the MoE all-to-all primitive. - ``backend="trtllm"`` is the default. Pass ``backend="cake"`` to size a - workspace for the explicitly selected Blackwell implementation. - Parameters ---------- ep_size : int @@ -869,9 +856,8 @@ def moe_a2a_get_workspace_size_per_rank( eplb_stats_num_experts : int Number of experts reserved for the EPLB gathered-stats region (``0`` disables it). - backend : MoeAlltoAllBackend - Communication backend used to size the workspace. Defaults to - ``"trtllm"``; pass ``"cake"`` for the Cake backend. + backend : {"cake"} + The generated Blackwell kernels (CC 10.0/10.3), the only backend. Returns ------- diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index c23a549a7ea..1a81757d19e 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -304,7 +304,7 @@ def gen_ulysses_a2a_module() -> JitSpec: ) -MoeAlltoAllTarget = Literal["legacy", "sm100a", "sm103a"] +MoeAlltoAllTarget = Literal["sm100a", "sm103a"] _MOE_ALLTOALL_GENERATED_SOURCE = { "sm100a": "mnnvl_moe_alltoall_sm100.cu", @@ -316,33 +316,23 @@ def gen_ulysses_a2a_module() -> JitSpec: } -def gen_moe_alltoall_module(target: MoeAlltoAllTarget = "legacy") -> JitSpec: - communication_kernels = ( - jit_env.FLASHINFER_CSRC_DIR - / "nv_internal" - / "tensorrt_llm" - / "kernels" - / "communicationKernels" - ) - if target == "legacy": - module_name = "mnnvl_moe_alltoall" - kernel_sources = [communication_kernels / "moeAlltoAllKernels.cu"] - arch_flags = [] - elif target in _MOE_ALLTOALL_GENERATED_SOURCE: - module_name = f"mnnvl_moe_alltoall_{target}" - kernel_sources = [ - communication_kernels / "moeAlltoAllFusedKernels.cu", - jit_env.FLASHINFER_CSRC_DIR - / "generated" - / _MOE_ALLTOALL_GENERATED_SOURCE[target], - ] - arch_flags = _MOE_ALLTOALL_NVCC_FLAGS[target] - else: +def gen_moe_alltoall_module(target: MoeAlltoAllTarget) -> JitSpec: + """Cake MoE all-to-all module behind the ``flashinfer.comm`` ``moe_a2a_*`` ops.""" + if target not in _MOE_ALLTOALL_GENERATED_SOURCE: raise ValueError(f"unsupported MNNVL MoE all-to-all target: {target}") + module_name = f"mnnvl_moe_alltoall_{target}" + kernel_sources = [ + jit_env.FLASHINFER_CSRC_DIR / "cake_moe_alltoall_dispatch.cu", + jit_env.FLASHINFER_CSRC_DIR + / "generated" + / "moe_alltoall" + / _MOE_ALLTOALL_GENERATED_SOURCE[target], + ] + arch_flags = _MOE_ALLTOALL_NVCC_FLAGS[target] return gen_jit_spec( module_name, [ - jit_env.FLASHINFER_CSRC_DIR / "trtllm_moe_alltoall.cu", + jit_env.FLASHINFER_CSRC_DIR / "cake_moe_alltoall.cu", *kernel_sources, jit_env.FLASHINFER_CSRC_DIR / "nv_internal" diff --git a/tests/comm/moe_a2a_driver.py b/tests/comm/moe_a2a_driver.py index ef516f4a1b2..f4c1afa115d 100644 --- a/tests/comm/moe_a2a_driver.py +++ b/tests/comm/moe_a2a_driver.py @@ -35,7 +35,7 @@ def workspace_bytes_per_rank( extra_payload_bytes_per_token: int = 0, eplb_stats_num_experts: int = 0, *, - backend: MoeAlltoAllBackend = "trtllm", + backend: MoeAlltoAllBackend = "cake", ) -> int: """Workspace bytes per rank for 16-bit hidden states, int32 expert ids and FP32 weights, plus ``extra_payload_bytes_per_token``.""" @@ -49,7 +49,7 @@ def workspace_bytes_per_rank( ) -def metainfo_index(backend: MoeAlltoAllBackend = "trtllm") -> dict: +def metainfo_index(backend: MoeAlltoAllBackend = "cake") -> dict: """Metainfo field indices, keyed by name without the ``MOE_A2A_`` prefix.""" names, values = get_moe_alltoall_module(backend).moe_a2a_get_metainfo_index_pairs() return { @@ -74,7 +74,7 @@ def __init__( eplb_stats_num_experts: int = 0, enable_rank_mask: bool = False, *, - backend: MoeAlltoAllBackend = "trtllm", + backend: MoeAlltoAllBackend = "cake", ): MnnvlMemory.initialize() if mnnvl_config: diff --git a/tests/comm/test_mnnvl_moe_alltoall_fused_adapter.py b/tests/comm/test_mnnvl_moe_alltoall_fused_adapter.py index b99f907fe78..c6d9631f275 100644 --- a/tests/comm/test_mnnvl_moe_alltoall_fused_adapter.py +++ b/tests/comm/test_mnnvl_moe_alltoall_fused_adapter.py @@ -100,14 +100,14 @@ def test_fused_module_keeps_the_public_python_contract(): } expected_module_defaults = { "moe_a2a_active_rank_mask": {}, - "moe_a2a_initialize": {"eplb_stats_num_experts": 0, "backend": "trtllm"}, + "moe_a2a_initialize": {"eplb_stats_num_experts": 0, "backend": "cake"}, "moe_a2a_wrap_payload_tensor_in_workspace": {}, "moe_a2a_dispatch": { "enable_pdl": None, "eplb_local_stats": None, "enable_rank_mask": False, "active_rank_mask": None, - "backend": "trtllm", + "backend": "cake", "recv_view_cache": None, }, "moe_a2a_combine": { @@ -121,12 +121,12 @@ def test_fused_module_keeps_the_public_python_contract(): "enable_pdl": None, "enable_rank_mask": False, "active_rank_mask": None, - "backend": "trtllm", + "backend": "cake", }, - "moe_a2a_sanitize_expert_ids": {"enable_pdl": None, "backend": "trtllm"}, + "moe_a2a_sanitize_expert_ids": {"enable_pdl": None, "backend": "cake"}, "moe_a2a_get_workspace_size_per_rank": { "eplb_stats_num_experts": 0, - "backend": "trtllm", + "backend": "cake", }, } assert set(api.__all__) == set(expected_module_parameters) @@ -157,13 +157,6 @@ def test_active_rank_mask_preserves_upper_u64_bits_in_dispatch(): assert mask.device.type == "cpu" assert mask.tolist() == [0x8000000180000001] - dispatch_source = ( - Path(__file__).resolve().parents[2] - / "csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllDispatch.cu" - ).read_text() - assert "1 << (unsigned long long)" not in dispatch_source - assert dispatch_source.count("1ULL << (unsigned long long)") == 5 - def test_backends_register_distinct_custom_ops(monkeypatch): import flashinfer.comm.trtllm_moe_alltoall as api @@ -195,11 +188,11 @@ def build_and_load(self): } try: expected_registrations = set() - for target in ("legacy", "sm100a", "sm103a"): + for target in ("sm100a", "sm103a"): module = api._get_moe_alltoall_module_for_target(target) assert api._get_moe_alltoall_module_for_target(target) is module assert set(vars(module)) == operation_names - suffix = "" if target == "legacy" else f"_{target}" + suffix = f"_{target}" expected_registrations.update( f"flashinfer::{name}{suffix}" for name in operation_names ) @@ -248,48 +241,15 @@ def capture(name, sources, **kwargs): names = {path.name for path in captured["sources"]} assert captured["name"] == module_name assert generated_name in names + assert {"cake_moe_alltoall.cu", "cake_moe_alltoall_dispatch.cu"} <= names assert { - "trtllm_moe_alltoall.cu", - "moeAlltoAllFusedKernels.cu", - } <= names - assert ( - not { - "moeAlltoAllKernels.cu", - "moeAlltoAllPrepareDispatch.cu", - "moeAlltoAllDispatch.cu", - "moeAlltoAllStageCombine.cu", - "moeAlltoAllPublishCombine.cu", - "moeAlltoAllCombine.cu", - "moeAlltoAllQuantizeCombine.cu", - "moeAlltoAllSanitize.cu", - } - & names - ) - assert { - path.name for path in captured["sources"] if path.parent.name == "generated" + path.name + for path in captured["sources"] + if path.parent.name == "moe_alltoall" and path.parent.parent.name == "generated" } == {generated_name} assert any(arch_flag in flag for flag in captured["kwargs"]["extra_cuda_cflags"]) -def test_legacy_jit_inventory_remains_available(monkeypatch): - import flashinfer.jit.comm as jit_comm - - captured = {} - - def capture(name, sources, **kwargs): - captured.update(name=name, sources=tuple(sources), kwargs=kwargs) - return object() - - monkeypatch.setattr(jit_comm, "gen_jit_spec", capture) - jit_comm.gen_moe_alltoall_module() - names = {path.name for path in captured["sources"]} - assert captured["name"] == "mnnvl_moe_alltoall" - assert "moeAlltoAllKernels.cu" in names - assert "moeAlltoAllFusedKernels.cu" not in names - assert not {name for name in names if name.startswith("mnnvl_moe_alltoall_sm")} - assert captured["kwargs"]["extra_cuda_cflags"] == ["-DENABLE_BF16"] - - @pytest.mark.parametrize( ("capability", "cake_target"), [ @@ -302,7 +262,7 @@ def capture(name, sources, **kwargs): ((12, 1), None), ], ) -def test_runtime_module_selection_requires_explicit_cake_opt_in( +def test_runtime_module_selection_follows_the_device( monkeypatch, capability, cake_target ): import flashinfer.comm.trtllm_moe_alltoall as api @@ -325,22 +285,23 @@ def select_target(target): try: api.get_moe_alltoall_module.cache_clear() - assert api.get_moe_alltoall_module() == "legacy" - assert api.get_moe_alltoall_module(backend="trtllm") == "legacy" + with pytest.raises(ValueError, match="NVLinkOneSidedAlltoAll"): + api.get_moe_alltoall_module(backend="trtllm") assert capability_queries == [] if cake_target is None: - with pytest.raises(ValueError, match="cake"): - api.get_moe_alltoall_module(backend="cake") - assert selected == ["legacy", "legacy"] + with pytest.raises(ValueError, match="compute capability"): + api.get_moe_alltoall_module() + assert selected == [] else: + assert api.get_moe_alltoall_module() == cake_target assert api.get_moe_alltoall_module(backend="cake") == cake_target - assert selected == ["legacy", "legacy", cake_target] + assert selected == [cake_target, cake_target] assert capability_queries == [7] finally: api.get_moe_alltoall_module.cache_clear() -@pytest.mark.parametrize("backend", ["auto", "legacy", "sm100a", "", None]) +@pytest.mark.parametrize("backend", ["auto", "trtllm", "legacy", "sm100a", "", None]) def test_runtime_module_selection_rejects_invalid_backends(backend): import flashinfer.comm.trtllm_moe_alltoall as api @@ -393,11 +354,8 @@ def spec(name): def test_quantized_combine_preserves_the_physical_accumulator_dtype(): repo_root = Path(__file__).resolve().parents[2] - adapter_source = (repo_root / "csrc/trtllm_moe_alltoall.cu").read_text() - launcher_source = ( - repo_root - / "csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllFusedKernels.cu" - ).read_text() + adapter_source = (repo_root / "csrc/cake_moe_alltoall.cu").read_text() + launcher_source = (repo_root / "csrc/cake_moe_alltoall_dispatch.cu").read_text() assert ( "alloc_tensor({localNumTokens, elementsPerToken}, payload.dtype(), payload.device())" @@ -409,26 +367,20 @@ def test_quantized_combine_preserves_the_physical_accumulator_dtype(): def test_combine_inventory_is_specialized_and_fail_closed(): - source_root = ( - Path(__file__).resolve().parents[2] - / "csrc/nv_internal/tensorrt_llm/kernels/communicationKernels" - ) - combine_source = (source_root / "moeAlltoAllCombine.cu").read_text() - launcher_source = (source_root / "moeAlltoAllFusedKernels.cu").read_text() + launcher_source = ( + Path(__file__).resolve().parents[2] / "csrc/cake_moe_alltoall_dispatch.cu" + ).read_text() supported_top_k = (1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 22) for top_k in supported_top_k: - assert f"DEFINE_COMBINE_TOP_K({top_k})" in combine_source assert f"PRELOAD_COMBINE_TOP_K({top_k})" in launcher_source assert f"LAUNCH_COMBINE_TOP_K({top_k})" in launcher_source assert "unsupported top_k for moe_a2a_combine" in launcher_source - assert "int ep_rank, int top_k" not in combine_source def test_bf16_vector_routes_and_stage_grid_keep_exact_boundaries(): launcher_source = ( - Path(__file__).resolve().parents[2] - / "csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllFusedKernels.cu" + Path(__file__).resolve().parents[2] / "csrc/cake_moe_alltoall_dispatch.cu" ).read_text() for top_k in (1, 2, 4): @@ -494,7 +446,7 @@ def test_bf16_vector_routes_and_stage_grid_keep_exact_boundaries(): " params.enable_rank_mask, params.active_rank_mask[0]" in launcher_source ) - generated_root = Path(__file__).resolve().parents[2] / "csrc/generated" + generated_root = Path(__file__).resolve().parents[2] / "csrc/generated/moe_alltoall" for generated_name in ( "mnnvl_moe_alltoall_sm100.cu", "mnnvl_moe_alltoall_sm103.cu", @@ -605,14 +557,12 @@ def _normalized_cxx(source: str) -> str: def test_fused_quantized_combine_routes_mirror_the_cake_host_rule(): repo_root = Path(__file__).resolve().parents[2] - kernel_root = ( - repo_root / "csrc/nv_internal/tensorrt_llm/kernels/communicationKernels" - ) - launcher_source = (kernel_root / "moeAlltoAllFusedKernels.cu").read_text() + launcher_source = (repo_root / "csrc/cake_moe_alltoall_dispatch.cu").read_text() launcher = _normalized_cxx(launcher_source) - header = _normalized_cxx((kernel_root / "moeAlltoAllKernels.h").read_text()) - legacy = _normalized_cxx((kernel_root / "moeAlltoAllKernels.cu").read_text()) - adapter = _normalized_cxx((repo_root / "csrc/trtllm_moe_alltoall.cu").read_text()) + header = _normalized_cxx( + (repo_root / "csrc/cake_moe_alltoall_dispatch.cuh").read_text() + ) + adapter = _normalized_cxx((repo_root / "csrc/cake_moe_alltoall.cu").read_text()) # one extern declaration, preload slot and launch case per exported TOP_K for top_k in _FUSED_QUANTIZED_COMBINE_TOP_K: @@ -678,8 +628,7 @@ def test_fused_quantized_combine_routes_mirror_the_cake_host_rule(): assert "int const grid = std::max(params.local_num_tokens, 1);" in launcher_source # both quantized routes take the scale geometry from one helper assert launcher_source.count("quantizedCombineLayout(params);") == 2 - # the host boundary skips the BF16/FP16 intermediate on the fused route; the trtllm - # backend never fuses, so backend="trtllm" keeps its allocation + # the host boundary skips the BF16/FP16 intermediate on the fused route assert ( "bool moe_a2a_combine_fuses_quantization(MoeA2ACombineParams const& params);" in header @@ -688,10 +637,6 @@ def test_fused_quantized_combine_routes_mirror_the_cake_host_rule(): "bool moe_a2a_combine_fuses_quantization(MoeA2ACombineParams const& params) { " "return useFusedQuantizedCombine(params); }" in launcher ) - assert ( - "bool moe_a2a_combine_fuses_quantization(MoeA2ACombineParams const&) " - "{ return false; }" in legacy - ) assert ( "if (!tl_throughput::moe_a2a_combine_fuses_quantization(params)) { " "accumulationStorage = alloc_tensor({localNumTokens, elementsPerToken}, " @@ -699,7 +644,7 @@ def test_fused_quantized_combine_routes_mirror_the_cake_host_rule(): "params.accumulation_data = accumulationStorage.data_ptr(); }" in adapter ) - generated_root = repo_root / "csrc/generated" + generated_root = repo_root / "csrc/generated/moe_alltoall" for generated_name in ( "mnnvl_moe_alltoall_sm100.cu", "mnnvl_moe_alltoall_sm103.cu", @@ -743,8 +688,7 @@ def test_fused_quantized_combine_routes_mirror_the_cake_host_rule(): def test_kernel_preload_cache_is_per_kernel_and_device_without_global_tree(): launcher_source = ( - Path(__file__).resolve().parents[2] - / "csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllFusedKernels.cu" + Path(__file__).resolve().parents[2] / "csrc/cake_moe_alltoall_dispatch.cu" ).read_text() assert "template " in launcher_source @@ -1065,7 +1009,7 @@ def _run_public_combine_round( ) -def _run_public_mpi2_cycle(backend): +def _run_public_mpi2_cycle(): comm = MPI.COMM_WORLD rank = comm.Get_rank() if comm.Get_size() != 2: @@ -1085,11 +1029,7 @@ def _run_public_mpi2_cycle(backend): from flashinfer.comm.trtllm_moe_alltoall import get_moe_alltoall_module - # Keep both implementations loaded before exercising either collective. - legacy_module = get_moe_alltoall_module() - cake_module = get_moe_alltoall_module(backend="cake") - assert get_moe_alltoall_module(backend="trtllm") is legacy_module - assert cake_module is not legacy_module + assert get_moe_alltoall_module() is get_moe_alltoall_module(backend="cake") routes = torch.tensor(_ROUTES_BY_RANK[rank], dtype=torch.int32, device="cuda") payloads = _payloads(rank, routes) @@ -1103,7 +1043,7 @@ def _run_public_mpi2_cycle(backend): _HIDDEN_SIZE, extra_payload_bytes_per_token=extra_payload_bytes, eplb_stats_num_experts=5, - backend=backend, + backend="cake", ) mapping = Mapping(rank=rank, moe_ep_size=2, tp_size=2, world_size=2) collective = MoeA2ADriver( @@ -1114,7 +1054,7 @@ def _run_public_mpi2_cycle(backend): workspace_size_per_rank=workspace_size, eplb_stats_num_experts=5, enable_rank_mask=True, - backend=backend, + backend="cake", ) active_mask = moe_a2a_active_rank_mask((0, 1), 2) _run_public_combine_round( @@ -1149,7 +1089,7 @@ def _run_public_mpi2_cycle(backend): max_tokens, _HIDDEN_SIZE, extra_payload_bytes_per_token=extra_payload_bytes, - backend=backend, + backend="cake", ) topk8_collective = MoeA2ADriver( mapping, @@ -1158,7 +1098,7 @@ def _run_public_mpi2_cycle(backend): num_experts=16, workspace_size_per_rank=topk8_workspace_size, enable_rank_mask=True, - backend=backend, + backend="cake", ) _run_public_combine_round( topk8_collective, @@ -1207,15 +1147,14 @@ def _run_public_mpi2_cycle(backend): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_public_mpi2_nondivisible_six_payload_eplb_and_external_output(): - """Exercise both backends in one process through public APIs only.""" - legacy = _run_public_mpi2_cycle(backend="trtllm") - cake = _run_public_mpi2_cycle(backend="cake") - for legacy_collective, cake_collective in zip(legacy, cake, strict=True): + """Run the public-API cycle twice; the second reuses the first's workspaces.""" + first = _run_public_mpi2_cycle() + second = _run_public_mpi2_cycle() + for first_collective, second_collective in zip(first, second, strict=True): assert ( - legacy_collective.workspace.data_ptr() - != cake_collective.workspace.data_ptr() + first_collective.workspace.data_ptr() + == second_collective.workspace.data_ptr() ) - _run_public_mpi2_cycle(backend="trtllm") # Check a dropped peer only after all full-rank reuse checks. A peer that # misses collective epochs cannot rejoin without coordinated reinitialization. @@ -1226,7 +1165,7 @@ def test_public_mpi2_nondivisible_six_payload_eplb_and_external_output(): routes = torch.tensor(_ROUTES_BY_RANK[rank], dtype=torch.int32, device="cuda") payloads = _payloads(rank, routes) rank_zero_mask = moe_a2a_active_rank_mask((0,), 2) - for collective in (legacy[0], cake[0]): + for collective in (first[0],): _run_public_combine_round( collective, rank, From 9b1bddfdd6c71a88859d8ac974772a3e2afb16dd Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Sun, 18 Oct 2026 16:22:26 +0000 Subject: [PATCH 04/17] refactor(moe_ep)!: move the Cake and two-sided all-to-all ops into their backends The Cake and NVLink two-sided backends drove their kernels through function layers in flashinfer.comm (trtllm_moe_alltoall.py and trtllm_alltoall.py), left over from before flashinfer.moe_ep. Each backend package now owns its kernels the way nvlink_one_sided does, completing the move announced when MoeAlltoAll / MnnvlMoe were deprecated (#5532). Cake: - cake/kernels.py loads the per-architecture module and holds the kernel-level initialize / dispatch / combine / sanitize_expert_ids / workspace_size_per_rank / wrap_workspace / active_rank_mask operations; CakeAlltoAll calls them directly. - flashinfer/comm/trtllm_moe_alltoall.py and the flashinfer.comm moe_a2a_* re-exports are removed, along with the backend argument and the per-architecture torch custom ops. - JIT: gen_moe_alltoall_module -> gen_moe_ep_cake_module, modules moe_ep_cake_alltoall_sm100a / _sm103a. NVLink two-sided: - nvlink_two_sided/kernels.py holds prepare / alltoallv and the workspace sizes; NVLinkTwoSidedAlltoAll calls them directly. - flashinfer/comm/trtllm_alltoall.py is removed, with the moe_comm_prepare_indices and moe_local_gather ops and kernels nothing but their tests used. - Sources: csrc/trtllm_alltoall{,_prepare}.cu and include/flashinfer/comm/trtllm_alltoall{,_prepare}.cuh become csrc/moe_ep_nvlink_two_sided_{alltoall,prepare}.{cu,cuh} in namespace flashinfer::moe_ep_nvlink_two_sided. JIT: gen_comm_alltoall_module -> gen_moe_ep_nvlink_two_sided_module, module moe_ep_nvlink_two_sided_alltoall. The workspace sharing of both backends is unchanged. Kernel-level tests import the new modules; test_trtllm_moe_alltoall.py and test_trtllm_alltoall.py are renamed test_cake_moe_alltoall.py and test_nvlink_two_sided_alltoall.py. BREAKING CHANGE: flashinfer.comm.moe_a2a_* and the flashinfer.comm.trtllm_moe_alltoall / trtllm_alltoall modules are removed; use flashinfer.moe_ep.CakeAlltoAll / NVLinkTwoSidedAlltoAll. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- ...cu => moe_ep_nvlink_two_sided_alltoall.cu} | 168 +--- .../moe_ep_nvlink_two_sided_alltoall.cuh | 404 +------- ....cu => moe_ep_nvlink_two_sided_prepare.cu} | 8 +- .../moe_ep_nvlink_two_sided_prepare.cuh | 4 +- flashinfer/aot.py | 10 +- flashinfer/comm/__init__.py | 15 - flashinfer/comm/trtllm_alltoall.py | 416 -------- flashinfer/comm/trtllm_moe_alltoall.py | 892 ------------------ flashinfer/jit/__init__.py | 6 +- flashinfer/jit/comm.py | 29 +- .../backends/split/comm/cake/communication.py | 67 +- .../backends/split/comm/cake/kernels.py | 373 ++++++++ .../comm/nvlink_two_sided/communication.py | 20 +- .../split/comm/nvlink_two_sided/kernels.py | 175 ++++ scripts/task_test_single_node_comm_kernels.sh | 2 +- tests/comm/moe_a2a_driver.py | 55 +- ..._alltoall.py => test_cake_moe_alltoall.py} | 92 +- tests/comm/test_mnnvl_memory.py | 10 +- .../test_mnnvl_moe_alltoall_fused_adapter.py | 260 +---- tests/comm/test_moe_a2a_recv_view_cache.py | 16 +- ...l.py => test_nvlink_two_sided_alltoall.py} | 236 +---- tests/moe_ep/test_communication.py | 34 +- 22 files changed, 778 insertions(+), 2514 deletions(-) rename csrc/{trtllm_alltoall.cu => moe_ep_nvlink_two_sided_alltoall.cu} (54%) rename include/flashinfer/comm/trtllm_alltoall.cuh => csrc/moe_ep_nvlink_two_sided_alltoall.cuh (51%) rename csrc/{trtllm_alltoall_prepare.cu => moe_ep_nvlink_two_sided_prepare.cu} (99%) rename include/flashinfer/comm/trtllm_alltoall_prepare.cuh => csrc/moe_ep_nvlink_two_sided_prepare.cuh (98%) delete mode 100644 flashinfer/comm/trtllm_alltoall.py delete mode 100644 flashinfer/comm/trtllm_moe_alltoall.py create mode 100644 flashinfer/moe_ep/backends/split/comm/cake/kernels.py create mode 100644 flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/kernels.py rename tests/comm/{test_trtllm_moe_alltoall.py => test_cake_moe_alltoall.py} (95%) rename tests/comm/{test_trtllm_alltoall.py => test_nvlink_two_sided_alltoall.py} (75%) diff --git a/csrc/trtllm_alltoall.cu b/csrc/moe_ep_nvlink_two_sided_alltoall.cu similarity index 54% rename from csrc/trtllm_alltoall.cu rename to csrc/moe_ep_nvlink_two_sided_alltoall.cu index 005fe10f1e0..7e8e136e854 100644 --- a/csrc/trtllm_alltoall.cu +++ b/csrc/moe_ep_nvlink_two_sided_alltoall.cu @@ -17,147 +17,15 @@ #include -#include "flashinfer/comm/trtllm_alltoall.cuh" -#include "flashinfer/comm/trtllm_alltoall_prepare.cuh" +#include "moe_ep_nvlink_two_sided_alltoall.cuh" +#include "moe_ep_nvlink_two_sided_prepare.cuh" #include "tvm_ffi_utils.h" -using namespace flashinfer::trtllm_alltoall; +using namespace flashinfer::moe_ep_nvlink_two_sided; using tvm::ffi::Optional; using tvm::ffi::Tuple; -void moeCommPrepareIndicesOp(TensorView gatheredTargetRankIds, - Optional realRankTokenCountCumSum, - TensorView localGatherIndices, TensorView sendRankCountCumSum, - TensorView sendRankLocalIndices, TensorView recvRankCountCumSum, - TensorView recvRankLocalIndices, - TensorView backwardRecvRankLocalIndices, int64_t maxTokenCountPerRank, - int64_t expertCount, int64_t topK, int64_t epRank, int64_t epSize) { - CHECK_INPUT_TYPE(gatheredTargetRankIds, dl_int32); - TVM_FFI_ICHECK_EQ(gatheredTargetRankIds.ndim(), 2) << "gatheredTargetRankIds must be a 2D tensor"; - TVM_FFI_ICHECK_EQ(gatheredTargetRankIds.size(1), topK) - << "gatheredTargetRankIds must have topK columns"; - - int const* realRankTokenCountCumSumPtr = nullptr; - if (realRankTokenCountCumSum.has_value()) { - TVM_FFI_ICHECK_EQ(realRankTokenCountCumSum.value().ndim(), 1) - << "realRankTokenCountCumSum must be a 1D tensor"; - CHECK_INPUT_TYPE(realRankTokenCountCumSum.value(), dl_int32) - TVM_FFI_ICHECK_EQ(realRankTokenCountCumSum.value().size(0), epSize) - << "realRankTokenCountCumSum must have epSize elements"; - realRankTokenCountCumSumPtr = static_cast(realRankTokenCountCumSum.value().data_ptr()); - } else { - TVM_FFI_ICHECK_EQ(gatheredTargetRankIds.size(0), epSize * maxTokenCountPerRank) - << "gatheredTargetRankIds should have shape (epSize * maxTokenCountPerRank, topK)"; - } - TVM_FFI_ICHECK_GT(maxTokenCountPerRank, 0) << "maxTokenCountPerRank must be greater than 0"; - TVM_FFI_ICHECK_GT(expertCount, 0) << "expertCount must be greater than 0"; - TVM_FFI_ICHECK_GT(topK, 0) << "topK must be greater than 0"; - TVM_FFI_ICHECK_LE(topK, expertCount) << "topK must be less than or equal to expertCount"; - TVM_FFI_ICHECK(epRank >= 0 && epRank < epSize) << "epRank must be in the range [0, epSize)"; - - auto stream = get_current_stream(); - - int maxSendRanksPerToken = std::max(static_cast(epSize), static_cast(topK)); - - CHECK_INPUT_TYPE(localGatherIndices, dl_int32); - TVM_FFI_ICHECK_EQ(localGatherIndices.ndim(), 1); - TVM_FFI_ICHECK_EQ(localGatherIndices.size(0), maxTokenCountPerRank * epSize); - - CHECK_INPUT_TYPE(sendRankCountCumSum, dl_int32); - TVM_FFI_ICHECK_EQ(sendRankCountCumSum.ndim(), 1); - TVM_FFI_ICHECK_EQ(sendRankCountCumSum.size(0), epSize); - - CHECK_INPUT_TYPE(sendRankLocalIndices, dl_int32); - TVM_FFI_ICHECK_EQ(sendRankLocalIndices.ndim(), 1); - TVM_FFI_ICHECK_EQ(sendRankLocalIndices.size(0), maxTokenCountPerRank * maxSendRanksPerToken); - - CHECK_INPUT_TYPE(recvRankCountCumSum, dl_int32); - TVM_FFI_ICHECK_EQ(recvRankCountCumSum.ndim(), 1); - TVM_FFI_ICHECK_EQ(recvRankCountCumSum.size(0), epSize); - - CHECK_INPUT_TYPE(recvRankLocalIndices, dl_int32); - TVM_FFI_ICHECK_EQ(recvRankLocalIndices.ndim(), 1); - TVM_FFI_ICHECK_EQ(recvRankLocalIndices.size(0), maxTokenCountPerRank * epSize); - - CHECK_INPUT_TYPE(backwardRecvRankLocalIndices, dl_int32); - TVM_FFI_ICHECK_EQ(backwardRecvRankLocalIndices.ndim(), 1); - TVM_FFI_ICHECK_EQ(backwardRecvRankLocalIndices.size(0), - maxTokenCountPerRank * maxSendRanksPerToken); - - flashinfer::trtllm_alltoall::MoeExpertParallelInfo expertParallelInfo; - expertParallelInfo.expertCount = static_cast(expertCount); - expertParallelInfo.topK = static_cast(topK); - - flashinfer::trtllm_alltoall::MoeEpWorldInfo worldInfo = {static_cast(epSize), - static_cast(epRank)}; - - auto cudaResult = flashinfer::trtllm_alltoall::moeAllToAllPrepareIndices( - worldInfo, expertParallelInfo, static_cast(maxTokenCountPerRank), - static_cast(gatheredTargetRankIds.data_ptr()), realRankTokenCountCumSumPtr, - static_cast(localGatherIndices.data_ptr()), - static_cast(sendRankCountCumSum.data_ptr()), - static_cast(sendRankLocalIndices.data_ptr()), - static_cast(recvRankCountCumSum.data_ptr()), - static_cast(recvRankLocalIndices.data_ptr()), - static_cast(backwardRecvRankLocalIndices.data_ptr()), stream); - TVM_FFI_ICHECK(cudaResult == cudaSuccess) - << "CUDA error in moeAllToAllPrepareIndices: " << cudaGetErrorString(cudaResult); -} - -void moeLocalGatherOp(TensorView recvRankCumSum, TensorView localGatherIndices, - TensorView gatheredExpertIds, TensorView gatheredScales, - TensorView localExpertIds, TensorView localScales, - int64_t maxTokenCountPerRank, int64_t expertCount, int64_t topK, - int64_t epRank, int64_t epSize) { - CHECK_INPUT_TYPE(recvRankCumSum, dl_int32); - CHECK_INPUT_TYPE(localGatherIndices, dl_int32); - CHECK_INPUT_TYPE(gatheredExpertIds, dl_int32); - CHECK_INPUT_TYPE(gatheredScales, dl_float32); - CHECK_INPUT_TYPE(localExpertIds, dl_int32); - CHECK_INPUT_TYPE(localScales, dl_float32); - - TVM_FFI_ICHECK_GT(maxTokenCountPerRank, 0) << "maxTokenCountPerRank must be greater than 0"; - TVM_FFI_ICHECK_GT(expertCount, 0) << "expertCount must be greater than 0"; - TVM_FFI_ICHECK_GT(topK, 0) << "topK must be greater than 0"; - TVM_FFI_ICHECK_LE(topK, expertCount) << "topK must be less than or equal to expertCount"; - TVM_FFI_ICHECK(epRank >= 0 && epRank < epSize) << "epRank must be in the range [0, epSize)"; - - TVM_FFI_ICHECK_EQ(recvRankCumSum.ndim(), 1) << "recvRankCumSum must be a 1D tensor"; - TVM_FFI_ICHECK_EQ(recvRankCumSum.size(0), epSize) << "recvRankCumSum must have epSize elements"; - TVM_FFI_ICHECK_EQ(localGatherIndices.ndim(), 1) << "localGatherIndices must be a 1D tensor"; - TVM_FFI_ICHECK_EQ(gatheredExpertIds.ndim(), 2) << "gatheredExpertIds must be a 2D tensor"; - TVM_FFI_ICHECK_EQ(gatheredScales.ndim(), 2) << "gatheredScales must be a 2D tensor"; - TVM_FFI_ICHECK_EQ(localExpertIds.ndim(), 2) << "localExpertIds must be a 2D tensor"; - TVM_FFI_ICHECK_EQ(localScales.ndim(), 2) << "localScales must be a 2D tensor"; - TVM_FFI_ICHECK_EQ(gatheredExpertIds.size(1), topK) << "gatheredExpertIds must have topK columns"; - TVM_FFI_ICHECK_EQ(gatheredScales.size(1), topK) << "gatheredScales must have topK columns"; - TVM_FFI_ICHECK_EQ(localExpertIds.size(1), topK) << "localExpertIds must have topK columns"; - TVM_FFI_ICHECK_EQ(localScales.size(1), topK) << "localScales must have topK columns"; - - int localMaxTokenCount = static_cast(localGatherIndices.size(0)); - TVM_FFI_ICHECK_EQ(localExpertIds.size(0), localMaxTokenCount) - << "localExpertIds must have localMaxTokenCount rows"; - TVM_FFI_ICHECK_EQ(localScales.size(0), localMaxTokenCount) - << "localScales must have localMaxTokenCount rows"; - - auto stream = get_current_stream(); - - flashinfer::trtllm_alltoall::MoeExpertParallelInfo expertParallelInfo; - expertParallelInfo.expertCount = static_cast(expertCount); - expertParallelInfo.topK = static_cast(topK); - - flashinfer::trtllm_alltoall::MoeEpWorldInfo worldInfo = {static_cast(epSize), - static_cast(epRank)}; - flashinfer::trtllm_alltoall::moeLocalGather( - worldInfo, expertParallelInfo, static_cast(maxTokenCountPerRank), localMaxTokenCount, - static_cast(recvRankCumSum.data_ptr()), - static_cast(localGatherIndices.data_ptr()), - static_cast(gatheredExpertIds.data_ptr()), - static_cast(gatheredScales.data_ptr()), static_cast(localExpertIds.data_ptr()), - static_cast(localScales.data_ptr()), stream); -} - void moeCommOp(TensorView input, TensorView sendRankCumSum, TensorView sendIndices, TensorView output, TensorView recvRankCumSum, TensorView recvIndices, TensorView allWorkspaces, int64_t epRank, int64_t epSize) { @@ -184,9 +52,9 @@ void moeCommOp(TensorView input, TensorView sendRankCumSum, TensorView sendIndic TVM_FFI_ICHECK(epRank >= 0 && epRank < epSize) << "epRank must be in the range [0, epSize)"; - flashinfer::trtllm_alltoall::MoeEpWorldInfo worldInfo = {static_cast(epSize), - static_cast(epRank)}; - flashinfer::trtllm_alltoall::SendRecvDataInfo sendRecvDataInfo; + flashinfer::moe_ep_nvlink_two_sided::MoeEpWorldInfo worldInfo = {static_cast(epSize), + static_cast(epRank)}; + flashinfer::moe_ep_nvlink_two_sided::SendRecvDataInfo sendRecvDataInfo; size_t eltSize = get_element_size(input); size_t eltCountPerU64 = sizeof(uint64_t) / eltSize; @@ -195,7 +63,7 @@ void moeCommOp(TensorView input, TensorView sendRankCumSum, TensorView sendIndic sendRecvDataInfo.vectorSizeInU64 = input.size(1) / eltCountPerU64; sendRecvDataInfo.DoPreCompute(); - flashinfer::trtllm_alltoall::SendRecvDispls sendDispls, recvDispls; + flashinfer::moe_ep_nvlink_two_sided::SendRecvDispls sendDispls, recvDispls; sendDispls.dataPtr = static_cast(input.data_ptr()); sendDispls.rankCountCumSum = static_cast(sendRankCumSum.data_ptr()); sendDispls.rankLocalIndices = static_cast(sendIndices.data_ptr()); @@ -206,24 +74,24 @@ void moeCommOp(TensorView input, TensorView sendRankCumSum, TensorView sendIndic recvDispls.rankLocalIndices = static_cast(recvIndices.data_ptr()); recvDispls.vectorStrideInU64 = output.stride(0) / eltCountPerU64; - flashinfer::trtllm_alltoall::MoeCommWorkspace workspace; + flashinfer::moe_ep_nvlink_two_sided::MoeCommWorkspace workspace; workspace.workspacePtr = static_cast(allWorkspaces.data_ptr()); workspace.rankStrideInU64 = allWorkspaces.stride(0); auto stream = get_current_stream(); - flashinfer::trtllm_alltoall::moeAllToAll(worldInfo, sendRecvDataInfo, sendDispls, recvDispls, - workspace, stream); + flashinfer::moe_ep_nvlink_two_sided::moeAllToAll(worldInfo, sendRecvDataInfo, sendDispls, + recvDispls, workspace, stream); } int64_t getWorkspaceSizePerRank(int64_t epSize) { int epSize32 = static_cast(epSize); - return flashinfer::trtllm_alltoall::getMoeCommWorkspaceSize(epSize32); + return flashinfer::moe_ep_nvlink_two_sided::getMoeCommWorkspaceSize(epSize32); } int64_t getPrepareWorkspaceSizePerRank(int64_t epSize) { int epSize32 = static_cast(epSize); - return flashinfer::trtllm_alltoall::moe_prepare::getMoePrepareWorkspaceSize(epSize32); + return flashinfer::moe_ep_nvlink_two_sided::moe_prepare::getMoePrepareWorkspaceSize(epSize32); } void moePrepareOp(TensorView expertsIds, Optional scales, @@ -315,13 +183,13 @@ void moePrepareOp(TensorView expertsIds, Optional scales, gatheredExpertStaticsPtr = static_cast(gatheredExpertStatics.value().data_ptr()); } - flashinfer::trtllm_alltoall::moe_prepare::MoeCommWorkspace workspace; + flashinfer::moe_ep_nvlink_two_sided::moe_prepare::MoeCommWorkspace workspace; workspace.workspacePtr = static_cast(allWorkspaces.data_ptr()); workspace.rankStrideInU64 = allWorkspaces.stride(0); auto stream = get_current_stream(); - flashinfer::trtllm_alltoall::moe_prepare::computeCountAndIndice( + flashinfer::moe_ep_nvlink_two_sided::moe_prepare::computeCountAndIndice( static_cast(expertsIds.data_ptr()), static_cast(sendRankCountCumSum.data_ptr()), static_cast(recvRankCountCumSum.data_ptr()), static_cast(sendRankIndices.data_ptr()), @@ -329,11 +197,11 @@ void moePrepareOp(TensorView expertsIds, Optional scales, static_cast(recvRankIndices.data_ptr()), workspace, tokenCount, maxTokenCountPerRank, topK, slotCount, epRank, epSize, stream); - flashinfer::trtllm_alltoall::moe_prepare::computeCumsum( + flashinfer::moe_ep_nvlink_two_sided::moe_prepare::computeCumsum( static_cast(sendRankCountCumSum.data_ptr()), static_cast(recvRankCountCumSum.data_ptr()), epRank, epSize, stream); - flashinfer::trtllm_alltoall::moe_prepare::moveIndice( + flashinfer::moe_ep_nvlink_two_sided::moe_prepare::moveIndice( static_cast(sendRankCountCumSum.data_ptr()), static_cast(recvRankCountCumSum.data_ptr()), static_cast(sendRankIndices.data_ptr()), @@ -344,7 +212,7 @@ void moePrepareOp(TensorView expertsIds, Optional scales, static_cast(gatherRecvRankIndices.data_ptr()), epRank, epSize, maxTokenCountPerRank, stream); - flashinfer::trtllm_alltoall::moe_prepare::allToAllMetadata( + flashinfer::moe_ep_nvlink_two_sided::moe_prepare::allToAllMetadata( static_cast(expertsIds.data_ptr()), static_cast(preparedLocalExpertIds.data_ptr()), scalesPtr, preparedLocalScalesPtr, localExpertStaticsPtr, gatheredExpertStaticsPtr, workspace, @@ -355,8 +223,6 @@ void moePrepareOp(TensorView expertsIds, Optional scales, expertCount, slotCount, epRank, epSize, stream); } -TVM_FFI_DLL_EXPORT_TYPED_FUNC(moe_comm_prepare_indices, moeCommPrepareIndicesOp); -TVM_FFI_DLL_EXPORT_TYPED_FUNC(moe_local_gather, moeLocalGatherOp); TVM_FFI_DLL_EXPORT_TYPED_FUNC(moe_comm, moeCommOp); TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_moe_max_usable_sm_count, setMaxUsableSmCount); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_moe_commworkspace_size_per_rank, getWorkspaceSizePerRank); diff --git a/include/flashinfer/comm/trtllm_alltoall.cuh b/csrc/moe_ep_nvlink_two_sided_alltoall.cuh similarity index 51% rename from include/flashinfer/comm/trtllm_alltoall.cuh rename to csrc/moe_ep_nvlink_two_sided_alltoall.cuh index 10c400528c3..412ac4fece2 100644 --- a/include/flashinfer/comm/trtllm_alltoall.cuh +++ b/csrc/moe_ep_nvlink_two_sided_alltoall.cuh @@ -14,21 +14,19 @@ * limitations under the License. */ -#include #include -#include +#include +#include -#include "../exception.h" -#include "../logging.h" -#include "../utils.cuh" -#include "../vec_dtypes.cuh" - -namespace cg = cooperative_groups; +#include "flashinfer/exception.h" +#include "flashinfer/logging.h" +#include "flashinfer/utils.cuh" +#include "flashinfer/vec_dtypes.cuh" namespace flashinfer { -namespace trtllm_alltoall { +namespace moe_ep_nvlink_two_sided { #ifdef __CUDACC__ #define ALIGN_256 __align__(256) @@ -157,11 +155,6 @@ struct MoeEpWorldInfo { int epRank; }; -struct MoeExpertParallelInfo { - int expertCount = -1; - int topK = 1; -}; - struct SendRecvDataInfo { int vectorSizeInU64; // pre-computed at host side for GPU kernel @@ -533,387 +526,6 @@ void moeAllToAll(MoeEpWorldInfo worldInfo, SendRecvDataInfo sendRecvDataInfo, recvDispls); } -template -__inline__ __device__ void computeSendRecvRankCountDevice( - MoeEpWorldInfo worldInfo, MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, - int const* realRankTokenCountCumSum, int const* gatheredTargetRankIds, - int* sharedSendRecvRankCount, int* sendRecvRankCount) { - cg::thread_block_tile tile = - cg::tiled_partition(cg::this_thread_block()); - int laneInTile = tile.thread_rank(); - int tileId = threadIdx.x / kThreadsGroupSize; - int tileCountPerBlock = blockDim.x / kThreadsGroupSize; - - int topK = expertParallelInfo.topK; - int epRank = worldInfo.epRank; - int epSize = worldInfo.epSize; - - if (threadIdx.x == 0) { - *sharedSendRecvRankCount = 0; - } - - __syncthreads(); - int readRank = isSend ? epRank : blockIdx.x; - int compareRankId = isSend ? blockIdx.x : epRank; - int const* readRankTargetRankIds = gatheredTargetRankIds + readRank * maxTokenCountPerRank * topK; - int readRankTokenCount = maxTokenCountPerRank; - if (realRankTokenCountCumSum != nullptr) { - int readRankStart = readRank == 0 ? 0 : realRankTokenCountCumSum[readRank - 1]; - readRankTargetRankIds = gatheredTargetRankIds + readRankStart * topK; - readRankTokenCount = realRankTokenCountCumSum[readRank] - readRankStart; - } - - for (int i = tileId + blockIdx.z * tileCountPerBlock; i < readRankTokenCount; - i += tileCountPerBlock * gridDim.z) { - int targetRankId = laneInTile < topK ? readRankTargetRankIds[i * topK + laneInTile] : epSize; - bool rankMatched = (targetRankId == compareRankId); - bool hasRankMatched = tile.any(rankMatched); - if (hasRankMatched && laneInTile == 0) { - atomicAdd_block(sharedSendRecvRankCount, 1); - } - tile.sync(); - } - __syncthreads(); - if (threadIdx.x == 0) { - atomicAdd_system(sendRecvRankCount + blockIdx.x, *sharedSendRecvRankCount); - } -} - -template -__global__ void computeSendRecvRankCountKernel(MoeEpWorldInfo worldInfo, - MoeExpertParallelInfo expertParallelInfo, - int maxTokenCountPerRank, - int const* realRankTokenCountCumSum, - int const* gatheredTargetRankIds, int* sendRankCount, - int* recvRankCount) { - static_assert(kThreadsGroupSize == 1 || kThreadsGroupSize == 2 || kThreadsGroupSize == 4 || - kThreadsGroupSize == 8 || kThreadsGroupSize == 16 || kThreadsGroupSize == 32, - "Only 1, 2, 4, 8, 16, 32 threads group size supported now."); - __shared__ int sharedSendRecvRankCount; - if (blockIdx.y == 0) { - // compute send rank count - computeSendRecvRankCountDevice( - worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, - gatheredTargetRankIds, &sharedSendRecvRankCount, sendRankCount); - } else { - // compute recv rank count - computeSendRecvRankCountDevice( - worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, - gatheredTargetRankIds, &sharedSendRecvRankCount, recvRankCount); - } -} - -void computeSendRecvRankCount(const MoeEpWorldInfo& worldInfo, - MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, - int const* realRankTokenCountCumSum, int const* gatheredTargetRankIds, - int* sendRankCount, int* recvRankCount, cudaStream_t stream) { - FLASHINFER_CHECK(expertParallelInfo.topK <= 32, - "Only topK less than or equal to 32 supported now."); - int threadsPerBlock = 1024; - auto* kernelPtr = computeSendRecvRankCountKernel<32>; - if (expertParallelInfo.topK <= 1) { - kernelPtr = computeSendRecvRankCountKernel<1>; - } else if (expertParallelInfo.topK <= 2) { - kernelPtr = computeSendRecvRankCountKernel<2>; - } else if (expertParallelInfo.topK <= 4) { - kernelPtr = computeSendRecvRankCountKernel<4>; - } else if (expertParallelInfo.topK <= 8) { - kernelPtr = computeSendRecvRankCountKernel<8>; - } else if (expertParallelInfo.topK <= 16) { - kernelPtr = computeSendRecvRankCountKernel<16>; - } - dim3 block(worldInfo.epSize, 2, 1); - kernelPtr<<>>( - worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, - gatheredTargetRankIds, sendRankCount, recvRankCount); -} - -template -__global__ void inplaceSendRecvRankCumSumKernel(MoeEpWorldInfo worldInfo, int* sendRankCount, - int* recvRankCount) { - int* inputOutputPtr = blockIdx.x == 0 ? sendRankCount : recvRankCount; - typedef cub::BlockScan BlockScan; - __shared__ typename BlockScan::TempStorage temp_storage; - - int tid = threadIdx.x; - int threadData = tid < worldInfo.epSize ? inputOutputPtr[tid] : 0; - - BlockScan(temp_storage).InclusiveSum(threadData, threadData); - if (tid < worldInfo.epSize) { - inputOutputPtr[tid] = threadData; - } -} - -void inplaceSendRecvRankCumSum(const MoeEpWorldInfo& worldInfo, int* sendRankCount, - int* recvRankCount, cudaStream_t stream) { - int epSize = worldInfo.epSize; - int epRank = worldInfo.epRank; - - FLASHINFER_CHECK(epSize <= 1024, "Only epSize less than or equal to 1024 supported now."); - - auto* kernelPtr = inplaceSendRecvRankCumSumKernel<1024>; - int blockSize = 1024; - if (epSize <= 32) { - kernelPtr = inplaceSendRecvRankCumSumKernel<32>; - blockSize = 32; - } else if (epSize <= 64) { - kernelPtr = inplaceSendRecvRankCumSumKernel<64>; - blockSize = 64; - } else if (epSize <= 128) { - kernelPtr = inplaceSendRecvRankCumSumKernel<128>; - blockSize = 128; - } else if (epSize <= 256) { - kernelPtr = inplaceSendRecvRankCumSumKernel<256>; - blockSize = 256; - } else if (epSize <= 512) { - kernelPtr = inplaceSendRecvRankCumSumKernel<512>; - blockSize = 512; - } - kernelPtr<<<2, blockSize, 0, stream>>>(worldInfo, sendRankCount, recvRankCount); -} - -template -__inline__ __device__ void computeSendRecvIndicesDevice( - MoeEpWorldInfo worldInfo, MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, - int const* realRankTokenCountCumSum, int const* gatheredTargetRankIds, - int const* sendRecvCumSum, - int* sendRecvIndices, // send or receive - int* localGatherIndices, // receive only - int* backwardRecvRankLocalIndices, // send only - int* sharedSendRecvRankStart, - typename cub::BlockScan::TempStorage& tempStorage) { - cg::thread_block_tile tile = - cg::tiled_partition(cg::this_thread_block()); - int laneInTile = tile.thread_rank(); - int tileId = threadIdx.x / kThreadsGroupSize; - int tileCountPerBlock = blockDim.x / kThreadsGroupSize; - - int topK = expertParallelInfo.topK; - int epRank = worldInfo.epRank; - int epSize = worldInfo.epSize; - - if (threadIdx.x == 0) { - *sharedSendRecvRankStart = blockIdx.x == 0 ? 0 : sendRecvCumSum[blockIdx.x - 1]; - } - - __syncthreads(); - int readRank = isSend ? epRank : blockIdx.x; - int compareRankId = isSend ? blockIdx.x : epRank; - int readRankStart = readRank * maxTokenCountPerRank; - int const* readRankTargetRankIds = gatheredTargetRankIds + readRankStart * topK; - int readRankTokenCount = maxTokenCountPerRank; - if (realRankTokenCountCumSum != nullptr) { - readRankStart = readRank == 0 ? 0 : realRankTokenCountCumSum[readRank - 1]; - readRankTargetRankIds = gatheredTargetRankIds + readRankStart * topK; - readRankTokenCount = realRankTokenCountCumSum[readRank] - readRankStart; - } - - for (int blockStartId = blockIdx.z * tileCountPerBlock; blockStartId < readRankTokenCount; - blockStartId += tileCountPerBlock * gridDim.z) { - int stepStartIndice = *sharedSendRecvRankStart; - int i = blockStartId + tileId; - int targetRankId = (laneInTile < topK && i < readRankTokenCount) - ? readRankTargetRankIds[i * topK + laneInTile] - : epSize; - bool rankMatched = (targetRankId == compareRankId); - bool hasRankMatched = tile.any(rankMatched); - unsigned int laneMask = tile.ballot(rankMatched); - int lowestLane = __ffs(laneMask) - 1; - int isMatchedLane = (hasRankMatched && laneInTile == lowestLane) ? 1 : 0; - int indice; - typedef cub::BlockScan BlockScan; - BlockScan(tempStorage).ExclusiveSum(isMatchedLane, indice); - indice += stepStartIndice; - __syncthreads(); - - if (isMatchedLane == 1) { - atomicAdd_block(sharedSendRecvRankStart, 1); - if (isSend) { - sendRecvIndices[indice] = i; - backwardRecvRankLocalIndices[indice] = i * topK + lowestLane; - } else { - sendRecvIndices[indice] = indice; - localGatherIndices[indice] = readRankStart + i; - } - } - __syncthreads(); - } -} - -template -__global__ void computeSendRecvIndicesKernel( - MoeEpWorldInfo worldInfo, MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, - int const* realRankTokenCountCumSum, int const* gatheredTargetRankIds, - int const* sendRankCountCumSum, int const* recvRankCountCumSum, int* localGatherIndices, - int* sendRankLocalIndices, int* recvRankLocalIndices, int* backwardRecvRankLocalIndices) { - static_assert(kThreadsGroupSize == 1 || kThreadsGroupSize == 2 || kThreadsGroupSize == 4 || - kThreadsGroupSize == 8 || kThreadsGroupSize == 16 || kThreadsGroupSize == 32, - "Only 1, 2, 4, 8, 16, 32 threads group size supported now."); - __shared__ int sharedSendRecvRankStart; - __shared__ typename cub::BlockScan::TempStorage tempStorage; - if (blockIdx.y == 0) { - // compute send rank count - computeSendRecvIndicesDevice( - worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, - gatheredTargetRankIds, sendRankCountCumSum, sendRankLocalIndices, localGatherIndices, - backwardRecvRankLocalIndices, &sharedSendRecvRankStart, tempStorage); - } else { - // compute recv rank count - computeSendRecvIndicesDevice( - worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, - gatheredTargetRankIds, recvRankCountCumSum, recvRankLocalIndices, localGatherIndices, - backwardRecvRankLocalIndices, &sharedSendRecvRankStart, tempStorage); - } -} - -void computeSendRecvIndices(const MoeEpWorldInfo& worldInfo, - MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, - int const* realRankTokenCountCumSum, int const* gatheredTargetRankIds, - int const* sendRankCountCumSum, int const* recvRankCountCumSum, - int* localGatherIndices, int* sendRankLocalIndices, - int* recvRankLocalIndices, int* backwardRecvRankLocalIndices, - cudaStream_t stream) { - FLASHINFER_CHECK(expertParallelInfo.topK <= 32, - "Only topK less than or equal to 32 supported now."); - int threadsPerBlock = 1024; - auto* kernelPtr = computeSendRecvIndicesKernel<32, 1024>; - if (expertParallelInfo.topK <= 1) { - kernelPtr = computeSendRecvIndicesKernel<1, 1024>; - } else if (expertParallelInfo.topK <= 2) { - kernelPtr = computeSendRecvIndicesKernel<2, 1024>; - } else if (expertParallelInfo.topK <= 4) { - kernelPtr = computeSendRecvIndicesKernel<4, 1024>; - } else if (expertParallelInfo.topK <= 8) { - kernelPtr = computeSendRecvIndicesKernel<8, 1024>; - } else if (expertParallelInfo.topK <= 16) { - kernelPtr = computeSendRecvIndicesKernel<16, 1024>; - } else if (expertParallelInfo.topK <= 32) { - kernelPtr = computeSendRecvIndicesKernel<32, 1024>; - } - dim3 block(worldInfo.epSize, 2, 1); - kernelPtr<<>>( - worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, - gatheredTargetRankIds, sendRankCountCumSum, recvRankCountCumSum, localGatherIndices, - sendRankLocalIndices, recvRankLocalIndices, backwardRecvRankLocalIndices); -} - -cudaError_t moeAllToAllPrepareIndices( - MoeEpWorldInfo worldInfo, MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, - int const* gatheredTargetRankIds, int const* realRankTokenCountCumSum, - // indices of gatheredTargetRankIds that has the local rank in topK - int* localGatherIndices, // max length = maxTokenCountPerRank * worldInfo.epSize when all ranks - // send to current rank - int* sendRankCountCumSum, // max length = worldInfo.epSize - int* sendRankLocalIndices, // max length = maxTokenCountPerRank * - // expertParallelInfo.expertCount when current rank has - // maxTokenCountPerRank tokens to send and all has expertCount dest - int* recvRankCountCumSum, // max length = worldInfo.epSize - int* recvRankLocalIndices, // max length = maxTokenCountPerRank * worldInfo.epSize when all - // ranks send to current rank - // the rankCountCumSum of combineRecv should be the same as sendRankCountCumSum - int* backwardRecvRankLocalIndices, // max length = maxTokenCountPerRank * - // expertParallelInfo.expertCount when current rank has - // maxTokenCountPerRank tokens to send and all has - // expertCount dest - cudaStream_t stream) { - FLASHINFER_CHECK(worldInfo.epSize <= 1024, - "Only worldInfo.epSize less than or equal to 1024 supported now."); - - FLASHINFER_CUDA_CALL( - cudaMemsetAsync(sendRankCountCumSum, 0, sizeof(int) * worldInfo.epSize, stream)); - FLASHINFER_CUDA_CALL( - cudaMemsetAsync(recvRankCountCumSum, 0, sizeof(int) * worldInfo.epSize, stream)); - int maxSendRanksPerToken = std::max(worldInfo.epSize, expertParallelInfo.topK); - - FLASHINFER_CUDA_CALL(cudaMemsetAsync( - localGatherIndices, -1, maxTokenCountPerRank * worldInfo.epSize * sizeof(int), stream)); - FLASHINFER_CUDA_CALL(cudaMemsetAsync( - sendRankLocalIndices, -1, maxTokenCountPerRank * maxSendRanksPerToken * sizeof(int), stream)); - FLASHINFER_CUDA_CALL(cudaMemsetAsync( - recvRankLocalIndices, -1, maxTokenCountPerRank * worldInfo.epSize * sizeof(int), stream)); - FLASHINFER_CUDA_CALL(cudaMemsetAsync(backwardRecvRankLocalIndices, -1, - maxTokenCountPerRank * maxSendRanksPerToken * sizeof(int), - stream)); - computeSendRecvRankCount(worldInfo, expertParallelInfo, maxTokenCountPerRank, - realRankTokenCountCumSum, gatheredTargetRankIds, sendRankCountCumSum, - recvRankCountCumSum, stream); - - inplaceSendRecvRankCumSum(worldInfo, sendRankCountCumSum, recvRankCountCumSum, stream); - computeSendRecvIndices(worldInfo, expertParallelInfo, maxTokenCountPerRank, - realRankTokenCountCumSum, gatheredTargetRankIds, sendRankCountCumSum, - recvRankCountCumSum, localGatherIndices, sendRankLocalIndices, - recvRankLocalIndices, backwardRecvRankLocalIndices, stream); - return cudaSuccess; -} - -template -__global__ void moeLocalGatherDevice(MoeEpWorldInfo worldInfo, - MoeExpertParallelInfo expertParallelInfo, - int maxTokenCountPerRank, int localMaxTokenCount, - int const* recvRankCountCumSum, int const* localGatherIndices, - int const* gatheredExpertIds, float const* gatheredScales, - int* localExpertIds, float* localScales) { - cg::thread_block_tile tile = - cg::tiled_partition(cg::this_thread_block()); - int laneInTile = tile.thread_rank(); - int tileId = threadIdx.x / kThreadsGroupSize; - int tileCountPerBlock = blockDim.x / kThreadsGroupSize; - - int epSize = worldInfo.epSize; - int rankTokenCount = recvRankCountCumSum[epSize - 1]; - bool needLoad = laneInTile < expertParallelInfo.topK; - - for (int index = tileId + blockIdx.x * tileCountPerBlock; index < localMaxTokenCount; - index += tileCountPerBlock * gridDim.x) { - int localTokenIndice = localGatherIndices[index]; - int expertId = needLoad && (index < rankTokenCount) - ? gatheredExpertIds[localTokenIndice * expertParallelInfo.topK + laneInTile] - : expertParallelInfo.expertCount; - float scale = needLoad && (index < rankTokenCount) - ? gatheredScales[localTokenIndice * expertParallelInfo.topK + laneInTile] - : 0.0f; - if (needLoad) { - localExpertIds[index * expertParallelInfo.topK + laneInTile] = expertId; - localScales[index * expertParallelInfo.topK + laneInTile] = scale; - } - } -} - -void moeLocalGather(MoeEpWorldInfo worldInfo, MoeExpertParallelInfo expertParallelInfo, - int maxTokenCountPerRank, int localMaxTokenCount, - int const* recvRankCountCumSum, int const* localGatherIndices, - int const* gatheredExpertIds, float const* gatheredScales, int* localExpertIds, - float* localScales, cudaStream_t stream) { - FLASHINFER_CHECK(expertParallelInfo.topK <= 32, - "Only topK less than or equal to 32 supported now."); - auto* kernelPtr = moeLocalGatherDevice<32>; - int paddedTopK = 32; - if (expertParallelInfo.topK <= 1) { - paddedTopK = 1; - kernelPtr = moeLocalGatherDevice<1>; - } else if (expertParallelInfo.topK <= 2) { - paddedTopK = 2; - kernelPtr = moeLocalGatherDevice<2>; - } else if (expertParallelInfo.topK <= 4) { - paddedTopK = 4; - kernelPtr = moeLocalGatherDevice<4>; - } else if (expertParallelInfo.topK <= 8) { - paddedTopK = 8; - kernelPtr = moeLocalGatherDevice<8>; - } else if (expertParallelInfo.topK <= 16) { - paddedTopK = 16; - kernelPtr = moeLocalGatherDevice<16>; - } - - int threadsPerBlock = 512; - int tokenPerBlock = threadsPerBlock / paddedTopK; - int blockCount = (localMaxTokenCount + tokenPerBlock - 1) / tokenPerBlock * 2; - - kernelPtr<<>>( - worldInfo, expertParallelInfo, maxTokenCountPerRank, localMaxTokenCount, recvRankCountCumSum, - localGatherIndices, gatheredExpertIds, gatheredScales, localExpertIds, localScales); -} - int AllToAllChannelCommunicatorBase::maxSmCount = -1; bool AllToAllChannelCommunicatorBase::maxSmCountUsed = false; @@ -921,5 +533,5 @@ void setMaxUsableSmCount(int smCount) { AllToAllChannelCommunicatorBase::setMaxUsableSmCount(smCount); } -} // namespace trtllm_alltoall +} // namespace moe_ep_nvlink_two_sided } // namespace flashinfer diff --git a/csrc/trtllm_alltoall_prepare.cu b/csrc/moe_ep_nvlink_two_sided_prepare.cu similarity index 99% rename from csrc/trtllm_alltoall_prepare.cu rename to csrc/moe_ep_nvlink_two_sided_prepare.cu index c6e8f1f6efd..ccfc98a411b 100644 --- a/csrc/trtllm_alltoall_prepare.cu +++ b/csrc/moe_ep_nvlink_two_sided_prepare.cu @@ -20,11 +20,11 @@ #include -#include "flashinfer/comm/trtllm_alltoall_prepare.cuh" #include "flashinfer/exception.h" #include "flashinfer/utils.cuh" +#include "moe_ep_nvlink_two_sided_prepare.cuh" -// Local definition to avoid multiple definition issues from trtllm_alltoall.cuh +// Local definition to avoid multiple definition issues from moe_ep_nvlink_two_sided_alltoall.cuh static int getMultiProcessorCount() { int device_id; int multi_processor_count; @@ -36,7 +36,7 @@ static int getMultiProcessorCount() { namespace cg = cooperative_groups; -namespace flashinfer::trtllm_alltoall { +namespace flashinfer::moe_ep_nvlink_two_sided { namespace moe_prepare { @@ -620,4 +620,4 @@ size_t getMoePrepareWorkspaceSize(int epSize) { } // namespace moe_prepare -} // namespace flashinfer::trtllm_alltoall +} // namespace flashinfer::moe_ep_nvlink_two_sided diff --git a/include/flashinfer/comm/trtllm_alltoall_prepare.cuh b/csrc/moe_ep_nvlink_two_sided_prepare.cuh similarity index 98% rename from include/flashinfer/comm/trtllm_alltoall_prepare.cuh rename to csrc/moe_ep_nvlink_two_sided_prepare.cuh index 2945eee9732..ee0f9613198 100644 --- a/include/flashinfer/comm/trtllm_alltoall_prepare.cuh +++ b/csrc/moe_ep_nvlink_two_sided_prepare.cuh @@ -21,7 +21,7 @@ #define DEBUG_PIPELINE 0 -namespace flashinfer::trtllm_alltoall { +namespace flashinfer::moe_ep_nvlink_two_sided { namespace moe_prepare { @@ -125,4 +125,4 @@ size_t getMoePrepareWorkspaceSize(int epSize); } // namespace moe_prepare -} // namespace flashinfer::trtllm_alltoall +} // namespace flashinfer::moe_ep_nvlink_two_sided diff --git a/flashinfer/aot.py b/flashinfer/aot.py index 58ec4e11f15..9f6f82e3b1b 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -993,11 +993,11 @@ def gen_all_modules( if add_comm: from .jit.comm import ( - gen_comm_alltoall_module, gen_dcp_alltoall_module, gen_dcp_lse_reduce_module, - gen_moe_alltoall_module, + gen_moe_ep_cake_module, gen_moe_ep_nvlink_one_sided_module, + gen_moe_ep_nvlink_two_sided_module, gen_pcie_ipc_ag_rs_module, gen_pcie_ipc_comm_module, gen_trtllm_comm_module, @@ -1005,7 +1005,7 @@ def gen_all_modules( gen_vllm_comm_module, ) - jit_specs.append(gen_comm_alltoall_module()) + jit_specs.append(gen_moe_ep_nvlink_two_sided_module()) if ( has_sm90 or has_sm100 @@ -1036,9 +1036,9 @@ def gen_all_modules( ): jit_specs.append(gen_dcp_lse_reduce_module()) if has_sm100a_exact: - jit_specs.append(gen_moe_alltoall_module("sm100a")) + jit_specs.append(gen_moe_ep_cake_module("sm100a")) if has_sm103a_exact: - jit_specs.append(gen_moe_alltoall_module("sm103a")) + jit_specs.append(gen_moe_ep_cake_module("sm103a")) # NVLink one-sided MoE all-to-all behind flashinfer.moe_ep. The fence path # needs SM90+; the CFT path compiles only for SM100+ and traps elsewhere. if has_sm90 or has_sm100 or has_sm100f or has_sm103 or has_sm107: diff --git a/flashinfer/comm/__init__.py b/flashinfer/comm/__init__.py index e178f3cb521..0348b25ce55 100644 --- a/flashinfer/comm/__init__.py +++ b/flashinfer/comm/__init__.py @@ -126,21 +126,6 @@ tp_moe_finalize_rmsnorm as tp_moe_finalize_rmsnorm, ) -# MNNVL A2A (Throughput Backend) -from .trtllm_moe_alltoall import moe_a2a_active_rank_mask as moe_a2a_active_rank_mask -from .trtllm_moe_alltoall import moe_a2a_combine as moe_a2a_combine -from .trtllm_moe_alltoall import moe_a2a_dispatch as moe_a2a_dispatch -from .trtllm_moe_alltoall import moe_a2a_initialize as moe_a2a_initialize -from .trtllm_moe_alltoall import ( - moe_a2a_get_workspace_size_per_rank as moe_a2a_get_workspace_size_per_rank, -) -from .trtllm_moe_alltoall import ( - moe_a2a_sanitize_expert_ids as moe_a2a_sanitize_expert_ids, -) -from .trtllm_moe_alltoall import ( - moe_a2a_wrap_payload_tensor_in_workspace as moe_a2a_wrap_payload_tensor_in_workspace, -) - # DCP A2A (Decode Context Parallel Attention Reduction) from .dcp_alltoall import decode_cp_a2a_alltoall as decode_cp_a2a_alltoall from .dcp_alltoall import ( diff --git a/flashinfer/comm/trtllm_alltoall.py b/flashinfer/comm/trtllm_alltoall.py deleted file mode 100644 index a8a8d03c2d7..00000000000 --- a/flashinfer/comm/trtllm_alltoall.py +++ /dev/null @@ -1,416 +0,0 @@ -""" -Copyright (c) 2025 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 ..jit import gen_comm_alltoall_module -from ..utils import register_custom_op - - -@functools.cache -def get_comm_alltoall_module(): - module = gen_comm_alltoall_module().build_and_load() - - @register_custom_op( - "flashinfer::moe_comm_prepare_indices", - mutates_args=[], - ) - def moe_comm_prepare_indices( - gathered_target_rank_ids: torch.Tensor, - real_rank_token_count_cum_sum: Optional[torch.Tensor], - max_token_count_per_rank: int, - expert_count: int, - top_k: int, - ep_rank: int, - ep_size: int, - ) -> Tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - ]: - device = gathered_target_rank_ids.device - max_send_ranks_per_token = max(top_k, ep_size) - local_gather_indices = torch.empty( - (max_token_count_per_rank * ep_size), device=device, dtype=torch.int - ) - send_rank_count_cum_sum = torch.empty( - (ep_size,), device=device, dtype=torch.int - ) - send_rank_local_indices = torch.empty( - (max_token_count_per_rank * max_send_ranks_per_token), - device=device, - dtype=torch.int, - ) - recv_rank_count_cum_sum = torch.empty((ep_size), device=device, dtype=torch.int) - recv_rank_local_indices = torch.empty( - (max_token_count_per_rank * ep_size), device=device, dtype=torch.int - ) - backward_recv_rank_local_indice = torch.empty( - (max_token_count_per_rank * max_send_ranks_per_token), - device=device, - dtype=torch.int, - ) - module.moe_comm_prepare_indices( - gathered_target_rank_ids, - real_rank_token_count_cum_sum, - local_gather_indices, - send_rank_count_cum_sum, - send_rank_local_indices, - recv_rank_count_cum_sum, - recv_rank_local_indices, - backward_recv_rank_local_indice, - max_token_count_per_rank, - expert_count, - top_k, - ep_rank, - ep_size, - ) - return ( - local_gather_indices, - send_rank_count_cum_sum, - send_rank_local_indices, - recv_rank_count_cum_sum, - recv_rank_local_indices, - backward_recv_rank_local_indice, - ) - - @register_custom_op( - "flashinfer::moe_local_gather", - mutates_args=["local_expert_ids", "local_scales"], - ) - def moe_local_gather( - recv_rank_cum_sum: torch.Tensor, - local_gather_indices: torch.Tensor, - gathered_expert_ids: torch.Tensor, - gathered_scales: torch.Tensor, - local_expert_ids: torch.Tensor, - local_scales: torch.Tensor, - max_token_count_per_rank: int, - expert_count: int, - top_k: int, - ep_rank: int, - ep_size: int, - ) -> None: - module.moe_local_gather( - recv_rank_cum_sum, - local_gather_indices, - gathered_expert_ids, - gathered_scales, - local_expert_ids, - local_scales, - max_token_count_per_rank, - expert_count, - top_k, - ep_rank, - ep_size, - ) - - @register_custom_op( - "flashinfer::moe_comm", - mutates_args=["output"], - ) - def moe_comm( - input: torch.Tensor, - send_rank_cum_sum: torch.Tensor, - send_indices: torch.Tensor, - output: torch.Tensor, - recv_rank_cum_sum: torch.Tensor, - recv_indices: torch.Tensor, - all_workspaces: torch.Tensor, - ep_rank: int, - ep_size: int, - ) -> None: - module.moe_comm( - input, - send_rank_cum_sum, - send_indices, - output, - recv_rank_cum_sum, - recv_indices, - all_workspaces, - ep_rank, - ep_size, - ) - - @register_custom_op( - "flashinfer::set_moe_max_usable_sm_count", - mutates_args=[], - ) - def set_moe_max_usable_sm_count( - max_sm_count: int, - ) -> None: - module.set_moe_max_usable_sm_count(max_sm_count) - - @register_custom_op( - "flashinfer::get_moe_commworkspace_size_per_rank", - mutates_args=[], - ) - def get_moe_commworkspace_size_per_rank( - ep_size: int, - ) -> int: - return module.get_moe_commworkspace_size_per_rank(ep_size) - - @register_custom_op( - "flashinfer::get_moe_prepare_workspace_size_per_rank", - mutates_args=[], - ) - def get_moe_prepare_workspace_size_per_rank( - ep_size: int, - ) -> int: - return module.get_moe_prepare_workspace_size_per_rank(ep_size) - - @register_custom_op( - "flashinfer::moe_prepare", - mutates_args=[], - ) - def moe_prepare( - experts_ids: torch.Tensor, - scales: Optional[torch.Tensor], - experts_statics: Optional[torch.Tensor], - workspace: torch.Tensor, - max_token_count_per_rank: int, - ep_rank: int, - ep_size: int, - expert_count: int, - slot_count: int, - top_k: int, - ) -> Tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - ]: - attrs = {"dtype": torch.int32, "device": experts_ids.device} - - prepared_local_expert_ids = torch.empty( - (max_token_count_per_rank * ep_size, top_k), **attrs - ) - send_rank_count_cum_sum = torch.empty((ep_size,), **attrs) - recv_rank_count_cum_sum = torch.empty((ep_size,), **attrs) - - gather_recv_rank_indices = torch.empty( - (max_token_count_per_rank * ep_size,), **attrs - ) - recv_rank_indices = torch.empty((max_token_count_per_rank * ep_size,), **attrs) - - max_send_ranks_per_token = max(ep_size, top_k) - - gather_backward_recv_rank_indices = torch.empty( - (max_token_count_per_rank * max_send_ranks_per_token,), **attrs - ) - backward_recv_rank_indices = torch.empty( - (max_token_count_per_rank * max_send_ranks_per_token,), **attrs - ) - gather_send_rank_indices = torch.empty( - (max_token_count_per_rank * max_send_ranks_per_token,), **attrs - ) - send_rank_indices = torch.empty( - (max_token_count_per_rank * max_send_ranks_per_token,), **attrs - ) - if scales is not None: - prepared_local_scales = torch.empty( - (max_token_count_per_rank * ep_size, top_k), - dtype=torch.float32, - device=attrs["device"], - ) - else: - prepared_local_scales = None - if experts_statics is not None: - gathered_expert_statics = torch.empty((ep_size, expert_count), **attrs) - else: - gathered_expert_statics = None - - module.moe_prepare( - experts_ids, - scales, - experts_statics, - workspace, - prepared_local_expert_ids, - send_rank_count_cum_sum, - recv_rank_count_cum_sum, - gather_recv_rank_indices, - recv_rank_indices, - gather_backward_recv_rank_indices, - backward_recv_rank_indices, - gather_send_rank_indices, - send_rank_indices, - prepared_local_scales, - gathered_expert_statics, - max_token_count_per_rank, - ep_rank, - ep_size, - expert_count, - slot_count, - top_k, - ) - return ( - prepared_local_expert_ids, - prepared_local_scales, - send_rank_count_cum_sum, - gather_send_rank_indices, - recv_rank_count_cum_sum, - gather_recv_rank_indices, - gather_backward_recv_rank_indices, - gathered_expert_statics, - ) - - return SimpleNamespace( - moe_comm_prepare_indices=moe_comm_prepare_indices, - moe_local_gather=moe_local_gather, - moe_comm=moe_comm, - set_moe_max_usable_sm_count=set_moe_max_usable_sm_count, - get_moe_commworkspace_size_per_rank=get_moe_commworkspace_size_per_rank, - get_moe_prepare_workspace_size_per_rank=get_moe_prepare_workspace_size_per_rank, - moe_prepare=moe_prepare, - ) - - -def moe_comm_prepare_indices( - gathered_target_rank_ids: torch.Tensor, - real_rank_token_count_cum_sum: Optional[torch.Tensor], - max_token_count_per_rank: int, - expert_count: int, - top_k: int, - ep_rank: int, - ep_size: int, -) -> Tuple[ - torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor -]: - return get_comm_alltoall_module().moe_comm_prepare_indices( - gathered_target_rank_ids, - real_rank_token_count_cum_sum, - max_token_count_per_rank, - expert_count, - top_k, - ep_rank, - ep_size, - ) - - -def moe_local_gather( - recv_rank_cum_sum: torch.Tensor, - local_gather_indices: torch.Tensor, - gathered_expert_ids: torch.Tensor, - gathered_scales: torch.Tensor, - local_expert_ids: torch.Tensor, - local_scales: torch.Tensor, - max_token_count_per_rank: int, - expert_count: int, - top_k: int, - ep_rank: int, - ep_size: int, -) -> None: - get_comm_alltoall_module().moe_local_gather( - recv_rank_cum_sum, - local_gather_indices, - gathered_expert_ids, - gathered_scales, - local_expert_ids, - local_scales, - max_token_count_per_rank, - expert_count, - top_k, - ep_rank, - ep_size, - ) - - -def moe_comm( - input: torch.Tensor, - send_rank_cum_sum: torch.Tensor, - send_indices: torch.Tensor, - output: torch.Tensor, - recv_rank_cum_sum: torch.Tensor, - recv_indices: torch.Tensor, - all_workspaces: torch.Tensor, - ep_rank: int, - ep_size: int, -) -> None: - get_comm_alltoall_module().moe_comm( - input, - send_rank_cum_sum, - send_indices, - output, - recv_rank_cum_sum, - recv_indices, - all_workspaces, - ep_rank, - ep_size, - ) - - -def set_moe_max_usable_sm_count( - max_sm_count: int, -) -> None: - get_comm_alltoall_module().set_moe_max_usable_sm_count(max_sm_count) - - -def get_moe_commworkspace_size_per_rank( - ep_size: int, -) -> int: - return get_comm_alltoall_module().get_moe_commworkspace_size_per_rank(ep_size) - - -def get_moe_prepare_workspace_size_per_rank( - ep_size: int, -) -> int: - return get_comm_alltoall_module().get_moe_prepare_workspace_size_per_rank(ep_size) - - -def moe_prepare( - experts_ids: torch.Tensor, - scales: Optional[torch.Tensor], - experts_statics: Optional[torch.Tensor], - workspace: torch.Tensor, - max_token_count_per_rank: int, - ep_rank: int, - ep_size: int, - expert_count: int, - slot_count: int, - top_k: int, -) -> Tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, -]: - return get_comm_alltoall_module().moe_prepare( - experts_ids, - scales, - experts_statics, - workspace, - max_token_count_per_rank, - ep_rank, - ep_size, - expert_count, - slot_count, - top_k, - ) diff --git a/flashinfer/comm/trtllm_moe_alltoall.py b/flashinfer/comm/trtllm_moe_alltoall.py deleted file mode 100644 index 4b417eaa8d3..00000000000 --- a/flashinfer/comm/trtllm_moe_alltoall.py +++ /dev/null @@ -1,892 +0,0 @@ -""" -MoE All-to-All Operations (Throughput Backend) - -This module provides the throughput-optimized all-to-all backend for MoE expert parallelism, -supporting multiple payloads per collective operation. - -The ops run the generated Cake kernels (``backend="cake"``, compute capability -10.0 or 10.3) and are driven by :class:`flashinfer.moe_ep.CakeAlltoAll`. The -TensorRT-LLM kernels (``backend="trtllm"``) were removed; use -:class:`flashinfer.moe_ep.NVLinkOneSidedAlltoAll`, which runs the current -TensorRT-LLM one-sided kernels. -""" - -from types import SimpleNamespace -from typing import Literal, Optional, Sequence - -import torch -import functools - -from ..api_logging import flashinfer_api - -from ..jit.comm import gen_moe_alltoall_module -from ..utils import register_custom_op, device_support_pdl -from ..tllm_enums import SfLayout - -# Number of uint64 words in the active_rank_mask ABI (must match kRankMaskWords in -# csrc/cake_moe_alltoall_dispatch.cuh, which is -# derived from kMaxRanks there). A single word covers up to 64 ranks. -MOE_A2A_RANK_MASK_WORDS = 1 -MoeAlltoAllTarget = Literal["sm100a", "sm103a"] -MoeAlltoAllBackend = Literal["cake"] - - -def moe_a2a_active_rank_mask(active_ranks: Sequence[int], ep_size: int) -> torch.Tensor: - r"""Build a CPU ``uint64`` active-rank bitmask for :func:`moe_a2a_dispatch` / - :func:`moe_a2a_combine`'s ``active_rank_mask`` argument. - - Parameters - ---------- - active_ranks : Sequence[int] - Ranks that are alive and should participate in the collective, e.g. ``[0, 1, 3]`` - for a 4-rank job where rank 2 has failed. Also accepts a 1D ``torch.Tensor`` of - rank indices. - ep_size : int - Total expert-parallel world size (all ranks outside ``[0, ep_size)`` are ignored). - - Returns - ------- - torch.Tensor - ``[MOE_A2A_RANK_MASK_WORDS]`` ``uint64`` CPU tensor with bit ``i`` set for each - active rank ``i``. - """ - mask = 0 - for rank in active_ranks: - rank = int(rank) - assert 0 <= rank < ep_size, f"rank {rank} out of range [0, {ep_size})" - mask |= 1 << rank - words = [ - (mask >> (64 * word)) & 0xFFFFFFFFFFFFFFFF - for word in range(MOE_A2A_RANK_MASK_WORDS) - ] - return torch.tensor(words, dtype=torch.uint64, device="cpu") - - -@functools.cache -def _moe_alltoall_target(device_index: int) -> MoeAlltoAllTarget: - capability = torch.cuda.get_device_capability(device_index) - if capability == (10, 0): - return "sm100a" - if capability == (10, 3): - return "sm103a" - raise ValueError( - "The MoE all-to-all ops require compute capability 10.0 or 10.3, got " - f"{capability}; use flashinfer.moe_ep.NVLinkOneSidedAlltoAll on other GPUs." - ) - - -@functools.cache -def _get_moe_alltoall_module_for_target(target: MoeAlltoAllTarget): - """Build or load the all-to-all module of one exact architecture.""" - module = gen_moe_alltoall_module(target).build_and_load() - # Each architecture registers its own custom ops. - op_suffix = f"_{target}" - - @register_custom_op( - f"flashinfer::moe_a2a_initialize{op_suffix}", - mutates_args=("workspace",), - ) - def moe_a2a_initialize( - workspace: torch.Tensor, - ep_rank: int, - ep_size: int, - max_num_tokens: int, - eplb_stats_num_experts: int = 0, - ): - return module.moe_a2a_initialize( - workspace, ep_rank, ep_size, max_num_tokens, eplb_stats_num_experts - ) - - @register_custom_op( - f"flashinfer::moe_a2a_dispatch{op_suffix}", - mutates_args=("workspace",), - ) - def moe_a2a_dispatch( - token_selected_experts: torch.Tensor, - input_payloads: list[torch.Tensor], - workspace: torch.Tensor, - metainfo: torch.Tensor, - runtime_max_tokens_per_rank: int, - ep_rank: int, - ep_size: int, - top_k: int, - num_experts: int, - enable_pdl: bool, - eplb_local_stats: Optional[torch.Tensor] = None, - enable_rank_mask: bool = False, - active_rank_mask: Optional[torch.Tensor] = None, - ): - """ - Dispatch tokens and payloads to expert ranks. - - Args: - token_selected_experts: [local_num_tokens, top_k] int32 tensor - input_payloads: List of [local_num_tokens, *] tensors to dispatch - workspace: [ep_size, size_per_rank] workspace tensor - metainfo: Metadata tensor from initialize - runtime_max_tokens_per_rank: Max tokens per rank in this batch - ep_rank: Current expert parallel rank - ep_size: Total expert parallel size - top_k: Number of experts per token - num_experts: Total number of experts - enable_pdl: Whether to use programmatic dependent launch - eplb_local_stats: Optional [eplb_stats_num_experts] int32 tensor of - this rank's local EPLB statistics to all-gather during dispatch - enable_rank_mask: Whether to instantiate the kernel variant that checks - active_rank_mask at all. False (default) compiles out every rank-mask - check for the common no-fault-tolerance case and requires - active_rank_mask to be omitted. - active_rank_mask: Optional CPU uint64 tensor of shape [MOE_A2A_RANK_MASK_WORDS] - (see :func:`moe_a2a_active_rank_mask`). Bit i set means rank i is alive and - participates in this collective; tokens routed to a masked-off rank are - dropped. Requires enable_rank_mask=True. - - Returns: - recv_offsets: List of offsets for each payload in the workspace - recv_sizes: List of sizes for each payload in the workspace - combine_payload_offset: Offset for combine payload region - eplb_gathered_stats_offset: Offset for the gathered EPLB stats - region, or -1 when EPLB is disabled - eplb_stats_num_experts: Number of experts in the EPLB stats, or 0 - when EPLB is disabled - """ - return module.moe_a2a_dispatch( - token_selected_experts, - input_payloads, - workspace, - metainfo, - runtime_max_tokens_per_rank, - ep_rank, - ep_size, - top_k, - num_experts, - enable_pdl, - eplb_local_stats, - enable_rank_mask, - active_rank_mask, - ) - - @register_custom_op( - f"flashinfer::moe_a2a_combine{op_suffix}", - mutates_args=("workspace",), - ) - def moe_a2a_combine( - payload: torch.Tensor, - local_num_tokens: int, - workspace: torch.Tensor, - metainfo: torch.Tensor, - runtime_max_tokens_per_rank: int, - ep_rank: int, - ep_size: int, - top_k: int, - combine_payload_offset: int, - payload_in_workspace: bool = False, - output_dtype: Optional[torch.dtype] = None, - output_scales: Optional[torch.Tensor] = None, - output_scalar_scale: float = 1.0, - sf_layout: Optional[SfLayout] = None, - use_low_precision: bool = False, - enable_pdl: bool = True, - enable_rank_mask: bool = False, - active_rank_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """ - Combine expert outputs back to originating tokens. - - Args: - payload: [ep_size, max_tokens, elements_per_token] tensor - local_num_tokens: Number of tokens on this rank - workspace: [ep_size, size_per_rank] workspace tensor - metainfo: Metadata tensor from initialize - runtime_max_tokens_per_rank: Max tokens per rank in this batch - ep_rank: Current expert parallel rank - ep_size: Total expert parallel size - top_k: Number of experts per token - combine_payload_offset: Offset from dispatch - payload_in_workspace: If True, payload is workspace-backed - output_dtype: Optional output data type. Supported types: - torch.bfloat16, - torch.float8_e4m3fn, - torch.uint8 (packed fp4) - output_scales: Optional output scale tensor for quantized outputs. Support types: - torch.uint8 (packed ue8m0), vector size of 32 - torch.float8_e4m3fn, vector size of 16 - output_scalar_scale: Per-tensor global scale applied before FP4 block scaling - (NVFP4 SFScaleVal). Defaults to 1.0; ignored by MXFP8/MXFP4 paths. - sf_layout: Output swizzle layout. Defaults to linear. - use_low_precision: If True, quantize payload to FP8 before combine - enable_pdl: Whether to use programmatic dependent launch - enable_rank_mask: Whether to instantiate the kernel variant that checks - active_rank_mask at all. False (default) compiles out the rank-mask checks - in peer synchronization for the common no-fault-tolerance case and requires - active_rank_mask to be omitted. - active_rank_mask: Optional CPU uint64 tensor of shape [MOE_A2A_RANK_MASK_WORDS] - (see :func:`moe_a2a_active_rank_mask`). Should match the mask passed to the - corresponding :func:`moe_a2a_dispatch` call (or be omitted from both). Requires - enable_rank_mask=True. - Returns: - output: [local_num_tokens, elements_per_token] tensor - """ - return module.moe_a2a_combine( - payload, - local_num_tokens, - workspace, - metainfo, - runtime_max_tokens_per_rank, - ep_rank, - ep_size, - top_k, - combine_payload_offset, - payload_in_workspace, - output_dtype, - output_scales, - output_scalar_scale, - sf_layout.value if sf_layout is not None else SfLayout.layout_linear.value, - use_low_precision, - enable_pdl, - enable_rank_mask, - active_rank_mask, - ) - - @register_custom_op( - f"flashinfer::moe_a2a_combine_into{op_suffix}", - mutates_args=("workspace", "output"), - ) - def moe_a2a_combine_into( - payload: torch.Tensor, - local_num_tokens: int, - workspace: torch.Tensor, - metainfo: torch.Tensor, - runtime_max_tokens_per_rank: int, - ep_rank: int, - ep_size: int, - top_k: int, - combine_payload_offset: int, - payload_in_workspace: bool, - output_dtype: Optional[torch.dtype], - output_scales: Optional[torch.Tensor], - output_scalar_scale: float, - sf_layout: Optional[SfLayout], - use_low_precision: bool, - enable_pdl: bool, - enable_rank_mask: bool, - active_rank_mask: Optional[torch.Tensor], - output: torch.Tensor, - ) -> None: - module.moe_a2a_combine_into( - payload, - local_num_tokens, - workspace, - metainfo, - runtime_max_tokens_per_rank, - ep_rank, - ep_size, - top_k, - combine_payload_offset, - payload_in_workspace, - output_dtype, - output_scales, - output_scalar_scale, - sf_layout.value if sf_layout is not None else SfLayout.layout_linear.value, - use_low_precision, - enable_pdl, - enable_rank_mask, - active_rank_mask, - output, - ) - - @register_custom_op( - f"flashinfer::moe_a2a_sanitize_expert_ids{op_suffix}", - mutates_args=("expert_ids",), - ) - def moe_a2a_sanitize_expert_ids( - expert_ids: torch.Tensor, - workspace: torch.Tensor, - metainfo: torch.Tensor, - ep_rank: int, - invalid_expert_id: int, - enable_pdl: bool, - ): - return module.moe_a2a_sanitize_expert_ids( - expert_ids, workspace, metainfo, ep_rank, invalid_expert_id, enable_pdl - ) - - @register_custom_op( - f"flashinfer::moe_a2a_get_metainfo_index_pairs{op_suffix}", - mutates_args=[], - ) - def moe_a2a_get_metainfo_index_pairs(): - """ - Get all metainfo index constants from C++. - - Returns: - Tuple of (names, values) where names is a list of constant names - and values is a list of their corresponding integer values - """ - return module.moe_a2a_get_metainfo_index_pairs() - - @register_custom_op( - f"flashinfer::moe_a2a_get_aux_data_size{op_suffix}", - mutates_args=[], - ) - def moe_a2a_get_aux_data_size( - ep_size: int, - max_num_tokens: int, - eplb_stats_num_experts: int = 0, - ): - """ - Get the auxilary datasize per rank of the MoE all-to-all workspace. - - Args: - ep_size: Total expert parallel size - max_num_tokens: Maximum number of tokens across all ranks - eplb_stats_num_experts: Number of experts reserved for EPLB stats - (0 disables the EPLB region) - - Returns: - aux_data_size: Size of the auxilary data per rank in bytes - """ - return module.moe_a2a_get_aux_data_size( - ep_size, max_num_tokens, eplb_stats_num_experts - ) - - return SimpleNamespace( - moe_a2a_initialize=moe_a2a_initialize, - moe_a2a_dispatch=moe_a2a_dispatch, - moe_a2a_combine=moe_a2a_combine, - moe_a2a_combine_into=moe_a2a_combine_into, - moe_a2a_sanitize_expert_ids=moe_a2a_sanitize_expert_ids, - moe_a2a_get_metainfo_index_pairs=moe_a2a_get_metainfo_index_pairs, - moe_a2a_get_aux_data_size=moe_a2a_get_aux_data_size, - ) - - -def get_moe_alltoall_module(backend: MoeAlltoAllBackend = "cake"): - """Return the Cake all-to-all module for the current device.""" - if backend == "trtllm": - raise ValueError( - 'The backend="trtllm" MoE all-to-all ops were removed; use ' - "flashinfer.moe_ep.NVLinkOneSidedAlltoAll, which runs the current " - "TensorRT-LLM one-sided kernels." - ) - if backend != "cake": - raise ValueError(f"Unknown MoE all-to-all backend {backend!r}; expected 'cake'") - device_index = torch.cuda.current_device() - return _get_moe_alltoall_module_for_target(_moe_alltoall_target(device_index)) - - -def _clear_moe_alltoall_module_cache() -> None: - _moe_alltoall_target.cache_clear() - _get_moe_alltoall_module_for_target.cache_clear() - - -get_moe_alltoall_module.cache_clear = ( # type: ignore[attr-defined] - _clear_moe_alltoall_module_cache -) - - -@flashinfer_api -def moe_a2a_initialize( - workspace: torch.Tensor, - ep_rank: int, - ep_size: int, - max_num_tokens: int, - eplb_stats_num_experts: int = 0, - *, - backend: MoeAlltoAllBackend = "cake", -): - r"""Initialize the MoE all-to-all workspace and return a metainfo tensor. - - The metainfo tensor encodes per-rank offsets and bookkeeping required by - :func:`moe_a2a_dispatch` and :func:`moe_a2a_combine`; it must be passed - back into those routines for the same workspace. ``moe_a2a_initialize`` - is idempotent and must be called once per workspace allocation before any - dispatch/combine. - - Parameters - ---------- - workspace : torch.Tensor - ``[ep_size, size_per_rank]`` shared workspace tensor. - ep_rank : int - Current expert-parallel rank. - ep_size : int - Total expert-parallel world size. - max_num_tokens : int - Maximum number of tokens any rank may dispatch in a single call; - used to size the metainfo allocation. - eplb_stats_num_experts : int - Number of experts to reserve space for in the EPLB gathered-stats - region. ``0`` (default) disables the EPLB region; when non-zero it - must match the ``eplb_local_stats`` length passed to - :func:`moe_a2a_dispatch`. - backend : {"cake"} - The generated Blackwell kernels (CC 10.0/10.3), the only backend. - - Returns - ------- - torch.Tensor - Metainfo tensor opaque to callers; pass it to subsequent - ``moe_a2a_*`` calls. - """ - return get_moe_alltoall_module(backend).moe_a2a_initialize( - workspace, ep_rank, ep_size, max_num_tokens, eplb_stats_num_experts - ) - - -@flashinfer_api -def moe_a2a_wrap_payload_tensor_in_workspace( - workspace: torch.Tensor, - leading_shape: list[int], - slice_start: int, - slice_end: int, - dtype: torch.dtype, -) -> torch.Tensor: - r"""Wrap a slice of the shared workspace as a typed tensor view. - - Parameters - ---------- - workspace : torch.Tensor - ``[ep_size, size_per_rank]`` (or ``[size_per_rank]``) workspace - tensor. - leading_shape : list[int] - Leading shape of the resulting view. The trailing dimension is - inferred from ``slice_end - slice_start`` and ``dtype``. - slice_start : int - Start offset (in bytes from the beginning of the workspace) of the - slice to wrap. - slice_end : int - End offset (in bytes) of the slice. Must lie within a single rank. - dtype : torch.dtype - Element dtype of the resulting view. - - Returns - ------- - torch.Tensor - A workspace-backed tensor of shape ``leading_shape + [-1]``. - """ - if workspace.ndim == 1: - workspace = workspace.unsqueeze(0) - workspace_base = workspace.view(dtype=torch.uint8) - assert workspace.ndim == 2, "workspace must be shape [ep_size, size_per_rank]" - assert slice_end - slice_start <= workspace_base.shape[1], ( - "slice_end - slice_start must belong to a single rank" - ) - slice_rank = slice_start // workspace_base.stride(0) - local_slice_start = slice_start % workspace_base.stride(0) - slice_length = slice_end - slice_start - local_slice_end = local_slice_start + slice_length - assert local_slice_end <= workspace_base.shape[1], ( - "slice must fall within the workspace size per rank" - ) - result = ( - workspace_base[slice_rank, local_slice_start:local_slice_end] - .view(dtype=dtype) - .view(*leading_shape, -1) - ) - return result - - -@flashinfer_api -def moe_a2a_dispatch( - token_selected_experts: torch.Tensor, - input_payloads: list[torch.Tensor], - workspace: torch.Tensor, - metainfo: torch.Tensor, - runtime_max_tokens_per_rank: int, - ep_rank: int, - ep_size: int, - top_k: int, - num_experts: int, - enable_pdl: Optional[bool] = None, - eplb_local_stats: Optional[torch.Tensor] = None, - enable_rank_mask: bool = False, - active_rank_mask: Optional[torch.Tensor] = None, - *, - backend: MoeAlltoAllBackend = "cake", - recv_view_cache: Optional[dict] = None, -): - r"""Dispatch tokens and payloads to their target expert ranks. - - Parameters - ---------- - token_selected_experts : torch.Tensor - ``[local_num_tokens, top_k]`` ``int32`` tensor of expert assignments. - input_payloads : list[torch.Tensor] - Per-token payload tensors, each shaped ``[local_num_tokens, *]``. - workspace : torch.Tensor - ``[ep_size, size_per_rank]`` shared workspace. - metainfo : torch.Tensor - Metainfo tensor returned by :func:`moe_a2a_initialize`. - runtime_max_tokens_per_rank : int - Maximum tokens per rank for this batch (must be ``<=`` the - ``max_num_tokens`` used at initialize time). - ep_rank : int - Current expert-parallel rank. - ep_size : int - Total expert-parallel world size. - top_k : int - Number of experts assigned per token. - num_experts : int - Total number of experts. - enable_pdl : Optional[bool] - Whether to use programmatic dependent launch. ``None`` auto-detects - from the device. - eplb_local_stats : Optional[torch.Tensor] - Optional ``[eplb_stats_num_experts]`` ``int32`` tensor of this rank's - local EPLB statistics. When provided, the dispatch all-gathers it - across ranks and returns the result as ``eplb_gathered_stats``. The - length must match the ``eplb_stats_num_experts`` passed to - :func:`moe_a2a_initialize`. - enable_rank_mask : bool - Whether to instantiate the kernel variant that checks ``active_rank_mask`` at - all. ``False`` (default) compiles out every rank-mask check for the common - no-fault-tolerance case and requires ``active_rank_mask`` to be omitted. - active_rank_mask : Optional[torch.Tensor] - Optional CPU ``uint64`` tensor of shape ``[MOE_A2A_RANK_MASK_WORDS]`` (see - :func:`moe_a2a_active_rank_mask`). Bit ``i`` set means rank ``i`` is alive and - participates in this collective; tokens routed to a masked-off rank are dropped - instead of hanging the collective. Requires ``enable_rank_mask=True``; the local - ``ep_rank``'s own bit must always be set. - Masking a peer does not advance that peer's transport epoch. Before a - skipped peer rejoins, all ranks must quiesce and coordinate workspace - reinitialization; merely restoring its mask bit is insufficient. - backend : {"cake"} - The generated Blackwell kernels (CC 10.0/10.3), the only backend; - must match workspace initialization, combine, and all peer ranks. - recv_view_cache : dict, optional - Opaque cache of receive views. Pass an empty dictionary to enable - caching; ``None`` (default) disables it. The cache retains a reference - to ``workspace`` and is cleared when a different workspace tensor is - passed. Keep the backing allocation alive and do not change workspace - or cached-view shape, strides, or storage in place. Clear the dictionary - to release its references; otherwise, leave its contents unchanged. - - Returns - ------- - Tuple[list[torch.Tensor], int, Optional[torch.Tensor]] - ``(output_payloads, combine_payload_offset, eplb_gathered_stats)``. - ``output_payloads`` is a list of workspace-backed views, one per - ``input_payloads`` entry, that contains the data routed to this rank. - ``combine_payload_offset`` is the workspace offset reserved for the - matching :func:`moe_a2a_combine` call. ``eplb_gathered_stats`` is a - workspace-backed ``[ep_size, eplb_stats_num_experts]`` ``int32`` view - (row ``r`` holds rank ``r``'s ``eplb_local_stats``) when - ``eplb_local_stats`` was provided, else ``None``. - """ - if enable_pdl is None: - enable_pdl = device_support_pdl(token_selected_experts.device) - ( - recv_offsets, - recv_sizes, - combine_payload_offset, - eplb_gathered_stats_offset, - eplb_stats_num_experts, - ) = get_moe_alltoall_module(backend).moe_a2a_dispatch( - token_selected_experts, - input_payloads, - workspace, - metainfo, - runtime_max_tokens_per_rank, - ep_rank, - ep_size, - top_k, - num_experts, - enable_pdl, - eplb_local_stats, - enable_rank_mask, - active_rank_mask, - ) - - # Bind once per dispatch without pointer/device queries in the hot path. - if ( - recv_view_cache is not None - and recv_view_cache.get("_workspace") is not workspace - ): - recv_view_cache.clear() - recv_view_cache["_workspace"] = workspace - - output_payloads = [] - for input_payload, offset, size in zip( - input_payloads, recv_offsets, recv_sizes, strict=True - ): - # This uses absolute offsets in the workspace, so skip indexing into the workspace - # Reuse views to avoid four tensor constructions per payload. Cached views - # still observe fresh writes to the workspace bound above. - # ``int()`` guards the key rather than fixing an observed miss: tvm-ffi - # materializes an ``Array`` element as a real Python ``int`` - # today, so these already hash by value. Coercing keeps that property - # local to this function instead of resting on an FFI implementation - # detail -- an int-like that hashed by identity would miss on every - # lookup, silently disabling the cache while the dict still grew. Same - # idiom as the ``int()`` on FFI scalars in ``_init_constants`` below. - key = ( - ep_size, - runtime_max_tokens_per_rank, - int(offset), - int(size), - input_payload.dtype, - ) - payload_view = None if recv_view_cache is None else recv_view_cache.get(key) - if payload_view is None: - payload_view = moe_a2a_wrap_payload_tensor_in_workspace( - workspace, - [ep_size, runtime_max_tokens_per_rank], - offset, - offset + size, - input_payload.dtype, - ) - if recv_view_cache is not None: - recv_view_cache[key] = payload_view - output_payloads.append(payload_view) - - eplb_gathered_stats = None - if eplb_gathered_stats_offset >= 0: - eplb_gathered_stats = moe_a2a_wrap_payload_tensor_in_workspace( - workspace, - [ep_size], - eplb_gathered_stats_offset, - eplb_gathered_stats_offset + ep_size * eplb_stats_num_experts * 4, - torch.int32, - ) - - return output_payloads, combine_payload_offset, eplb_gathered_stats - - -@flashinfer_api -def moe_a2a_combine( - payload: torch.Tensor, - local_num_tokens: int, - workspace: torch.Tensor, - metainfo: torch.Tensor, - runtime_max_tokens_per_rank: int, - ep_rank: int, - ep_size: int, - top_k: int, - combine_payload_offset: int, - payload_in_workspace: bool = False, - output_dtype: Optional[torch.dtype] = None, - output_scales: Optional[torch.Tensor] = None, - output_scalar_scale: float = 1.0, - sf_layout: SfLayout = SfLayout.layout_linear, - output: Optional[torch.Tensor] = None, - *, - use_low_precision: bool = False, - enable_pdl: Optional[bool] = None, - enable_rank_mask: bool = False, - active_rank_mask: Optional[torch.Tensor] = None, - backend: MoeAlltoAllBackend = "cake", -) -> torch.Tensor: - r"""Combine per-expert outputs back to the originating ranks. - - Inverse of :func:`moe_a2a_dispatch`: scatters the rank-local expert - output rows back to the ranks that supplied the original tokens. - - Parameters - ---------- - payload : torch.Tensor - Output payload to send back to the source ranks. Shape - ``[ep_size, runtime_max_tokens_per_rank, *]`` regardless of - ``payload_in_workspace``: in both cases the payload holds the - per-expert-rank outputs to be combined back to the source ranks. - Only the backing memory differs (caller-supplied vs. workspace-backed - view produced by :func:`moe_a2a_wrap_payload_tensor_in_workspace`). - local_num_tokens : int - Number of tokens originally dispatched from this rank. - workspace : torch.Tensor - Shared workspace tensor (same one passed to dispatch). - metainfo : torch.Tensor - Metainfo tensor returned by :func:`moe_a2a_initialize`. - runtime_max_tokens_per_rank : int - Same value passed to :func:`moe_a2a_dispatch`. - ep_rank : int - Current expert-parallel rank. - ep_size : int - Total expert-parallel world size. - top_k : int - Number of experts assigned per token. - combine_payload_offset : int - Offset returned by :func:`moe_a2a_dispatch`. - payload_in_workspace : bool - ``True`` if ``payload`` is already a workspace-backed view (skips - the staging copy). Defaults to ``False``. - output_dtype : Optional[torch.dtype] - Optional output data type. Currently supports ``torch.bfloat16``, - ``torch.float8_e4m3fn``, and ``torch.uint8`` (packed fp4). - output_scales : Optional[torch.Tensor] - Contiguous CUDA scale tensor for quantized outputs. MXFP8 and MXFP4 use - UE8M0 scales packed in ``torch.uint8`` with vector size 32; NVFP4 uses - UE4M3 scales in ``torch.float8_e4m3fn`` with vector size 16. Its extent - must exactly match ``sf_layout``, including layout padding. - output_scalar_scale : float - Per-tensor global scale applied before FP4 block scaling - (NVFP4 SFScaleVal). Defaults to ``1.0``; ignored by MXFP8/MXFP4 - paths. - sf_layout : SfLayout - Output swizzle layout. Defaults to ``SfLayout.layout_linear``. - output : Optional[torch.Tensor] - Caller-provided contiguous output tensor. Its shape and dtype must - match the requested combine output, and it must be on the same device - as ``payload``. - use_low_precision : bool - If ``True``, quantize the recv-buffer payload to FP8 (e4m3) before - accumulating; the combine upcasts to a bf16 output. - enable_pdl : Optional[bool] - Whether to use programmatic dependent launch. ``None`` auto-detects - from the device. - enable_rank_mask : bool - Whether to instantiate the kernel variant that checks ``active_rank_mask`` at - all. ``False`` (default) compiles out the rank-mask checks in peer - synchronization for the common no-fault-tolerance case and requires - ``active_rank_mask`` to be omitted. - active_rank_mask : Optional[torch.Tensor] - Optional CPU ``uint64`` tensor of shape ``[MOE_A2A_RANK_MASK_WORDS]`` (see - :func:`moe_a2a_active_rank_mask`). Should match the mask passed to the - corresponding :func:`moe_a2a_dispatch` call (or be omitted from both). Requires - ``enable_rank_mask=True``. - backend : {"cake"} - The generated Blackwell kernels (CC 10.0/10.3), the only backend. - - Returns - ------- - torch.Tensor - ``[local_num_tokens, *]`` tensor with the combined outputs. - """ - if enable_pdl is None: - enable_pdl = device_support_pdl(payload.device) - if output is not None: - if not output.is_cuda: - raise ValueError( - f"output must be a CUDA tensor, got device={output.device}" - ) - if not output.is_contiguous(): - raise ValueError(f"output must be contiguous, got stride={output.stride()}") - - module = get_moe_alltoall_module(backend) - args = ( - payload, - local_num_tokens, - workspace, - metainfo, - runtime_max_tokens_per_rank, - ep_rank, - ep_size, - top_k, - combine_payload_offset, - payload_in_workspace, - output_dtype, - output_scales, - output_scalar_scale, - sf_layout, - use_low_precision, - enable_pdl, - enable_rank_mask, - active_rank_mask, - ) - if output is None: - return module.moe_a2a_combine(*args) - module.moe_a2a_combine_into(*args, output) - return output - - -@flashinfer_api -def moe_a2a_sanitize_expert_ids( - expert_ids: torch.Tensor, - workspace: torch.Tensor, - metainfo: torch.Tensor, - ep_rank: int, - invalid_expert_id: int, - enable_pdl: Optional[bool] = None, - *, - backend: MoeAlltoAllBackend = "cake", -): - r"""Sanitize invalid slots that contain no token routed to this rank by - setting their expert IDs to ``invalid_expert_id``. - - Parameters - ---------- - expert_ids : torch.Tensor - ``[local_num_tokens, top_k]`` ``int32`` tensor of expert assignments - (mutated in place). - workspace : torch.Tensor - Shared workspace tensor. - metainfo : torch.Tensor - Metainfo tensor returned by :func:`moe_a2a_initialize`. - ep_rank : int - Current expert-parallel rank. - invalid_expert_id : int - Value to write into slots that received no token (per - ``recv_counters``, e.g. padding beyond a source rank's valid count). - enable_pdl : Optional[bool] - Whether to use programmatic dependent launch. ``None`` auto-detects - from the device. - backend : {"cake"} - The generated Blackwell kernels (CC 10.0/10.3), the only backend. - """ - if enable_pdl is None: - enable_pdl = device_support_pdl(expert_ids.device) - return get_moe_alltoall_module(backend).moe_a2a_sanitize_expert_ids( - expert_ids, workspace, metainfo, ep_rank, invalid_expert_id, enable_pdl - ) - - -@flashinfer_api -def moe_a2a_get_workspace_size_per_rank( - ep_size: int, - max_num_tokens: int, - total_dispatch_payload_size_per_token: int, - combine_payload_size_per_token: int, - eplb_stats_num_experts: int = 0, - *, - backend: MoeAlltoAllBackend = "cake", -): - r"""Compute the per-rank workspace size for the MoE all-to-all primitive. - - Parameters - ---------- - ep_size : int - Total expert-parallel world size. - max_num_tokens : int - Maximum number of tokens across all ranks. - total_dispatch_payload_size_per_token : int - Sum (in bytes) of all per-token payloads sent during the dispatch - phase. - combine_payload_size_per_token : int - Per-token payload size (in bytes) sent back during the combine - phase. - eplb_stats_num_experts : int - Number of experts reserved for the EPLB gathered-stats region - (``0`` disables it). - backend : {"cake"} - The generated Blackwell kernels (CC 10.0/10.3), the only backend. - - Returns - ------- - int - Required workspace size per rank, in bytes. - """ - aux_data_size = get_moe_alltoall_module(backend).moe_a2a_get_aux_data_size( - ep_size, - max_num_tokens, - eplb_stats_num_experts, - ) - - def pad_up(x, y): - return ((x + y - 1) // y) * y - - # Pad to 128 bytes to ensure alignment. This matches the implementation of C++ torch OP code. - return ( - pad_up(aux_data_size, 128) - + pad_up(ep_size * max_num_tokens * total_dispatch_payload_size_per_token, 128) - + pad_up(ep_size * max_num_tokens * combine_payload_size_per_token, 128) - ) - - -__all__ = [ - "moe_a2a_active_rank_mask", - "moe_a2a_combine", - "moe_a2a_dispatch", - "moe_a2a_get_workspace_size_per_rank", - "moe_a2a_initialize", - "moe_a2a_sanitize_expert_ids", - "moe_a2a_wrap_payload_tensor_in_workspace", -] diff --git a/flashinfer/jit/__init__.py b/flashinfer/jit/__init__.py index de9c7ea5b38..f7d0607bb45 100644 --- a/flashinfer/jit/__init__.py +++ b/flashinfer/jit/__init__.py @@ -78,15 +78,17 @@ from .core import sm121a_nvcc_flags as sm121a_nvcc_flags from .core import current_compilation_context as current_compilation_context from .cubin_loader import setup_cubin_loader -from .comm import gen_comm_alltoall_module as gen_comm_alltoall_module from .comm import gen_trtllm_mnnvl_comm_module as gen_trtllm_mnnvl_comm_module from .comm import gen_trtllm_comm_module as gen_trtllm_comm_module from .comm import gen_vllm_comm_module as gen_vllm_comm_module from .comm import gen_pcie_ipc_comm_module as gen_pcie_ipc_comm_module -from .comm import gen_moe_alltoall_module as gen_moe_alltoall_module +from .comm import gen_moe_ep_cake_module as gen_moe_ep_cake_module from .comm import ( gen_moe_ep_nvlink_one_sided_module as gen_moe_ep_nvlink_one_sided_module, ) +from .comm import ( + gen_moe_ep_nvlink_two_sided_module as gen_moe_ep_nvlink_two_sided_module, +) from .comm import gen_dcp_alltoall_module as gen_dcp_alltoall_module from .dsv3_optimizations import ( gen_dsv3_router_gemm_module as gen_dsv3_router_gemm_module, diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index 1a81757d19e..fa999d7b119 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -30,12 +30,13 @@ from typing import Any, Dict, List, Literal, Tuple -def gen_comm_alltoall_module() -> JitSpec: +def gen_moe_ep_nvlink_two_sided_module() -> JitSpec: + """NVLink two-sided MoE all-to-all behind ``flashinfer.moe_ep.NVLinkTwoSidedAlltoAll``.""" return gen_jit_spec( - "comm", + "moe_ep_nvlink_two_sided_alltoall", [ - jit_env.FLASHINFER_CSRC_DIR / "trtllm_alltoall.cu", - jit_env.FLASHINFER_CSRC_DIR / "trtllm_alltoall_prepare.cu", + jit_env.FLASHINFER_CSRC_DIR / "moe_ep_nvlink_two_sided_alltoall.cu", + jit_env.FLASHINFER_CSRC_DIR / "moe_ep_nvlink_two_sided_prepare.cu", ], ) @@ -304,31 +305,31 @@ def gen_ulysses_a2a_module() -> JitSpec: ) -MoeAlltoAllTarget = Literal["sm100a", "sm103a"] +CakeAlltoAllTarget = Literal["sm100a", "sm103a"] -_MOE_ALLTOALL_GENERATED_SOURCE = { +_CAKE_ALLTOALL_GENERATED_SOURCE = { "sm100a": "mnnvl_moe_alltoall_sm100.cu", "sm103a": "mnnvl_moe_alltoall_sm103.cu", } -_MOE_ALLTOALL_NVCC_FLAGS = { +_CAKE_ALLTOALL_NVCC_FLAGS = { "sm100a": sm100a_nvcc_flags, "sm103a": sm103a_nvcc_flags, } -def gen_moe_alltoall_module(target: MoeAlltoAllTarget) -> JitSpec: - """Cake MoE all-to-all module behind the ``flashinfer.comm`` ``moe_a2a_*`` ops.""" - if target not in _MOE_ALLTOALL_GENERATED_SOURCE: - raise ValueError(f"unsupported MNNVL MoE all-to-all target: {target}") - module_name = f"mnnvl_moe_alltoall_{target}" +def gen_moe_ep_cake_module(target: CakeAlltoAllTarget) -> JitSpec: + """Cake MoE all-to-all of one architecture behind ``flashinfer.moe_ep.CakeAlltoAll``.""" + if target not in _CAKE_ALLTOALL_GENERATED_SOURCE: + raise ValueError(f"unsupported Cake MoE all-to-all target: {target}") + module_name = f"moe_ep_cake_alltoall_{target}" kernel_sources = [ jit_env.FLASHINFER_CSRC_DIR / "cake_moe_alltoall_dispatch.cu", jit_env.FLASHINFER_CSRC_DIR / "generated" / "moe_alltoall" - / _MOE_ALLTOALL_GENERATED_SOURCE[target], + / _CAKE_ALLTOALL_GENERATED_SOURCE[target], ] - arch_flags = _MOE_ALLTOALL_NVCC_FLAGS[target] + arch_flags = _CAKE_ALLTOALL_NVCC_FLAGS[target] return gen_jit_spec( module_name, [ diff --git a/flashinfer/moe_ep/backends/split/comm/cake/communication.py b/flashinfer/moe_ep/backends/split/comm/cake/communication.py index 983fd0adf44..5dfd9d718b4 100644 --- a/flashinfer/moe_ep/backends/split/comm/cake/communication.py +++ b/flashinfer/moe_ep/backends/split/comm/cake/communication.py @@ -1,18 +1,16 @@ """NVLink one-sided MoE communication with the generated Cake kernels. The Cake kernels are generated for Blackwell (compute capability 10.0 or -10.3) and implement the one-sided put/get protocol of the ``moe_a2a_*`` ops in -:mod:`flashinfer.comm.trtllm_moe_alltoall`, selected with ``backend="cake"``: -the dispatch kernel writes each token into the receive buffers of the ranks -that own its experts and the combine kernel reads the expert outputs back and -reduces them locally. The symmetric workspace grows with -``ep_size * max_tokens_per_rank``. +10.3) and implement a one-sided put/get protocol: the dispatch kernel writes +each token into the receive buffers of the ranks that own its experts and the +combine kernel reads the expert outputs back and reduces them locally. The +symmetric workspace grows with ``ep_size * max_tokens_per_rank``. """ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Optional from .....core.comm.communication import ( MoEEpCommParams, @@ -39,7 +37,6 @@ logger = logging.getLogger(__name__) _CAKE_COMPUTE_CAPABILITIES = ((10, 0), (10, 3)) -_BACKEND: Final = "cake" # Alignment the Cake ops give the workspace regions. _ALIGNMENT = 128 @@ -73,9 +70,7 @@ def __init__( mapping, mnnvl_config = mnnvl_mapping_and_config(bootstrap, config.comm_backend) super().__init__(bootstrap, params) from ......comm.mnnvl import MnnvlMemory - from ......comm.trtllm_moe_alltoall import ( - moe_a2a_get_workspace_size_per_rank, - ) + from . import kernels if config.extra_payload_bytes_per_token < 0: raise ValueError("extra_payload_bytes_per_token must be non-negative") @@ -93,13 +88,12 @@ def __init__( combine_bytes_per_token = params.hidden_size * max( params.token_dtype.itemsize, 2 ) - workspace_size_per_rank = moe_a2a_get_workspace_size_per_rank( + workspace_size_per_rank = kernels.workspace_size_per_rank( self.ep_size, params.max_tokens_per_rank, dispatch_bytes_per_token, combine_bytes_per_token, config.eplb_stats_num_experts, - backend=_BACKEND, ) MnnvlMemory.initialize() @@ -141,10 +135,7 @@ def _acquire_workspace( import torch from ......comm.mnnvl import MnnvlMemory - from ......comm.trtllm_moe_alltoall import ( - get_moe_alltoall_module, - moe_a2a_initialize, - ) + from . import kernels key = ( self.ep_rank, @@ -158,23 +149,17 @@ def _acquire_workspace( if state is None: mnnvl_mem = MnnvlMemory(mapping, workspace_size_per_rank) workspace = mnnvl_mem.as_torch_strided_tensor(torch.uint8) - metainfo = moe_a2a_initialize( + metainfo = kernels.initialize( workspace, self.ep_rank, self.ep_size, self.params.max_tokens_per_rank, self.config.eplb_stats_num_experts, - backend=_BACKEND, ) # No peer may publish into this workspace until every rank has # cleared its own slice. MnnvlMemory.allocated_map[mnnvl_mem.ptr].comm.barrier() - names, values = get_moe_alltoall_module( - _BACKEND - ).moe_a2a_get_metainfo_index_pairs() - payload_index = dict(zip(names, values, strict=True))[ - "MOE_A2A_PAYLOAD_DATA_OFFSET_INDEX" - ] + payload_index = kernels.metainfo_index()["PAYLOAD_DATA_OFFSET_INDEX"] payload_offset = int(metainfo[payload_index]) # The ops place the combine payload right after the current round's # dispatch payloads, so it moves with the token count and payload @@ -240,15 +225,13 @@ def dispatch( """See :meth:`MoEEpCommunication.dispatch`. ``active_rank_mask`` is a CPU ``uint64`` mask from - :func:`flashinfer.comm.moe_a2a_active_rank_mask`; tokens routed to a - masked-off rank are dropped. Requires ``enable_rank_mask=True``. + :func:`flashinfer.moe_ep.backends.split.comm.cake.kernels.active_rank_mask`; + tokens routed to a masked-off rank are dropped. Requires + ``enable_rank_mask=True``. """ import torch - from ......comm.trtllm_moe_alltoall import ( - moe_a2a_dispatch, - moe_a2a_sanitize_expert_ids, - ) + from . import kernels state = self._live_state() check_round_owner(self, state) @@ -291,7 +274,7 @@ def dispatch( "(MoEEpCommParams.dispatch_format, extra_payload_bytes_per_token)" ) - recv, combine_offset, eplb_gathered_stats = moe_a2a_dispatch( + recv, combine_offset, eplb_gathered_stats = kernels.dispatch( topk_ids, payloads, self.workspace, @@ -305,17 +288,15 @@ def dispatch( enable_rank_mask=self.config.enable_rank_mask, active_rank_mask=active_rank_mask, enable_pdl=self.config.enable_pdl, - backend=_BACKEND, recv_view_cache=state["views"], ) - moe_a2a_sanitize_expert_ids( + kernels.sanitize_expert_ids( recv[expert_id_index], self.workspace, self.metainfo, self.ep_rank, self.params.invalid_expert_id, enable_pdl=self.config.enable_pdl, - backend=_BACKEND, ) self._round = { "local_num_tokens": topk_ids.shape[0], @@ -339,9 +320,7 @@ def dispatch( def get_combine_input_buffer(self, dtype: "torch.dtype") -> "torch.Tensor": """``[ep_size * tokens_per_rank, hidden_size]`` view of the combine payload region of this rank's workspace.""" - from ......comm.trtllm_moe_alltoall import ( - moe_a2a_wrap_payload_tensor_in_workspace, - ) + from .kernels import wrap_workspace self._live_state() round_state = self._round @@ -355,7 +334,7 @@ def get_combine_input_buffer(self, dtype: "torch.dtype") -> "torch.Tensor": f"a {dtype} combine input of {end - start} bytes exceeds the " "combine region" ) - buffer = moe_a2a_wrap_payload_tensor_in_workspace( + buffer = wrap_workspace( self.workspace[self.ep_rank, :], [rows], start, end, dtype ) round_state["combine_buffer"] = buffer @@ -372,7 +351,7 @@ def combine( ``active_rank_mask`` must match the mask passed to :meth:`dispatch`. """ - from ......comm.trtllm_moe_alltoall import moe_a2a_combine + from . import kernels self._live_state() round_state = self._round @@ -392,7 +371,7 @@ def combine( in_workspace = ( buffer is not None and expert_output.data_ptr() == buffer.data_ptr() ) - combined = moe_a2a_combine( + combined = kernels.combine( payload, round_state["local_num_tokens"], self.workspace, @@ -408,7 +387,6 @@ def combine( active_rank_mask=active_rank_mask, enable_pdl=self.config.enable_pdl, output=output, - backend=_BACKEND, ) self._round = None self._state["round_owner"] = None @@ -437,18 +415,17 @@ def checkpoint_restore(self, comm_backend: "CommBackend") -> None: """ import torch - from ......comm.trtllm_moe_alltoall import moe_a2a_initialize + from . import kernels state = self._live_state() if not remap_mnnvl_memory(state["mnnvl_mem"], comm_backend): return - metainfo = moe_a2a_initialize( + metainfo = kernels.initialize( self.workspace, self.ep_rank, self.ep_size, self.params.max_tokens_per_rank, self.config.eplb_stats_num_experts, - backend=_BACKEND, ) if not torch.equal(metainfo, self.metainfo): raise RuntimeError( diff --git a/flashinfer/moe_ep/backends/split/comm/cake/kernels.py b/flashinfer/moe_ep/backends/split/comm/cake/kernels.py new file mode 100644 index 00000000000..7b01898d59c --- /dev/null +++ b/flashinfer/moe_ep/backends/split/comm/cake/kernels.py @@ -0,0 +1,373 @@ +"""Kernel-level operations of the generated Cake MoE all-to-all kernels. + +These wrap the JIT module behind :class:`flashinfer.moe_ep.CakeAlltoAll`: they +size and initialize a workspace and run one dispatch / combine round on it. +The kernels are generated for compute capability 10.0 (``sm100a``) and 10.3 +(``sm103a``); each architecture has its own module. +""" + +from __future__ import annotations + +import functools +from typing import Any, Optional, Sequence + +import torch + +from ......tllm_enums import SfLayout +from ......utils import device_support_pdl + +# Number of uint64 words in the active_rank_mask ABI (must match kRankMaskWords in +# csrc/cake_moe_alltoall_dispatch.cuh, which is derived from kMaxRanks there). A +# single word covers up to 64 ranks. +RANK_MASK_WORDS = 1 + +_TARGETS = {(10, 0): "sm100a", (10, 3): "sm103a"} + + +@functools.cache +def get_cake_module(target: str) -> Any: + """Build (on first use) and load the module of one architecture.""" + from ......jit.comm import gen_moe_ep_cake_module + + return gen_moe_ep_cake_module(target).build_and_load() # type: ignore[arg-type] + + +@functools.cache +def _device_target(device_index: int) -> str: + capability = torch.cuda.get_device_capability(device_index) + target = _TARGETS.get(capability) + if target is None: + raise ValueError( + "The Cake MoE all-to-all kernels require compute capability 10.0 or " + f"10.3, got {capability}; use flashinfer.moe_ep.NVLinkOneSidedAlltoAll " + "on other GPUs." + ) + return target + + +def current_module() -> Any: + """The module of the current CUDA device's architecture.""" + return get_cake_module(_device_target(torch.cuda.current_device())) + + +def metainfo_index() -> dict[str, int]: + """Metainfo field indices, keyed by name without the ``MOE_A2A_`` prefix.""" + names, values = current_module().moe_a2a_get_metainfo_index_pairs() + return { + str(name).removeprefix("MOE_A2A_"): int(value) + for name, value in zip(names, values, strict=True) + } + + +def active_rank_mask(active_ranks: Sequence[int], ep_size: int) -> torch.Tensor: + """Build the CPU ``uint64`` ``active_rank_mask`` of :func:`dispatch` and + :func:`combine`. + + ``active_ranks`` lists the ranks that are alive and participate in the + collective, e.g. ``[0, 1, 3]`` for a 4-rank job whose rank 2 failed; a 1D + tensor of rank indices works too. Returns a ``[RANK_MASK_WORDS]`` tensor + with bit ``i`` set for each active rank ``i``. + """ + mask = 0 + for rank in active_ranks: + rank = int(rank) + if not 0 <= rank < ep_size: + raise ValueError(f"rank {rank} out of range [0, {ep_size})") + mask |= 1 << rank + words = [ + (mask >> (64 * word)) & 0xFFFFFFFFFFFFFFFF for word in range(RANK_MASK_WORDS) + ] + return torch.tensor(words, dtype=torch.uint64, device="cpu") + + +def workspace_size_per_rank( + ep_size: int, + max_num_tokens: int, + dispatch_bytes_per_token: int, + combine_bytes_per_token: int, + eplb_stats_num_experts: int = 0, +) -> int: + """Workspace bytes per rank for rounds of up to ``max_num_tokens`` tokens per + rank whose dispatch payloads total ``dispatch_bytes_per_token`` and whose + combine payload is ``combine_bytes_per_token``. ``eplb_stats_num_experts`` + reserves the gathered EPLB statistics region (0 disables it).""" + aux_data_size = current_module().moe_a2a_get_aux_data_size( + ep_size, max_num_tokens, eplb_stats_num_experts + ) + + def pad_up(x: int) -> int: + return (x + 127) // 128 * 128 + + # Each region is padded to 128 bytes, as the module lays them out. + return ( + pad_up(aux_data_size) + + pad_up(ep_size * max_num_tokens * dispatch_bytes_per_token) + + pad_up(ep_size * max_num_tokens * combine_bytes_per_token) + ) + + +def initialize( + workspace: torch.Tensor, + ep_rank: int, + ep_size: int, + max_num_tokens: int, + eplb_stats_num_experts: int = 0, +) -> torch.Tensor: + """Clear this rank's slice of the ``[ep_size, bytes_per_rank]`` workspace + and return its metainfo tensor. + + Call it once per workspace allocation; no peer may dispatch into the + workspace until every rank has initialized its slice. + ``eplb_stats_num_experts`` must match the ``eplb_local_stats`` length later + passed to :func:`dispatch` (0 disables EPLB statistics). + """ + return current_module().moe_a2a_initialize( + workspace, ep_rank, ep_size, max_num_tokens, eplb_stats_num_experts + ) + + +def wrap_workspace( + workspace: torch.Tensor, + leading_shape: Sequence[int], + slice_start: int, + slice_end: int, + dtype: torch.dtype, +) -> torch.Tensor: + """View the workspace bytes ``[slice_start, slice_end)`` as a ``dtype`` + tensor of shape ``leading_shape + [-1]``. + + The offsets count from the start of the ``[ep_size, bytes_per_rank]`` (or + ``[bytes_per_rank]``) workspace and must fall within a single rank. + """ + if workspace.ndim == 1: + workspace = workspace.unsqueeze(0) + workspace_base = workspace.view(dtype=torch.uint8) + assert workspace.ndim == 2, "workspace must be shape [ep_size, size_per_rank]" + assert slice_end - slice_start <= workspace_base.shape[1], ( + "slice_end - slice_start must belong to a single rank" + ) + slice_rank = slice_start // workspace_base.stride(0) + local_slice_start = slice_start % workspace_base.stride(0) + local_slice_end = local_slice_start + slice_end - slice_start + assert local_slice_end <= workspace_base.shape[1], ( + "slice must fall within the workspace size per rank" + ) + return ( + workspace_base[slice_rank, local_slice_start:local_slice_end] + .view(dtype=dtype) + .view(*leading_shape, -1) + ) + + +def dispatch( + token_selected_experts: torch.Tensor, + input_payloads: list[torch.Tensor], + workspace: torch.Tensor, + metainfo: torch.Tensor, + runtime_max_tokens_per_rank: int, + ep_rank: int, + ep_size: int, + top_k: int, + num_experts: int, + enable_pdl: Optional[bool] = None, + eplb_local_stats: Optional[torch.Tensor] = None, + enable_rank_mask: bool = False, + active_rank_mask: Optional[torch.Tensor] = None, + *, + recv_view_cache: Optional[dict] = None, +) -> tuple[list[torch.Tensor], int, Optional[torch.Tensor]]: + """Send each token's payloads to the ranks that own its experts. + + ``token_selected_experts`` is the ``[local_num_tokens, top_k]`` ``int32`` + routing and ``input_payloads`` the ``[local_num_tokens, *]`` tensors to + send. ``runtime_max_tokens_per_rank`` bounds every rank's token count in + this round and must not exceed the ``max_num_tokens`` of + :func:`initialize`. ``enable_pdl=None`` uses programmatic dependent launch + where the device supports it. + + ``eplb_local_stats``, an ``[eplb_stats_num_experts]`` ``int32`` tensor, is + all-gathered across the ranks during the dispatch. + + ``enable_rank_mask`` selects the kernel variant that honors + ``active_rank_mask`` (see :func:`active_rank_mask`); without it the mask + checks are compiled out and the mask must be omitted. Tokens routed to a + masked-off rank are dropped instead of hanging the collective, and the + local rank's bit must be set. Masking a peer does not advance its + transport epoch: before a skipped peer rejoins, all ranks must quiesce and + re-initialize the workspace. + + ``recv_view_cache``, an initially empty dict owned by the caller, reuses + the receive views across rounds on the same workspace tensor. It holds a + reference to the workspace and is cleared when another workspace is + passed. + + Returns ``(recv_payloads, combine_payload_offset, eplb_gathered_stats)``: + one ``[ep_size, runtime_max_tokens_per_rank, *]`` workspace view per input + payload holding the rows routed to this rank, the workspace offset of the + matching :func:`combine` payload, and the ``[ep_size, + eplb_stats_num_experts]`` gathered statistics (``None`` without + ``eplb_local_stats``). + """ + if enable_pdl is None: + enable_pdl = device_support_pdl(token_selected_experts.device) + ( + recv_offsets, + recv_sizes, + combine_payload_offset, + eplb_gathered_stats_offset, + eplb_stats_num_experts, + ) = current_module().moe_a2a_dispatch( + token_selected_experts, + input_payloads, + workspace, + metainfo, + runtime_max_tokens_per_rank, + ep_rank, + ep_size, + top_k, + num_experts, + enable_pdl, + eplb_local_stats, + enable_rank_mask, + active_rank_mask, + ) + + if ( + recv_view_cache is not None + and recv_view_cache.get("_workspace") is not workspace + ): + recv_view_cache.clear() + recv_view_cache["_workspace"] = workspace + + recv_payloads = [] + for input_payload, offset, size in zip( + input_payloads, recv_offsets, recv_sizes, strict=True + ): + # The offsets index the whole workspace. ``int()`` keeps the key hashing + # by value whatever integer type the FFI returns. + key = ( + ep_size, + runtime_max_tokens_per_rank, + int(offset), + int(size), + input_payload.dtype, + ) + view = None if recv_view_cache is None else recv_view_cache.get(key) + if view is None: + view = wrap_workspace( + workspace, + [ep_size, runtime_max_tokens_per_rank], + offset, + offset + size, + input_payload.dtype, + ) + if recv_view_cache is not None: + recv_view_cache[key] = view + recv_payloads.append(view) + + eplb_gathered_stats = None + if eplb_gathered_stats_offset >= 0: + eplb_gathered_stats = wrap_workspace( + workspace, + [ep_size], + eplb_gathered_stats_offset, + eplb_gathered_stats_offset + ep_size * eplb_stats_num_experts * 4, + torch.int32, + ) + return recv_payloads, combine_payload_offset, eplb_gathered_stats + + +def combine( + payload: torch.Tensor, + local_num_tokens: int, + workspace: torch.Tensor, + metainfo: torch.Tensor, + runtime_max_tokens_per_rank: int, + ep_rank: int, + ep_size: int, + top_k: int, + combine_payload_offset: int, + payload_in_workspace: bool = False, + output_dtype: Optional[torch.dtype] = None, + output_scales: Optional[torch.Tensor] = None, + output_scalar_scale: float = 1.0, + sf_layout: SfLayout = SfLayout.layout_linear, + output: Optional[torch.Tensor] = None, + *, + use_low_precision: bool = False, + enable_pdl: Optional[bool] = None, + enable_rank_mask: bool = False, + active_rank_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Return each source token's expert outputs, summed over its top-k slots. + + ``payload`` is the ``[ep_size, runtime_max_tokens_per_rank, *]`` expert + output of the preceding :func:`dispatch` round on this rank, whose + ``combine_payload_offset`` and ``runtime_max_tokens_per_rank`` it takes. + ``payload_in_workspace`` marks a payload already written to that offset, + e.g. through :func:`wrap_workspace`, which skips the staging copy. + + ``output_dtype`` quantizes the ``[local_num_tokens, *]`` result: bfloat16, + ``float8_e4m3fn`` or packed FP4 ``uint8``. Quantized outputs write their + block scales to the contiguous ``output_scales``: UE8M0 ``uint8`` scales + per 32 elements for MXFP8 / MXFP4, UE4M3 ``float8_e4m3fn`` scales per 16 + for NVFP4, laid out by ``sf_layout`` including its padding. NVFP4 applies + ``output_scalar_scale`` before block scaling. ``use_low_precision`` + transfers the payload as FP8 (e4m3) and accumulates into a bfloat16 + result. ``output``, a contiguous CUDA tensor of the result's shape and + dtype, receives the result in place. + + ``active_rank_mask`` must match the one of the dispatch. + """ + if enable_pdl is None: + enable_pdl = device_support_pdl(payload.device) + if output is not None: + if not output.is_cuda: + raise ValueError( + f"output must be a CUDA tensor, got device={output.device}" + ) + if not output.is_contiguous(): + raise ValueError(f"output must be contiguous, got stride={output.stride()}") + module = current_module() + args = ( + payload, + local_num_tokens, + workspace, + metainfo, + runtime_max_tokens_per_rank, + ep_rank, + ep_size, + top_k, + combine_payload_offset, + payload_in_workspace, + output_dtype, + output_scales, + output_scalar_scale, + sf_layout.value, + use_low_precision, + enable_pdl, + enable_rank_mask, + active_rank_mask, + ) + if output is None: + return module.moe_a2a_combine(*args) + module.moe_a2a_combine_into(*args, output) + return output + + +def sanitize_expert_ids( + expert_ids: torch.Tensor, + workspace: torch.Tensor, + metainfo: torch.Tensor, + ep_rank: int, + invalid_expert_id: int, + enable_pdl: Optional[bool] = None, +) -> None: + """Set the expert ids of the received rows beyond each source rank's token + count to ``invalid_expert_id``, in place; ``expert_ids`` is the expert-id + receive payload of the preceding :func:`dispatch`.""" + if enable_pdl is None: + enable_pdl = device_support_pdl(expert_ids.device) + current_module().moe_a2a_sanitize_expert_ids( + expert_ids, workspace, metainfo, ep_rank, invalid_expert_id, enable_pdl + ) diff --git a/flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/communication.py b/flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/communication.py index 626c687579b..5b6d5131cda 100644 --- a/flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/communication.py +++ b/flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/communication.py @@ -76,19 +76,17 @@ def _acquire_workspaces(self, mapping: "Mapping") -> dict[str, Any]: import torch from ......comm.mnnvl import MnnvlMemory - from ......comm.trtllm_alltoall import ( - get_moe_commworkspace_size_per_rank, - get_moe_prepare_workspace_size_per_rank, + from .kernels import ( + comm_workspace_size_per_rank, + prepare_workspace_size_per_rank, ) key = (self.ep_rank, self.ep_size) workspaces = self._WORKSPACES.get(key) if workspaces is None: - memory = MnnvlMemory( - mapping, get_moe_commworkspace_size_per_rank(self.ep_size) - ) + memory = MnnvlMemory(mapping, comm_workspace_size_per_rank(self.ep_size)) prepare_memory = MnnvlMemory( - mapping, get_moe_prepare_workspace_size_per_rank(self.ep_size) + mapping, prepare_workspace_size_per_rank(self.ep_size) ) workspaces = { "memory": memory, @@ -113,11 +111,11 @@ def _alltoallv( """Move the rows of ``x`` selected by the send indices to the receive indices of ``output`` on the peer ranks; rows nobody sends to keep their contents.""" - from ......comm.trtllm_alltoall import moe_comm + from .kernels import alltoallv if x.dim() != 2: raise ValueError(f"expected a 2D tensor, got shape {tuple(x.shape)}") - moe_comm( + alltoallv( x, send_cumsum, send_indices, @@ -142,7 +140,7 @@ def dispatch( ) -> MoEEpDispatchResult: import torch - from ......comm.trtllm_alltoall import moe_prepare + from .kernels import prepare if self._workspace is None: raise RuntimeError("NVLinkTwoSidedAlltoAll has been destroyed") @@ -173,7 +171,7 @@ def dispatch( recv_indices, backward_recv_indices, gathered_stats, - ) = moe_prepare( + ) = prepare( topk_ids, topk_weights, eplb_local_stats, diff --git a/flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/kernels.py b/flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/kernels.py new file mode 100644 index 00000000000..4d2fd29503c --- /dev/null +++ b/flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/kernels.py @@ -0,0 +1,175 @@ +"""Kernel-level operations of the NVLink two-sided MoE all-to-all. + +These wrap the JIT module behind :class:`flashinfer.moe_ep.NVLinkTwoSidedAlltoAll`. +:func:`prepare` exchanges the routing and computes the send / receive +indices of a round over the prepare workspace; :func:`alltoallv` then moves +rows along those indices through the FIFO channels of the communication +workspace. +""" + +from __future__ import annotations + +import functools +from typing import Any, Optional + +import torch + + +@functools.cache +def get_nvlink_two_sided_module() -> Any: + """Build (on first use) and load the JIT module backing ``NVLinkTwoSidedAlltoAll``.""" + from ......jit.comm import gen_moe_ep_nvlink_two_sided_module + + return gen_moe_ep_nvlink_two_sided_module().build_and_load() + + +def set_max_usable_sm_count(max_sm_count: int) -> None: + """Cap the SMs :func:`alltoallv` spreads its FIFO channels over, e.g. so + that several ranks can share one GPU. Must precede the first + :func:`comm_workspace_size_per_rank` or :func:`alltoallv`, which fix the + channel count.""" + get_nvlink_two_sided_module().set_moe_max_usable_sm_count(max_sm_count) + + +def comm_workspace_size_per_rank(ep_size: int) -> int: + """Bytes per rank of the FIFO workspace used by :func:`alltoallv`.""" + return get_nvlink_two_sided_module().get_moe_commworkspace_size_per_rank(ep_size) + + +def prepare_workspace_size_per_rank(ep_size: int) -> int: + """Bytes per rank of the workspace used by :func:`prepare`.""" + return get_nvlink_two_sided_module().get_moe_prepare_workspace_size_per_rank( + ep_size + ) + + +def alltoallv( + input: torch.Tensor, + send_rank_cum_sum: torch.Tensor, + send_indices: torch.Tensor, + output: torch.Tensor, + recv_rank_cum_sum: torch.Tensor, + recv_indices: torch.Tensor, + workspace: torch.Tensor, + ep_rank: int, + ep_size: int, +) -> None: + """Send the rows of the 2D ``input`` listed in ``send_indices`` to the + peer ranks and write the rows received from them to the rows of + ``output`` listed in ``recv_indices``. + + The ``int32`` cumulative per-rank counts split the index lists by peer. + ``workspace`` is the ``[ep_size, words_per_rank]`` ``uint64`` FIFO + workspace of :func:`comm_workspace_size_per_rank` bytes per rank. + """ + get_nvlink_two_sided_module().moe_comm( + input, + send_rank_cum_sum, + send_indices, + output, + recv_rank_cum_sum, + recv_indices, + workspace, + ep_rank, + ep_size, + ) + + +def prepare( + expert_ids: torch.Tensor, + scales: Optional[torch.Tensor], + expert_statistics: Optional[torch.Tensor], + workspace: torch.Tensor, + max_token_count_per_rank: int, + ep_rank: int, + ep_size: int, + expert_count: int, + slot_count: int, + top_k: int, +) -> tuple[ + torch.Tensor, + Optional[torch.Tensor], + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + Optional[torch.Tensor], +]: + """Exchange the ``[num_tokens, top_k]`` ``int32`` routing ``expert_ids``, + their optional FP32 ``scales`` and the optional per-expert + ``expert_statistics`` across the ranks, and compute the indices of the + round. + + ``workspace`` is the ``[ep_size, words_per_rank]`` ``uint64`` workspace of + :func:`prepare_workspace_size_per_rank` bytes per rank. ``expert_count`` + and ``slot_count`` must be divisible by 4. + + Returns ``(recv_expert_ids, recv_scales, send_rank_cum_sum, send_indices, + recv_rank_cum_sum, recv_indices, backward_recv_indices, + gathered_statistics)``. The received routing has ``ep_size * + max_token_count_per_rank`` rank-major rows; rows past a source rank's + token count carry ``slot_count`` as their expert id. The index tensors + drive :func:`alltoallv` forward (send / recv) and back (recv / + backward recv), and ``gathered_statistics`` is ``[ep_size, + expert_count]``. + """ + attrs = {"dtype": torch.int32, "device": expert_ids.device} + recv_rows = max_token_count_per_rank * ep_size + send_rows = max_token_count_per_rank * max(ep_size, top_k) + + recv_expert_ids = torch.empty((recv_rows, top_k), **attrs) + send_rank_cum_sum = torch.empty((ep_size,), **attrs) + recv_rank_cum_sum = torch.empty((ep_size,), **attrs) + gather_recv_indices = torch.empty((recv_rows,), **attrs) + recv_indices = torch.empty((recv_rows,), **attrs) + gather_backward_recv_indices = torch.empty((send_rows,), **attrs) + backward_recv_indices = torch.empty((send_rows,), **attrs) + gather_send_indices = torch.empty((send_rows,), **attrs) + send_indices = torch.empty((send_rows,), **attrs) + recv_scales = ( + None + if scales is None + else torch.empty( + (recv_rows, top_k), dtype=torch.float32, device=expert_ids.device + ) + ) + gathered_statistics = ( + None + if expert_statistics is None + else torch.empty((ep_size, expert_count), **attrs) + ) + + get_nvlink_two_sided_module().moe_prepare( + expert_ids, + scales, + expert_statistics, + workspace, + recv_expert_ids, + send_rank_cum_sum, + recv_rank_cum_sum, + gather_recv_indices, + recv_indices, + gather_backward_recv_indices, + backward_recv_indices, + gather_send_indices, + send_indices, + recv_scales, + gathered_statistics, + max_token_count_per_rank, + ep_rank, + ep_size, + expert_count, + slot_count, + top_k, + ) + return ( + recv_expert_ids, + recv_scales, + send_rank_cum_sum, + gather_send_indices, + recv_rank_cum_sum, + gather_recv_indices, + gather_backward_recv_indices, + gathered_statistics, + ) diff --git a/scripts/task_test_single_node_comm_kernels.sh b/scripts/task_test_single_node_comm_kernels.sh index fbdad0398ec..23ad0e9b617 100644 --- a/scripts/task_test_single_node_comm_kernels.sh +++ b/scripts/task_test_single_node_comm_kernels.sh @@ -56,7 +56,7 @@ pytest -s tests/comm/test_trtllm_allreduce_fusion.py pytest -s tests/moe/test_trtllm_cutlass_fused_moe.py pytest -s tests/comm/test_trtllm_moe_allreduce_fusion.py pytest -s tests/comm/test_trtllm_moe_allreduce_fusion_finalize.py -pytest -s tests/comm/test_trtllm_moe_alltoall.py +pytest -s tests/comm/test_cake_moe_alltoall.py # nvshmem ar pytest -s tests/comm/test_nvshmem.py pytest -s tests/comm/test_nvshmem_allreduce.py diff --git a/tests/comm/moe_a2a_driver.py b/tests/comm/moe_a2a_driver.py index f4c1afa115d..015ac16a607 100644 --- a/tests/comm/moe_a2a_driver.py +++ b/tests/comm/moe_a2a_driver.py @@ -1,10 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 -"""Test driver for the functional MoE all-to-all ops of flashinfer.comm. +"""Test driver for the kernel-level operations of the Cake MoE all-to-all. ``MoeA2ADriver`` allocates and initializes an MNNVL workspace, shared by -drivers with the same geometry and backend, and runs ``moe_a2a_dispatch``, -``moe_a2a_sanitize_expert_ids`` and ``moe_a2a_combine`` on it, carrying the +drivers with the same geometry, and runs the ``dispatch``, +``sanitize_expert_ids`` and ``combine`` operations of +:mod:`flashinfer.moe_ep.backends.split.comm.cake.kernels` on it, carrying the per-round state that combine needs from dispatch. """ @@ -14,18 +15,12 @@ from flashinfer.comm.mapping import Mapping from flashinfer.comm.mnnvl import MnnvlConfig, MnnvlMemory -from flashinfer.comm.trtllm_moe_alltoall import ( - MoeAlltoAllBackend, - get_moe_alltoall_module, - moe_a2a_combine, - moe_a2a_dispatch, - moe_a2a_get_workspace_size_per_rank, - moe_a2a_initialize, - moe_a2a_sanitize_expert_ids, - moe_a2a_wrap_payload_tensor_in_workspace, -) +from flashinfer.moe_ep.backends.split.comm.cake import kernels as cake_kernels +from flashinfer.moe_ep.backends.split.comm.cake.kernels import metainfo_index from flashinfer.tllm_enums import SfLayout +__all__ = ["MoeA2ADriver", "metainfo_index", "workspace_bytes_per_rank"] + def workspace_bytes_per_rank( ep_size: int, @@ -34,33 +29,21 @@ def workspace_bytes_per_rank( hidden_size: int, extra_payload_bytes_per_token: int = 0, eplb_stats_num_experts: int = 0, - *, - backend: MoeAlltoAllBackend = "cake", ) -> int: """Workspace bytes per rank for 16-bit hidden states, int32 expert ids and FP32 weights, plus ``extra_payload_bytes_per_token``.""" - return moe_a2a_get_workspace_size_per_rank( + return cake_kernels.workspace_size_per_rank( ep_size, max_num_tokens, hidden_size * 2 + top_k * 4 + top_k * 4 + extra_payload_bytes_per_token, hidden_size * 2, eplb_stats_num_experts, - backend=backend, ) -def metainfo_index(backend: MoeAlltoAllBackend = "cake") -> dict: - """Metainfo field indices, keyed by name without the ``MOE_A2A_`` prefix.""" - names, values = get_moe_alltoall_module(backend).moe_a2a_get_metainfo_index_pairs() - return { - name.removeprefix("MOE_A2A_"): int(value) - for name, value in zip(names, values, strict=True) - } - - class MoeA2ADriver: # Workspaces keyed by (bytes per rank, ep_rank, ep_size, max tokens, - # EPLB experts, backend). + # EPLB experts). _WORKSPACES: dict = {} def __init__( @@ -73,8 +56,6 @@ def __init__( mnnvl_config: Optional[MnnvlConfig] = None, eplb_stats_num_experts: int = 0, enable_rank_mask: bool = False, - *, - backend: MoeAlltoAllBackend = "cake", ): MnnvlMemory.initialize() if mnnvl_config: @@ -85,25 +66,22 @@ def __init__( self.top_k = top_k self.num_experts = num_experts self.enable_rank_mask = enable_rank_mask - self.backend = backend key = ( workspace_size_per_rank, self.ep_rank, self.ep_size, max_num_tokens, eplb_stats_num_experts, - backend, ) if key not in self._WORKSPACES: mnnvl_mem = MnnvlMemory(mapping, workspace_size_per_rank) workspace = mnnvl_mem.as_torch_strided_tensor(torch.uint8) - metainfo = moe_a2a_initialize( + metainfo = cake_kernels.initialize( workspace, self.ep_rank, self.ep_size, max_num_tokens, eplb_stats_num_experts, - backend=backend, ) # No peer may publish into this workspace until every rank has # cleared its own slice. @@ -125,7 +103,7 @@ def dispatch( active_rank_mask: Optional[torch.Tensor] = None, ) -> list: assert self._round is None, "dispatch called twice without combine" - recv, combine_offset, self.eplb_gathered_stats = moe_a2a_dispatch( + recv, combine_offset, self.eplb_gathered_stats = cake_kernels.dispatch( token_selected_experts, input_payloads, self.workspace, @@ -138,17 +116,15 @@ def dispatch( eplb_local_stats=eplb_local_stats, enable_rank_mask=self.enable_rank_mask, active_rank_mask=active_rank_mask, - backend=self.backend, recv_view_cache=self._recv_view_cache, ) if invalid_token_expert_id is not None: - moe_a2a_sanitize_expert_ids( + cake_kernels.sanitize_expert_ids( recv[expert_id_payload_index], self.workspace, self.metainfo, self.ep_rank, invalid_token_expert_id, - backend=self.backend, ) self._round = { "local_num_tokens": token_selected_experts.size(0), @@ -165,7 +141,7 @@ def get_combine_payload_tensor_in_workspace( start + self.ep_size * runtime_max_tokens_per_rank * hidden_size * dtype.itemsize ) - return moe_a2a_wrap_payload_tensor_in_workspace( + return cake_kernels.wrap_workspace( self.workspace[self.ep_rank, :], [self.ep_size, runtime_max_tokens_per_rank], start, @@ -188,7 +164,7 @@ def combine( active_rank_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: assert self._round is not None, "combine called before dispatch" - combined = moe_a2a_combine( + combined = cake_kernels.combine( payload, self._round["local_num_tokens"], self.workspace, @@ -207,7 +183,6 @@ def combine( use_low_precision=use_low_precision, enable_rank_mask=self.enable_rank_mask, active_rank_mask=active_rank_mask, - backend=self.backend, ) self._round = None self.eplb_gathered_stats = None diff --git a/tests/comm/test_trtllm_moe_alltoall.py b/tests/comm/test_cake_moe_alltoall.py similarity index 95% rename from tests/comm/test_trtllm_moe_alltoall.py rename to tests/comm/test_cake_moe_alltoall.py index 3dd11d614d4..14c41493c11 100644 --- a/tests/comm/test_trtllm_moe_alltoall.py +++ b/tests/comm/test_cake_moe_alltoall.py @@ -27,7 +27,7 @@ from flashinfer.fp4_quantization import e2m1_and_ufp8sf_scale_to_float from tests.utils_fp8 import mxfp8_quantize_reference -import flashinfer.comm.trtllm_moe_alltoall as trtllm_moe_alltoall +from flashinfer.moe_ep.backends.split.comm.cake import kernels as cake_kernels from .conftest import mnnvl_available @@ -217,7 +217,7 @@ def test_moe_combine_rejects_invalid_output_scales(invalid_kind, error_match): output_scales = torch.empty(scale_extent + 1, dtype=torch.uint8, device="cuda") output = torch.empty((num_tokens, output_width), dtype=output_dtype, device="cuda") with pytest.raises(Exception, match=error_match): - trtllm_moe_alltoall.moe_a2a_combine( + cake_kernels.combine( combine_payload, num_tokens, workspaces, @@ -262,14 +262,14 @@ def make_payload(num_tokens, vector_dim, dtype): def test_moe_a2a_combine_rejects_cpu_output_before_jit(monkeypatch): monkeypatch.setattr( - trtllm_moe_alltoall, - "get_moe_alltoall_module", + cake_kernels, + "current_module", lambda: pytest.fail("invalid output should be rejected before JIT loading"), ) payload = torch.empty((1, 1, 1)) with pytest.raises(ValueError, match="output must be a CUDA tensor"): - trtllm_moe_alltoall.moe_a2a_combine( + cake_kernels.combine( payload, 1, torch.empty(0), @@ -286,15 +286,15 @@ def test_moe_a2a_combine_rejects_cpu_output_before_jit(monkeypatch): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_moe_a2a_combine_rejects_noncontiguous_output_before_jit(monkeypatch): monkeypatch.setattr( - trtllm_moe_alltoall, - "get_moe_alltoall_module", + cake_kernels, + "current_module", lambda: pytest.fail("invalid output should be rejected before JIT loading"), ) payload = torch.empty((1, 1, 2), device="cuda") output = torch.empty((2, 2), device="cuda").T with pytest.raises(ValueError, match="output must be contiguous"): - trtllm_moe_alltoall.moe_a2a_combine( + cake_kernels.combine( payload, 1, torch.empty(0, device="cuda"), @@ -344,7 +344,7 @@ def test_moe_alltoall_single_gpu( payload_size_per_token = sum([x[0].numel() * x.itemsize for x in input_tensors]) - workspace_size = trtllm_moe_alltoall.moe_a2a_get_workspace_size_per_rank( + workspace_size = cake_kernels.workspace_size_per_rank( 1, num_tokens, payload_size_per_token, @@ -353,9 +353,9 @@ def test_moe_alltoall_single_gpu( workspace = torch.zeros( 1, workspace_size, dtype=torch.uint8, device=torch.device("cuda") ) - metainfo = trtllm_moe_alltoall.moe_a2a_initialize(workspace, 0, 1, num_tokens) + metainfo = cake_kernels.initialize(workspace, 0, 1, num_tokens) - output_tensors, combine_payload_offset, _ = trtllm_moe_alltoall.moe_a2a_dispatch( + output_tensors, combine_payload_offset, _ = cake_kernels.dispatch( token_selected_experts, input_tensors, workspace, @@ -367,9 +367,7 @@ def test_moe_alltoall_single_gpu( num_experts=num_experts, ) # Tokens assigned to invalid experts are set to -3. - trtllm_moe_alltoall.moe_a2a_sanitize_expert_ids( - output_tensors[2], workspace, metainfo, 0, -3 - ) + cake_kernels.sanitize_expert_ids(output_tensors[2], workspace, metainfo, 0, -3) # Sort to undo the shuffling that happens in the dispatch kernel. for input_tensor, output_tensor in zip(input_tensors, output_tensors, strict=True): @@ -379,7 +377,7 @@ def test_moe_alltoall_single_gpu( if payload_in_workspace: hidden_state = input_tensors[hidden_state_index] - combine_tensor = trtllm_moe_alltoall.moe_a2a_wrap_payload_tensor_in_workspace( + combine_tensor = cake_kernels.wrap_workspace( workspace[0, :], [1, num_tokens], combine_payload_offset, @@ -396,7 +394,7 @@ def test_moe_alltoall_single_gpu( if use_output_buffer else None ) - output = trtllm_moe_alltoall.moe_a2a_combine( + output = cake_kernels.combine( combine_tensor, num_tokens, workspace, @@ -438,7 +436,7 @@ def dispatch_from_single_rank( * input_tensors[hidden_state_index].itemsize ) - workspace_size = trtllm_moe_alltoall.moe_a2a_get_workspace_size_per_rank( + workspace_size = cake_kernels.workspace_size_per_rank( world_size, num_tokens * world_size, total_payload_size_per_element, @@ -453,7 +451,7 @@ def dispatch_from_single_rank( metainfo = [] for rank in range(world_size): metainfo.append( - trtllm_moe_alltoall.moe_a2a_initialize( + cake_kernels.initialize( all_workspaces, rank, world_size, @@ -476,7 +474,7 @@ def dispatch_from_single_rank( rank_token_selected_experts = token_selected_experts[ rank * num_tokens : (rank + 1) * num_tokens ] - output, offset, _ = trtllm_moe_alltoall.moe_a2a_dispatch( + output, offset, _ = cake_kernels.dispatch( rank_token_selected_experts, rank_payloads, all_workspaces, @@ -509,7 +507,7 @@ def sanitize_expert_ids_from_single_rank( enable_pdl=True, ): for rank in range(world_size): - trtllm_moe_alltoall.moe_a2a_sanitize_expert_ids( + cake_kernels.sanitize_expert_ids( output_tensors[rank][expert_ids_index], all_workspaces, metainfo[rank], @@ -560,7 +558,7 @@ def combine_from_single_rank( cuda_streams_all_ranks = [torch.cuda.Stream() for _ in range(world_size)] for rank in range(world_size): with torch.cuda.stream(cuda_streams_all_ranks[rank]): - result = trtllm_moe_alltoall.moe_a2a_combine( + result = cake_kernels.combine( combine_payload[rank], num_tokens, all_workspaces, @@ -630,7 +628,7 @@ def test_moe_combine_k4_matches_source_reduction_tree(dtype, scalar_values): physical = torch.tensor(scalar_values, dtype=dtype, device="cuda") combine_payload = [] for rank in range(world_size): - payload = trtllm_moe_alltoall.moe_a2a_wrap_payload_tensor_in_workspace( + payload = cake_kernels.wrap_workspace( all_workspaces[rank, :], [world_size, num_tokens], combine_payload_offsets[rank], @@ -692,7 +690,7 @@ def test_moe_combine_rejects_unsupported_top_k(): enable_pdl=False, ) with pytest.raises(Exception, match="[Uu]nsupported top_k"): - trtllm_moe_alltoall.moe_a2a_combine( + cake_kernels.combine( received[0][0].clone(), num_tokens, workspace, @@ -841,7 +839,7 @@ def test_moe_alltoall_active_rank_mask_single_gpu( num_experts = world_size # top_k=1, one expert per rank: expert_id == target rank active_ranks = [r for r in range(world_size) if r != dead_rank] - active_mask = trtllm_moe_alltoall.moe_a2a_active_rank_mask(active_ranks, world_size) + active_mask = cake_kernels.active_rank_mask(active_ranks, world_size) total_tokens = num_tokens * world_size token_selected_experts = make_random_topk_ids( @@ -850,7 +848,7 @@ def test_moe_alltoall_active_rank_mask_single_gpu( payload = make_payload(total_tokens, vector_dim, torch.bfloat16) total_payload_size = payload[0].numel() * payload.itemsize - workspace_size = trtllm_moe_alltoall.moe_a2a_get_workspace_size_per_rank( + workspace_size = cake_kernels.workspace_size_per_rank( world_size, total_tokens, total_payload_size, vector_dim * payload.itemsize ) all_workspaces = torch.zeros( @@ -861,7 +859,7 @@ def test_moe_alltoall_active_rank_mask_single_gpu( # a crashed rank; only active_ranks participate. metainfo = {} for rank in active_ranks: - metainfo[rank] = trtllm_moe_alltoall.moe_a2a_initialize( + metainfo[rank] = cake_kernels.initialize( all_workspaces, rank, world_size, total_tokens ) torch.cuda.synchronize() @@ -875,7 +873,7 @@ def test_moe_alltoall_active_rank_mask_single_gpu( rank_experts = token_selected_experts[ rank * num_tokens : (rank + 1) * num_tokens ] - out, offset, _ = trtllm_moe_alltoall.moe_a2a_dispatch( + out, offset, _ = cake_kernels.dispatch( rank_experts, [rank_payload], all_workspaces, @@ -916,7 +914,7 @@ def test_moe_alltoall_active_rank_mask_single_gpu( world_size, num_tokens, vector_dim, dtype=torch.bfloat16, device="cuda" ) combine_payload.copy_(output_tensors[rank][0]) - combine_results[rank] = trtllm_moe_alltoall.moe_a2a_combine( + combine_results[rank] = cake_kernels.combine( combine_payload, num_tokens, all_workspaces, @@ -955,22 +953,22 @@ def test_moe_alltoall_active_rank_mask_rejects_self_inactive(): world_size = 4 num_tokens = 8 local_rank = 2 - mask = trtllm_moe_alltoall.moe_a2a_active_rank_mask( + mask = cake_kernels.active_rank_mask( [r for r in range(world_size) if r != local_rank], world_size ) - workspace_size = trtllm_moe_alltoall.moe_a2a_get_workspace_size_per_rank( + workspace_size = cake_kernels.workspace_size_per_rank( world_size, num_tokens * world_size, 4 * 2, 0 ) ws = torch.zeros(world_size, workspace_size, dtype=torch.uint8, device="cuda") - metainfo = trtllm_moe_alltoall.moe_a2a_initialize( + metainfo = cake_kernels.initialize( ws, local_rank, world_size, num_tokens * world_size ) experts = torch.zeros(num_tokens, 1, dtype=torch.int32, device="cuda") payload = torch.randn(num_tokens, 4, dtype=torch.bfloat16, device="cuda") with pytest.raises(Exception, match="active_rank_mask"): - trtllm_moe_alltoall.moe_a2a_dispatch( + cake_kernels.dispatch( experts, [payload], ws, @@ -992,20 +990,20 @@ def test_moe_alltoall_active_rank_mask_requires_enable_rank_mask(): world_size = 4 num_tokens = 8 local_rank = 0 - mask = trtllm_moe_alltoall.moe_a2a_active_rank_mask(range(world_size), world_size) + mask = cake_kernels.active_rank_mask(range(world_size), world_size) - workspace_size = trtllm_moe_alltoall.moe_a2a_get_workspace_size_per_rank( + workspace_size = cake_kernels.workspace_size_per_rank( world_size, num_tokens * world_size, 4 * 2, 0 ) ws = torch.zeros(world_size, workspace_size, dtype=torch.uint8, device="cuda") - metainfo = trtllm_moe_alltoall.moe_a2a_initialize( + metainfo = cake_kernels.initialize( ws, local_rank, world_size, num_tokens * world_size ) experts = torch.zeros(num_tokens, 1, dtype=torch.int32, device="cuda") payload = torch.randn(num_tokens, 4, dtype=torch.bfloat16, device="cuda") with pytest.raises(Exception, match="enable_rank_mask"): - trtllm_moe_alltoall.moe_a2a_dispatch( + cake_kernels.dispatch( experts, [payload], ws, @@ -1021,7 +1019,7 @@ def test_moe_alltoall_active_rank_mask_requires_enable_rank_mask(): # enable_rank_mask=True without a mask should also raise (no silent all-active default). with pytest.raises(Exception, match="active_rank_mask"): - trtllm_moe_alltoall.moe_a2a_dispatch( + cake_kernels.dispatch( experts, [payload], ws, @@ -1071,7 +1069,7 @@ def test_moe_alltoall_eplb_stats_single_gpu( ] total_payload_size_per_element = payload[0].numel() * payload.itemsize - workspace_size = trtllm_moe_alltoall.moe_a2a_get_workspace_size_per_rank( + workspace_size = cake_kernels.workspace_size_per_rank( world_size, num_tokens * world_size, total_payload_size_per_element, @@ -1085,7 +1083,7 @@ def test_moe_alltoall_eplb_stats_single_gpu( metainfo = [] for rank in range(world_size): metainfo.append( - trtllm_moe_alltoall.moe_a2a_initialize( + cake_kernels.initialize( all_workspaces, rank, world_size, @@ -1103,7 +1101,7 @@ def test_moe_alltoall_eplb_stats_single_gpu( rank_experts = token_selected_experts[ rank * num_tokens : (rank + 1) * num_tokens ] - _, _, eplb_gathered = trtllm_moe_alltoall.moe_a2a_dispatch( + _, _, eplb_gathered = cake_kernels.dispatch( rank_experts, [rank_payload], all_workspaces, @@ -1441,7 +1439,7 @@ def test_moe_combine_multi_rank_single_gpu( for rank in range(world_size): if payload_in_workspace: inplace_combine_tensors.append( - trtllm_moe_alltoall.moe_a2a_wrap_payload_tensor_in_workspace( + cake_kernels.wrap_workspace( all_workspaces[rank, :], [world_size, num_tokens], combine_payload_offsets[rank], @@ -1684,7 +1682,7 @@ def test_moe_combine_fp8( for rank in range(world_size): if payload_in_workspace: inplace_combine_tensors.append( - trtllm_moe_alltoall.moe_a2a_wrap_payload_tensor_in_workspace( + cake_kernels.wrap_workspace( all_workspaces[rank, :], [world_size, num_tokens], combine_payload_offsets[rank], @@ -1747,7 +1745,7 @@ def test_moe_workspace_size_per_rank(): num_tokens = 10 hidden_size = 128 topk = 2 - raw_workspace_size = trtllm_moe_alltoall.moe_a2a_get_workspace_size_per_rank( + raw_workspace_size = cake_kernels.workspace_size_per_rank( ep_size, num_tokens, (hidden_size * torch.bfloat16.itemsize + topk * 4 + topk * 4), @@ -1755,19 +1753,17 @@ def test_moe_workspace_size_per_rank(): ) assert raw_workspace_size > 0 - empty_workspace_size = trtllm_moe_alltoall.moe_a2a_get_workspace_size_per_rank( + empty_workspace_size = cake_kernels.workspace_size_per_rank( ep_size, num_tokens, 0, 0 ) assert empty_workspace_size > 0 assert ( empty_workspace_size - == trtllm_moe_alltoall.get_moe_alltoall_module().moe_a2a_get_aux_data_size( - ep_size, num_tokens - ) + == cake_kernels.current_module().moe_a2a_get_aux_data_size(ep_size, num_tokens) ) - non_empty_workspace_size = trtllm_moe_alltoall.moe_a2a_get_workspace_size_per_rank( + non_empty_workspace_size = cake_kernels.workspace_size_per_rank( ep_size, num_tokens, hidden_size, hidden_size ) diff --git a/tests/comm/test_mnnvl_memory.py b/tests/comm/test_mnnvl_memory.py index d263fa6348c..f941acdebfb 100644 --- a/tests/comm/test_mnnvl_memory.py +++ b/tests/comm/test_mnnvl_memory.py @@ -20,9 +20,9 @@ from flashinfer.comm.mapping import Mapping from flashinfer.comm.mnnvl import MnnvlMemory, MpiComm -from flashinfer.comm.trtllm_alltoall import ( - get_moe_commworkspace_size_per_rank, - moe_comm, +from flashinfer.moe_ep.backends.split.comm.nvlink_two_sided.kernels import ( + alltoallv, + comm_workspace_size_per_rank, ) from .conftest import mnnvl_available @@ -247,7 +247,7 @@ def test_moe_alltoall_multi_rank_single_gpu(self): recv_ids_all_ranks.append(local_recv_ids) workspace_memory = MnnvlMemory( - self.mapping, get_moe_commworkspace_size_per_rank(self.world_size) + self.mapping, comm_workspace_size_per_rank(self.world_size) ) alltoall_workspace = workspace_memory.as_torch_strided_tensor(torch.uint64) output = torch.empty( @@ -259,7 +259,7 @@ def test_moe_alltoall_multi_rank_single_gpu(self): self.comm.Barrier() - moe_comm( + alltoallv( input_tensors_all_ranks[self.rank], send_cumsum_all_ranks[self.rank], send_ids_all_ranks[self.rank], diff --git a/tests/comm/test_mnnvl_moe_alltoall_fused_adapter.py b/tests/comm/test_mnnvl_moe_alltoall_fused_adapter.py index c6d9631f275..598fdd4a9a2 100644 --- a/tests/comm/test_mnnvl_moe_alltoall_fused_adapter.py +++ b/tests/comm/test_mnnvl_moe_alltoall_fused_adapter.py @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 -import inspect from pathlib import Path from types import SimpleNamespace @@ -10,9 +9,9 @@ import torch from flashinfer import mxfp4_quantize, nvfp4_quantize -from flashinfer.comm import moe_a2a_active_rank_mask from flashinfer.comm.mapping import Mapping from flashinfer.comm.mnnvl import MnnvlMemory +from flashinfer.moe_ep.backends.split.comm.cake.kernels import active_rank_mask from flashinfer.tllm_enums import SfLayout from tests.utils_fp8 import mxfp8_quantize_reference @@ -20,192 +19,46 @@ from .moe_a2a_driver import MoeA2ADriver, workspace_bytes_per_rank -def test_fused_module_keeps_the_public_python_contract(): +def test_cake_ops_are_not_public_comm_api(): import flashinfer.comm as public_api - import flashinfer.comm.trtllm_moe_alltoall as api - - expected_module_parameters = { - "moe_a2a_active_rank_mask": ("active_ranks", "ep_size"), - "moe_a2a_initialize": ( - "workspace", - "ep_rank", - "ep_size", - "max_num_tokens", - "eplb_stats_num_experts", - "backend", - ), - "moe_a2a_wrap_payload_tensor_in_workspace": ( - "workspace", - "leading_shape", - "slice_start", - "slice_end", - "dtype", - ), - "moe_a2a_dispatch": ( - "token_selected_experts", - "input_payloads", - "workspace", - "metainfo", - "runtime_max_tokens_per_rank", - "ep_rank", - "ep_size", - "top_k", - "num_experts", - "enable_pdl", - "eplb_local_stats", - "enable_rank_mask", - "active_rank_mask", - "backend", - "recv_view_cache", - ), - "moe_a2a_combine": ( - "payload", - "local_num_tokens", - "workspace", - "metainfo", - "runtime_max_tokens_per_rank", - "ep_rank", - "ep_size", - "top_k", - "combine_payload_offset", - "payload_in_workspace", - "output_dtype", - "output_scales", - "output_scalar_scale", - "sf_layout", - "output", - "use_low_precision", - "enable_pdl", - "enable_rank_mask", - "active_rank_mask", - "backend", - ), - "moe_a2a_sanitize_expert_ids": ( - "expert_ids", - "workspace", - "metainfo", - "ep_rank", - "invalid_expert_id", - "enable_pdl", - "backend", - ), - "moe_a2a_get_workspace_size_per_rank": ( - "ep_size", - "max_num_tokens", - "total_dispatch_payload_size_per_token", - "combine_payload_size_per_token", - "eplb_stats_num_experts", - "backend", - ), - } - expected_module_defaults = { - "moe_a2a_active_rank_mask": {}, - "moe_a2a_initialize": {"eplb_stats_num_experts": 0, "backend": "cake"}, - "moe_a2a_wrap_payload_tensor_in_workspace": {}, - "moe_a2a_dispatch": { - "enable_pdl": None, - "eplb_local_stats": None, - "enable_rank_mask": False, - "active_rank_mask": None, - "backend": "cake", - "recv_view_cache": None, - }, - "moe_a2a_combine": { - "payload_in_workspace": False, - "output_dtype": None, - "output_scales": None, - "output_scalar_scale": 1.0, - "sf_layout": SfLayout.layout_linear, - "output": None, - "use_low_precision": False, - "enable_pdl": None, - "enable_rank_mask": False, - "active_rank_mask": None, - "backend": "cake", - }, - "moe_a2a_sanitize_expert_ids": {"enable_pdl": None, "backend": "cake"}, - "moe_a2a_get_workspace_size_per_rank": { - "eplb_stats_num_experts": 0, - "backend": "cake", - }, - } - assert set(api.__all__) == set(expected_module_parameters) - for name, expected_parameters in expected_module_parameters.items(): - assert getattr(public_api, name) is getattr(api, name) - signature = inspect.signature(getattr(api, name)) - assert tuple(signature.parameters) == expected_parameters - assert { - parameter_name: parameter.default - for parameter_name, parameter in signature.parameters.items() - if parameter.default is not inspect.Parameter.empty - } == expected_module_defaults[name] - if "backend" in signature.parameters: - assert ( - signature.parameters["backend"].kind is inspect.Parameter.KEYWORD_ONLY - ) - combine = inspect.signature(api.moe_a2a_combine) - assert ( - combine.parameters["use_low_precision"].kind is inspect.Parameter.KEYWORD_ONLY - ) + assert not [name for name in dir(public_api) if name.startswith("moe_a2a_")] assert not hasattr(public_api, "MoeAlltoAll") def test_active_rank_mask_preserves_upper_u64_bits_in_dispatch(): - mask = moe_a2a_active_rank_mask((0, 31, 32, 63), 64) + mask = active_rank_mask((0, 31, 32, 63), 64) assert mask.dtype is torch.uint64 assert mask.device.type == "cpu" assert mask.tolist() == [0x8000000180000001] -def test_backends_register_distinct_custom_ops(monkeypatch): - import flashinfer.comm.trtllm_moe_alltoall as api - - registrations = [] - - def register(name, *, mutates_args): - def decorate(function): - registrations.append((name, tuple(mutates_args), function.__name__)) - return function +def test_each_target_loads_its_own_module(monkeypatch): + import flashinfer.jit.comm as jit_comm + from flashinfer.moe_ep.backends.split.comm.cake import kernels - return decorate + built = [] class FakeSpec: + def __init__(self, target): + self.target = target + def build_and_load(self): - return SimpleNamespace() - - api.get_moe_alltoall_module.cache_clear() - monkeypatch.setattr(api, "register_custom_op", register) - monkeypatch.setattr(api, "gen_moe_alltoall_module", lambda target: FakeSpec()) - operation_names = { - "moe_a2a_initialize", - "moe_a2a_dispatch", - "moe_a2a_combine", - "moe_a2a_combine_into", - "moe_a2a_sanitize_expert_ids", - "moe_a2a_get_metainfo_index_pairs", - "moe_a2a_get_aux_data_size", - } + built.append(self.target) + return SimpleNamespace(target=self.target) + + kernels.get_cake_module.cache_clear() + monkeypatch.setattr(jit_comm, "gen_moe_ep_cake_module", FakeSpec) try: - expected_registrations = set() - for target in ("sm100a", "sm103a"): - module = api._get_moe_alltoall_module_for_target(target) - assert api._get_moe_alltoall_module_for_target(target) is module - assert set(vars(module)) == operation_names - suffix = f"_{target}" - expected_registrations.update( - f"flashinfer::{name}{suffix}" for name in operation_names - ) - combine_into = next( - row - for row in registrations - if row[0] == f"flashinfer::moe_a2a_combine_into{suffix}" - ) - assert combine_into[1] == ("workspace", "output") - assert {name for name, _, _ in registrations} == expected_registrations - assert len(registrations) == len(expected_registrations) + modules = { + target: kernels.get_cake_module(target) for target in ("sm100a", "sm103a") + } + for target, module in modules.items(): + assert module.target == target + assert kernels.get_cake_module(target) is module + assert built == ["sm100a", "sm103a"] finally: - api.get_moe_alltoall_module.cache_clear() + kernels.get_cake_module.cache_clear() @pytest.mark.parametrize( @@ -213,13 +66,13 @@ def build_and_load(self): [ ( "sm100a", - "mnnvl_moe_alltoall_sm100a", + "moe_ep_cake_alltoall_sm100a", "mnnvl_moe_alltoall_sm100.cu", "compute_100a,code=sm_100a", ), ( "sm103a", - "mnnvl_moe_alltoall_sm103a", + "moe_ep_cake_alltoall_sm103a", "mnnvl_moe_alltoall_sm103.cu", "compute_103a,code=sm_103a", ), @@ -237,7 +90,7 @@ def capture(name, sources, **kwargs): return object() monkeypatch.setattr(jit_comm, "gen_jit_spec", capture) - jit_comm.gen_moe_alltoall_module(target) + jit_comm.gen_moe_ep_cake_module(target) names = {path.name for path in captured["sources"]} assert captured["name"] == module_name assert generated_name in names @@ -265,48 +118,31 @@ def capture(name, sources, **kwargs): def test_runtime_module_selection_follows_the_device( monkeypatch, capability, cake_target ): - import flashinfer.comm.trtllm_moe_alltoall as api + from flashinfer.moe_ep.backends.split.comm.cake import kernels - monkeypatch.setattr(api.torch.cuda, "current_device", lambda: 7) - selected = [] + monkeypatch.setattr(kernels.torch.cuda, "current_device", lambda: 7) capability_queries = [] def get_device_capability(device): capability_queries.append(device) return capability - def select_target(target): - selected.append(target) - return target - - select_target.cache_clear = lambda: None - monkeypatch.setattr(api, "_get_moe_alltoall_module_for_target", select_target) - monkeypatch.setattr(api.torch.cuda, "get_device_capability", get_device_capability) + monkeypatch.setattr(kernels, "get_cake_module", lambda target: target) + monkeypatch.setattr( + kernels.torch.cuda, "get_device_capability", get_device_capability + ) try: - api.get_moe_alltoall_module.cache_clear() - with pytest.raises(ValueError, match="NVLinkOneSidedAlltoAll"): - api.get_moe_alltoall_module(backend="trtllm") - assert capability_queries == [] + kernels._device_target.cache_clear() if cake_target is None: with pytest.raises(ValueError, match="compute capability"): - api.get_moe_alltoall_module() - assert selected == [] + kernels.current_module() else: - assert api.get_moe_alltoall_module() == cake_target - assert api.get_moe_alltoall_module(backend="cake") == cake_target - assert selected == [cake_target, cake_target] + assert kernels.current_module() == cake_target + assert kernels.current_module() == cake_target assert capability_queries == [7] finally: - api.get_moe_alltoall_module.cache_clear() - - -@pytest.mark.parametrize("backend", ["auto", "trtllm", "legacy", "sm100a", "", None]) -def test_runtime_module_selection_rejects_invalid_backends(backend): - import flashinfer.comm.trtllm_moe_alltoall as api - - with pytest.raises(ValueError, match="backend"): - api.get_moe_alltoall_module(backend=backend) + kernels._device_target.cache_clear() def test_aot_registers_each_exact_mnnvl_moe_target(monkeypatch): @@ -321,7 +157,9 @@ def spec(name): monkeypatch.setattr(aot, "gen_spdlog_module", lambda: spec("spdlog")) monkeypatch.setattr(aot, "gen_attention", lambda *args: ()) monkeypatch.setattr(aot, "gen_cudnn_fmha_module", lambda: spec("cudnn")) - monkeypatch.setattr(jit_comm, "gen_comm_alltoall_module", lambda: spec("comm")) + monkeypatch.setattr( + jit_comm, "gen_moe_ep_nvlink_two_sided_module", lambda: spec("two_sided") + ) monkeypatch.setattr(jit_comm, "gen_vllm_comm_module", lambda: spec("vllm")) monkeypatch.setattr(jit_comm, "gen_pcie_ipc_comm_module", lambda: spec("pcie")) monkeypatch.setattr( @@ -329,8 +167,8 @@ def spec(name): ) monkeypatch.setattr( jit_comm, - "gen_moe_alltoall_module", - lambda target: selected.append(target) or spec(f"mnnvl_{target}"), + "gen_moe_ep_cake_module", + lambda target: selected.append(target) or spec(f"cake_{target}"), ) aot.gen_all_modules( @@ -1027,10 +865,6 @@ def _run_public_mpi2_cycle(): if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): pytest.skip("the Cake backend requires SM100 or SM103") - from flashinfer.comm.trtllm_moe_alltoall import get_moe_alltoall_module - - assert get_moe_alltoall_module() is get_moe_alltoall_module(backend="cake") - routes = torch.tensor(_ROUTES_BY_RANK[rank], dtype=torch.int32, device="cuda") payloads = _payloads(rank, routes) max_tokens = routes.shape[0] @@ -1043,7 +877,6 @@ def _run_public_mpi2_cycle(): _HIDDEN_SIZE, extra_payload_bytes_per_token=extra_payload_bytes, eplb_stats_num_experts=5, - backend="cake", ) mapping = Mapping(rank=rank, moe_ep_size=2, tp_size=2, world_size=2) collective = MoeA2ADriver( @@ -1054,9 +887,8 @@ def _run_public_mpi2_cycle(): workspace_size_per_rank=workspace_size, eplb_stats_num_experts=5, enable_rank_mask=True, - backend="cake", ) - active_mask = moe_a2a_active_rank_mask((0, 1), 2) + active_mask = active_rank_mask((0, 1), 2) _run_public_combine_round( collective, rank, @@ -1089,7 +921,6 @@ def _run_public_mpi2_cycle(): max_tokens, _HIDDEN_SIZE, extra_payload_bytes_per_token=extra_payload_bytes, - backend="cake", ) topk8_collective = MoeA2ADriver( mapping, @@ -1098,7 +929,6 @@ def _run_public_mpi2_cycle(): num_experts=16, workspace_size_per_rank=topk8_workspace_size, enable_rank_mask=True, - backend="cake", ) _run_public_combine_round( topk8_collective, @@ -1164,7 +994,7 @@ def test_public_mpi2_nondivisible_six_payload_eplb_and_external_output(): if rank == 0: routes = torch.tensor(_ROUTES_BY_RANK[rank], dtype=torch.int32, device="cuda") payloads = _payloads(rank, routes) - rank_zero_mask = moe_a2a_active_rank_mask((0,), 2) + rank_zero_mask = active_rank_mask((0,), 2) for collective in (first[0],): _run_public_combine_round( collective, diff --git a/tests/comm/test_moe_a2a_recv_view_cache.py b/tests/comm/test_moe_a2a_recv_view_cache.py index 21f26eaaa01..f9230dfea21 100644 --- a/tests/comm/test_moe_a2a_recv_view_cache.py +++ b/tests/comm/test_moe_a2a_recv_view_cache.py @@ -12,17 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""``moe_a2a_dispatch``'s optional receive-view cache. +"""The optional receive-view cache of the Cake kernels' ``dispatch``. The cache is a pure host-side optimization: it skips rebuilding the workspace-backed views on every dispatch. These tests pin the three properties that make that safe, and they run without a GPU or an MPI harness by stubbing the FFI module -- the real -``tests/comm/test_trtllm_moe_alltoall.py`` path needs multi-rank NVLink hardware. +``tests/comm/test_cake_moe_alltoall.py`` path needs multi-rank NVLink hardware. """ import torch -import flashinfer.comm.trtllm_moe_alltoall as a2a +from flashinfer.moe_ep.backends.split.comm.cake import kernels as a2a _RECV_OFFSETS = [0, 512] @@ -41,7 +41,7 @@ def moe_a2a_dispatch(*args, **kwargs): def _dispatch(workspace, cache, wrap_calls): payloads = [torch.zeros(2, 4), torch.zeros(2, 2)] - return a2a.moe_a2a_dispatch( + return a2a.dispatch( torch.zeros(2, 1, dtype=torch.int32), payloads, workspace, @@ -57,13 +57,13 @@ def _dispatch(workspace, cache, wrap_calls): def _patch(monkeypatch, wrap_calls): - monkeypatch.setattr(a2a, "get_moe_alltoall_module", lambda *a, **k: _FakeModule) + monkeypatch.setattr(a2a, "current_module", lambda: _FakeModule) def fake_wrap(workspace, leading_shape, slice_start, slice_end, dtype): wrap_calls.append((int(slice_start), int(slice_end), dtype)) return torch.zeros(1) - monkeypatch.setattr(a2a, "moe_a2a_wrap_payload_tensor_in_workspace", fake_wrap) + monkeypatch.setattr(a2a, "wrap_workspace", fake_wrap) def test_cache_none_rebuilds_every_call(monkeypatch): @@ -125,7 +125,7 @@ class _IdentityInt(int): def test_cache_key_survives_an_identity_hashing_ffi_int(monkeypatch): """The key must hash by value even if the FFI hands back a wrapper. - ``moe_a2a_dispatch`` returns ``tvm::ffi::Array``, and each read boxes a + The module's ``moe_a2a_dispatch`` returns ``tvm::ffi::Array``, and each read boxes a fresh object (``a[0] is a[0]`` is False outside CPython's small-int intern range). Value-based hashing is what makes that harmless. Were the elements ever boxed in something that hashed by identity, every lookup would miss: @@ -148,7 +148,7 @@ def moe_a2a_dispatch(*args, **kwargs): 0, ) - monkeypatch.setattr(a2a, "get_moe_alltoall_module", lambda *a, **k: _WrapperModule) + monkeypatch.setattr(a2a, "current_module", lambda: _WrapperModule) ws = torch.zeros(4) cache = {} diff --git a/tests/comm/test_trtllm_alltoall.py b/tests/comm/test_nvlink_two_sided_alltoall.py similarity index 75% rename from tests/comm/test_trtllm_alltoall.py rename to tests/comm/test_nvlink_two_sided_alltoall.py index d275d0de256..efff933217b 100644 --- a/tests/comm/test_trtllm_alltoall.py +++ b/tests/comm/test_nvlink_two_sided_alltoall.py @@ -17,7 +17,7 @@ import pytest import torch -import flashinfer.comm.trtllm_alltoall as tllm_alltoall +from flashinfer.moe_ep.backends.split.comm.nvlink_two_sided import kernels from flashinfer.utils import get_compute_capability # Skip all tests on SM110 (Thor) devices - these tests hang indefinitely on this architecture @@ -38,7 +38,7 @@ def setup_test_environment(): # Set up SM count once for all tests sm_count = torch.cuda.get_device_properties(0).multi_processor_count max_sm_count = sm_count // 8 # Maximum world size is 8 - tllm_alltoall.set_moe_max_usable_sm_count(max_sm_count) + kernels.set_max_usable_sm_count(max_sm_count) has_setup_max_sm_count = True torch.manual_seed(0x1234) @@ -60,20 +60,6 @@ def setup_test_environment(): (8, 16384, 128, torch.float16), # Many small vectors, 8 ranks ] -PREPARE_INDICES_PARAMS = [ - (0, 8, 256, 4, 3, False), # Rank 0, small config - (1, 8, 256, 4, 3, True), # Rank 1, small config with real cumsum - (7, 8, 256, 8, 1025, False), # High rank, medium config - (7, 64, 1024, 32, 1029, True), # High rank, large config with real cumsum -] - -LOCAL_GATHER_PARAMS = [ - (0, 8, 256, 4, 3), # Rank 0, small config - (7, 8, 256, 8, 32), # High rank, medium config - (7, 64, 1024, 32, 1029), # High rank, large config -] - - # Real cross-GPU communication test parameters CROSS_GPU_PARAMS = [ (2, 100, 256, torch.float16), # 2 GPUs, 2 ranks @@ -139,12 +125,12 @@ def test_moe_alltoall_single_gpu( ) ref_output_tensor[recv_indices] = input_tensor[send_indices] - workspace_size = tllm_alltoall.get_moe_commworkspace_size_per_rank(1) + workspace_size = kernels.comm_workspace_size_per_rank(1) all_workspaces = torch.zeros( 1, workspace_size, dtype=torch.uint64, device=torch.device("cuda") ) - tllm_alltoall.moe_comm( + kernels.alltoallv( input_tensor, send_cumsum, send_indices, @@ -288,7 +274,7 @@ def test_moe_alltoall_multi_rank_single_gpu( cuda_streams_all_ranks = [torch.cuda.Stream() for _ in range(world_size)] - workspace_size = tllm_alltoall.get_moe_commworkspace_size_per_rank(world_size) + workspace_size = kernels.comm_workspace_size_per_rank(world_size) all_workspaces = torch.zeros( world_size, workspace_size, dtype=torch.uint64, device=torch.device("cuda") ) @@ -299,7 +285,7 @@ def test_moe_alltoall_multi_rank_single_gpu( # do alltoall in parallel for rank in range(world_size): with torch.cuda.stream(cuda_streams_all_ranks[rank]): - tllm_alltoall.moe_comm( + kernels.alltoallv( input_tensors_all_ranks[rank], send_cumsum_all_ranks[rank], send_ids_all_ranks[rank], @@ -316,208 +302,6 @@ def test_moe_alltoall_multi_rank_single_gpu( torch.testing.assert_close(output_tensor, ref_output_tensor, atol=1e-5, rtol=1e-5) -@pytest.mark.parametrize( - "ep_rank,ep_size,expert_count,top_k,max_token_count_per_rank,use_real_rank_token_count_cumsum", - PREPARE_INDICES_PARAMS, -) -def test_moe_alltoall_prepare_indices( - ep_rank, - ep_size, - expert_count, - top_k, - max_token_count_per_rank, - use_real_rank_token_count_cumsum, -): - """Test MOE alltoall prepare indices functionality.""" - torch.cuda.set_device(0) - - def generate_references(): - rank_token_count = max_token_count_per_rank - if use_real_rank_token_count_cumsum: - # Make sure we have at least 1 token in each rank except last rank - rank_token_counts = [ - max(1, torch.randint(1, max_token_count_per_rank + 1, (1,)).item()) - for _ in range(ep_size - 1) - ] - rank_token_counts.append( - max_token_count_per_rank - ) # last rank has max tokens - real_rank_token_count_cumsum = ( - torch.tensor( - rank_token_counts, dtype=torch.int32, device=torch.device("cuda") - ) - .cumsum(dim=0) - .to(torch.int32) - ) - rank_token_count = rank_token_counts[ep_rank] - else: - real_rank_token_count_cumsum = None - - # Generate target rank ids for this rank - target_rank_ids = torch.randint( - 0, - ep_size, - (rank_token_count, top_k), - dtype=torch.int32, - device=torch.device("cuda"), - ) - - if not use_real_rank_token_count_cumsum: - gathered_target_rank_ids = torch.zeros( - ep_size * max_token_count_per_rank, - top_k, - dtype=torch.int32, - device=torch.device("cuda"), - ) - gathered_target_rank_ids[ - ep_rank * max_token_count_per_rank : ep_rank * max_token_count_per_rank - + rank_token_count - ] = target_rank_ids - else: - total_tokens = real_rank_token_count_cumsum[-1].item() - gathered_target_rank_ids = torch.zeros( - total_tokens, top_k, dtype=torch.int32, device=torch.device("cuda") - ) - start_pos = ( - 0 if ep_rank == 0 else real_rank_token_count_cumsum[ep_rank - 1].item() - ) - gathered_target_rank_ids[start_pos : start_pos + rank_token_count] = ( - target_rank_ids - ) - - return gathered_target_rank_ids, real_rank_token_count_cumsum, target_rank_ids - - gathered_target_rank_ids, real_rank_token_count_cumsum, target_rank_ids = ( - generate_references() - ) - - ( - local_gather_indices, - send_rank_count_cumsum, - send_rank_local_indices, - recv_rank_count_cumsum, - recv_rank_local_indices, - backward_recv_rank_local_indices, - ) = tllm_alltoall.moe_comm_prepare_indices( - gathered_target_rank_ids, - real_rank_token_count_cumsum, - max_token_count_per_rank, - expert_count, - top_k, - ep_rank, - ep_size, - ) - - # Validate shapes - assert local_gather_indices.shape[0] <= max_token_count_per_rank * ep_size - assert send_rank_count_cumsum.shape[0] == ep_size - assert recv_rank_count_cumsum.shape[0] == ep_size - assert send_rank_local_indices.shape[0] <= max_token_count_per_rank * max( - ep_size, top_k - ) - assert recv_rank_local_indices.shape[0] <= max_token_count_per_rank * ep_size - assert backward_recv_rank_local_indices.shape[0] <= max_token_count_per_rank * max( - ep_size, top_k - ) - - # Basic validation - cumulative sums should be non-decreasing - assert torch.all(send_rank_count_cumsum[1:] >= send_rank_count_cumsum[:-1]) - assert torch.all(recv_rank_count_cumsum[1:] >= recv_rank_count_cumsum[:-1]) - - -@pytest.mark.parametrize( - "ep_rank,ep_size,expert_count,top_k,max_token_count_per_rank", LOCAL_GATHER_PARAMS -) -def test_moe_local_gather( - ep_rank, - ep_size, - expert_count, - top_k, - max_token_count_per_rank, -): - """Test MOE local gather functionality.""" - torch.cuda.set_device(0) - - # Generate test data using the original method - rank_token_count_cumsum = torch.randint( - 0, - max_token_count_per_rank + 1, - (ep_size,), - dtype=torch.int32, - device=torch.device("cuda"), - ) - rank_token_count_cumsum = torch.cumsum(rank_token_count_cumsum, dim=0).to( - torch.int32 - ) - local_token_count = rank_token_count_cumsum[ep_size - 1].cpu().item() - local_max_token_count = max_token_count_per_rank * ep_size - local_gather_indices = torch.randint( - 0, - max_token_count_per_rank * ep_size, - (local_max_token_count,), - dtype=torch.int32, - device=torch.device("cuda"), - ) - - gathered_expert_ids = torch.randint( - 0, - expert_count, - (max_token_count_per_rank * ep_size, top_k), - dtype=torch.int32, - device=torch.device("cuda"), - ) - gathered_scales = torch.rand( - (max_token_count_per_rank * ep_size, top_k), - dtype=torch.float32, - device=torch.device("cuda"), - ) - - ref_local_expert_ids = torch.zeros( - local_max_token_count, top_k, dtype=torch.int32, device=torch.device("cuda") - ) - ref_local_scales = torch.zeros( - local_max_token_count, - top_k, - dtype=torch.float32, - device=torch.device("cuda"), - ) - - # compute reference - ref_local_expert_ids += expert_count - valid_local_gather_indices = local_gather_indices[:local_token_count] - ref_local_expert_ids[:local_token_count] = gathered_expert_ids[ - valid_local_gather_indices - ] - ref_local_scales[:local_token_count] = gathered_scales[valid_local_gather_indices] - - local_expert_ids = torch.empty( - local_max_token_count, top_k, dtype=torch.int32, device=torch.device("cuda") - ) - local_scales = torch.empty( - local_max_token_count, - top_k, - dtype=torch.float32, - device=torch.device("cuda"), - ) - - tllm_alltoall.moe_local_gather( - rank_token_count_cumsum, - local_gather_indices, - gathered_expert_ids, - gathered_scales, - local_expert_ids, - local_scales, - max_token_count_per_rank, - expert_count, - top_k, - ep_rank, - ep_size, - ) - - assert torch.equal(local_expert_ids, ref_local_expert_ids) - assert torch.equal(local_scales, ref_local_scales) - - @pytest.mark.parametrize( "ep_rank, ep_size, expert_count, slot_count, top_k, max_token_count_per_rank", [ @@ -703,7 +487,7 @@ def generate_references(): cuda_streams_all_ranks = [torch.cuda.Stream() for _ in range(ep_size)] - workspace_size = tllm_alltoall.get_moe_prepare_workspace_size_per_rank(ep_size) + workspace_size = kernels.prepare_workspace_size_per_rank(ep_size) all_workspaces = torch.zeros( ep_size, workspace_size, dtype=torch.uint64, device=torch.device("cuda") @@ -713,7 +497,7 @@ def generate_references(): cur_stream = torch.cuda.current_stream() stream.wait_stream(cur_stream) with torch.cuda.stream(stream): - tllm_alltoall.moe_prepare( + kernels.prepare( expert_ids_all_ranks[0], scales_all_ranks[0], experter_count_lists[0], @@ -801,7 +585,7 @@ def generate_references(): local_recv_rank_indices, backward_local_recv_rank_indices, gathered_expert_statics, - ) = tllm_alltoall.moe_prepare( + ) = kernels.prepare( expert_ids_all_ranks[rank], scales_all_ranks[rank], experter_count_lists[rank], @@ -814,7 +598,7 @@ def generate_references(): top_k, ) else: - tllm_alltoall.moe_prepare( + kernels.prepare( expert_ids_all_ranks[rank], scales_all_ranks[rank], experter_count_lists[rank], diff --git a/tests/moe_ep/test_communication.py b/tests/moe_ep/test_communication.py index fcce30accc1..f37d3a91e49 100644 --- a/tests/moe_ep/test_communication.py +++ b/tests/moe_ep/test_communication.py @@ -289,7 +289,7 @@ def test_graph_state_requires_a_capturable_backend( class _FakeCakeOps: - """Records the ``moe_a2a_*`` ops CakeAlltoAll calls; dispatch returns zeros.""" + """Records the kernel operations CakeAlltoAll calls; dispatch returns zeros.""" def __init__(self) -> None: self.calls: dict = {} @@ -315,17 +315,15 @@ def combine(self, payload, local_num_tokens, *args, **kwargs): @pytest.fixture def fake_cake(monkeypatch): import flashinfer.comm.mnnvl as mnnvl - import flashinfer.comm.trtllm_moe_alltoall as a2a + from flashinfer.moe_ep.backends.split.comm.cake import kernels as a2a from flashinfer.comm.mapping import Mapping from flashinfer.moe_ep.backends.split.comm.cake import communication as cake ops = _FakeCakeOps() - monkeypatch.setattr( - a2a, "moe_a2a_get_workspace_size_per_rank", ops.workspace_size_per_rank - ) - monkeypatch.setattr(a2a, "moe_a2a_dispatch", ops.dispatch) - monkeypatch.setattr(a2a, "moe_a2a_sanitize_expert_ids", ops.sanitize) - monkeypatch.setattr(a2a, "moe_a2a_combine", ops.combine) + monkeypatch.setattr(a2a, "workspace_size_per_rank", ops.workspace_size_per_rank) + monkeypatch.setattr(a2a, "dispatch", ops.dispatch) + monkeypatch.setattr(a2a, "sanitize_expert_ids", ops.sanitize) + monkeypatch.setattr(a2a, "combine", ops.combine) monkeypatch.setattr(mnnvl.MnnvlMemory, "initialize", staticmethod(lambda: None)) monkeypatch.setattr(torch.cuda, "synchronize", lambda *args, **kwargs: None) @@ -371,8 +369,7 @@ def test_cake_sizes_dispatch_by_format(fake_cake) -> None: _params(hidden_size=64, dispatch_format=QuantFormat.NVFP4), "cake", ) - args, kwargs = fake_cake.calls["workspace_size"] - assert kwargs["backend"] == "cake" + args, _ = fake_cake.calls["workspace_size"] _, _, dispatch_bytes, combine_bytes, _ = args # NVFP4 values and scales, then int32 ids and FP32 weights for top_k=2. assert dispatch_bytes == 32 + 4 + 2 * 4 * 2 @@ -403,12 +400,11 @@ def test_cake_payload_plumbing(fake_cake) -> None: torch.int32, torch.float32, ] - assert kwargs["backend"] == "cake" # enable_pdl=None leaves the choice to the ops (device_support_pdl). assert kwargs["enable_pdl"] is None sanitized_ids, invalid_id, kwargs = fake_cake.calls["sanitize"] assert sanitized_ids.dtype == torch.int32 and sanitized_ids.shape == (2, 2, 2) - assert invalid_id == -1 and kwargs["backend"] == "cake" + assert invalid_id == -1 assert kwargs["enable_pdl"] is None assert result.tokens_per_rank == 2 assert result.hidden_states.shape == (4, 8) @@ -425,7 +421,7 @@ def test_cake_payload_plumbing(fake_cake) -> None: # is the region pinned at allocation, not the offset dispatch returned. assert args[-2] == comm._state["combine_offset"] != 0 assert args[-1] is True - assert kwargs["use_low_precision"] is True and kwargs["backend"] == "cake" + assert kwargs["use_low_precision"] is True comm.dispatch(hidden, ids, weights) comm.combine(torch.zeros(6, 8, dtype=torch.bfloat16)) @@ -816,26 +812,28 @@ def test_nvlink_one_sided_validates_config(fake_one_sided) -> None: def test_nvlink_two_sided_marks_padding_rows_invalid(monkeypatch) -> None: import flashinfer.comm.mnnvl as mnnvl - import flashinfer.comm.trtllm_alltoall as two_sided_ops from flashinfer.comm.mapping import Mapping from flashinfer.moe_ep.backends.split.comm.nvlink_two_sided import ( communication as two_sided, ) + from flashinfer.moe_ep.backends.split.comm.nvlink_two_sided import ( + kernels as two_sided_ops, + ) calls: dict = {"comm": []} - def moe_prepare(*args): + def prepare(*args): calls["prepare"] = args recv_ids = torch.tensor([[1, 3], [4, 4], [4, 4]], dtype=torch.int32) indices = [torch.zeros(1, dtype=torch.int32) for _ in range(5)] return (recv_ids, torch.ones(3, 2), *indices, None) - def moe_comm(x, send_cumsum, send_indices, output, *args): + def alltoallv(x, send_cumsum, send_indices, output, *args): calls["comm"].append((x, output, args)) output.fill_(1.0) - monkeypatch.setattr(two_sided_ops, "moe_prepare", moe_prepare) - monkeypatch.setattr(two_sided_ops, "moe_comm", moe_comm) + monkeypatch.setattr(two_sided_ops, "prepare", prepare) + monkeypatch.setattr(two_sided_ops, "alltoallv", alltoallv) monkeypatch.setattr(mnnvl.MnnvlMemory, "initialize", staticmethod(lambda: None)) monkeypatch.setattr( NVLinkTwoSidedAlltoAll, From aeb7d1898ddbddeb1ff41e6cf0ecdea609f8893e Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Sun, 18 Oct 2026 16:22:32 +0000 Subject: [PATCH 05/17] docs(moe_ep): document the expert-parallel communication backends Add docs/api/moe_ep.rst with MoEEpCommunication, create_communication and the NVLink one-sided, Cake and two-sided backends, and move the checkpoint/restore notes there from the removed flashinfer.comm MNNVL A2A section. The architecture doc now points at the backends' own kernels and says that instances of one EP group share a workspace per layout. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- docs/api/comm.rst | 40 --------------------- docs/api/moe_ep.rst | 46 +++++++++++++++++++++++++ docs/design_docs/moe_ep_architecture.md | 9 ++--- docs/index.rst | 1 + 4 files changed, 52 insertions(+), 44 deletions(-) create mode 100644 docs/api/moe_ep.rst diff --git a/docs/api/comm.rst b/docs/api/comm.rst index 12caa76062d..34d0ccd4966 100644 --- a/docs/api/comm.rst +++ b/docs/api/comm.rst @@ -600,46 +600,6 @@ TensorRT-LLM MNNVL AllReduce trtllm_mnnvl_fused_allreduce_rmsnorm mpi_barrier -MNNVL A2A (Throughput Backend) -------------------------------- - -These are the kernel-level primitives of the NVLink one-sided MoE all-to-all. -For the dispatch/combine interface shared by all expert-parallel communication -backends, use :class:`flashinfer.moe_ep.NVLinkOneSidedAlltoAll`, which runs -newer one-sided kernels (CFT counted writes, up to 256 ranks), or -:class:`flashinfer.moe_ep.CakeAlltoAll`, which runs these ops with -``backend="cake"``. - -.. currentmodule:: flashinfer.comm - -.. autosummary:: - :toctree: ../generated - - moe_a2a_initialize - moe_a2a_dispatch - moe_a2a_combine - moe_a2a_sanitize_expert_ids - moe_a2a_get_workspace_size_per_rank - moe_a2a_wrap_payload_tensor_in_workspace - -``NVLinkOneSidedAlltoAll`` and ``CakeAlltoAll`` preserve their workspace's CUDA -virtual addresses across process checkpoint/restore. After quiescing all work, -call ``checkpoint_prepare`` to release the non-checkpointable physical MNNVL -handles, and for ``NVLinkOneSidedAlltoAll`` its CFT logical endpoints. Then -call ``checkpoint_restore`` with a fresh communication backend before replaying -a captured CUDA graph: - -.. code-block:: python - - communication.checkpoint_prepare() - communication.checkpoint_restore(comm_backend) - -Both methods are collective. Every rank must call them in the same order, and -``comm_backend`` must reproduce the original rank and world size. -Repeated calls are no-ops after the workspace reaches the requested state. -If an exception occurs after physical handle unmapping or remapping begins, -do not retry or reuse the workspace; restart the affected rank. - DCP All-to-All (Context-Parallel Attention Reduction) ----------------------------------------------------- diff --git a/docs/api/moe_ep.rst b/docs/api/moe_ep.rst new file mode 100644 index 00000000000..11a98184d8f --- /dev/null +++ b/docs/api/moe_ep.rst @@ -0,0 +1,46 @@ +.. _apimoe_ep: + +flashinfer.moe_ep +================= + +.. currentmodule:: flashinfer.moe_ep + +Expert-parallel communication +----------------------------- + +A :class:`MoEEpCommunication` dispatches each token to the ranks that own its +experts and combines the expert outputs back on the token's source rank. +:func:`create_communication` builds one from a backend config; the backends +below run over NVLink (MNNVL) symmetric memory. + +.. autosummary:: + :toctree: ../generated + + MoEEpCommunication + MoEEpCommParams + MoEEpDispatchResult + create_communication + NVLinkOneSidedAlltoAll + NVLinkOneSidedConfig + CakeAlltoAll + CakeAlltoAllConfig + NVLinkTwoSidedAlltoAll + NVLinkTwoSidedConfig + +``NVLinkOneSidedAlltoAll`` and ``CakeAlltoAll`` preserve their workspace's CUDA +virtual addresses across process checkpoint/restore. After quiescing all work, +call ``checkpoint_prepare`` to release the non-checkpointable physical MNNVL +handles, and for ``NVLinkOneSidedAlltoAll`` its CFT logical endpoints. Then +call ``checkpoint_restore`` with a fresh communication backend before replaying +a captured CUDA graph: + +.. code-block:: python + + communication.checkpoint_prepare() + communication.checkpoint_restore(comm_backend) + +Both methods are collective. Every rank must call them in the same order, and +``comm_backend`` must reproduce the original rank and world size. +Repeated calls are no-ops after the workspace reaches the requested state. +If an exception occurs after physical handle unmapping or remapping begins, +do not retry or reuse the workspace; restart the affected rank. diff --git a/docs/design_docs/moe_ep_architecture.md b/docs/design_docs/moe_ep_architecture.md index 5a1cfb1f055..c27f5c3369b 100644 --- a/docs/design_docs/moe_ep_architecture.md +++ b/docs/design_docs/moe_ep_architecture.md @@ -67,8 +67,9 @@ and the MoEEpCommunication path); a backend must be registered as exactly one kind, and both paths hand the inner kernel the same `SplitKernelContext`: - **`MoEEpCommunication`** (`core/comm/communication.py`) is a self-contained - dispatch/combine object: one long-lived instance per EP group owns its - workspace, and per-step routing state lives in that workspace. The NVLink + dispatch/combine object: long-lived instances of one EP group (e.g. one per + MoE layer) share a workspace per layout, and per-step routing state lives in + that workspace, so rounds on it must not interleave. The NVLink backends implement it. `dispatch(hidden_states, topk_ids, topk_weights, ...)` returns `ep_size * tokens_per_rank` receive rows with their GLOBAL top-k routing (rows without a token carry @@ -91,8 +92,8 @@ kind, and both paths hand the inner kernel the same `SplitKernelContext`: | Backend | Config | Interface | Transport | |---|---|---|---| | `nvlink_one_sided` | `NVLinkOneSidedConfig` | `MoEEpCommunication` | MNNVL symmetric memory; dispatch puts tokens into peers' receive buffers, combine gets results back; optional CFT counted writes over the NVLink fabric (SM100+, CUDA 13.4+, driver 615+); up to 256 ranks. Kernels follow TensorRT-LLM's (`kernels/moe/communication/`) | -| `cake` | `CakeAlltoAllConfig` | `MoEEpCommunication` | The one-sided protocol of the `flashinfer.comm` `moe_a2a_*` ops running the generated Cake kernels (`backend="cake"`); SM100/SM103 only | -| `nvlink_two_sided` | `NVLinkTwoSidedConfig` | `MoEEpCommunication` | MNNVL FIFO channels, all-to-all-v (`moe_prepare` / `moe_comm` of `flashinfer.comm.trtllm_alltoall`); `num_experts % 4 == 0` | +| `cake` | `CakeAlltoAllConfig` | `MoEEpCommunication` | One-sided put/get like `nvlink_one_sided`, running generated Cake kernels (`cake/kernels.py`, `csrc/cake_moe_alltoall*`); SM100/SM103 only | +| `nvlink_two_sided` | `NVLinkTwoSidedConfig` | `MoEEpCommunication` | MNNVL FIFO channels, all-to-all-v (`nvlink_two_sided/kernels.py`, `csrc/moe_ep_nvlink_two_sided_*`); `num_experts % 4 == 0` | | `nccl_ep` | `NcclEpConfig` | Fleet/Handle | see below | | `nixl_ep` | `NvepConfig` | Fleet/Handle | see below | diff --git a/docs/index.rst b/docs/index.rst index 406d90bd747..23b9d0e705e 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -44,6 +44,7 @@ FlashInfer is a library and kernel generator for Large Language Models that prov api/grouped_mm api/fused_moe api/megamoe + api/moe_ep api/cascade api/comm api/sparse From 604fae8f52f46cd602bb0abc65022fda50576258 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:57:34 +0000 Subject: [PATCH 06/17] test(comm): name the Cake all-to-all tests and driver after Cake test_mnnvl_moe_alltoall.py is the multi-GPU MNNVL test of the Cake kernels, the counterpart of the single-GPU test_cake_moe_alltoall.py. Rename it to test_cake_moe_alltoall_multi_gpu.py, the fused-adapter and receive-view-cache tests to the same prefix, and moe_a2a_driver.MoeA2ADriver to cake_a2a_driver.CakeA2ADriver. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- scripts/task_test_multi_node_comm_kernels.sh | 2 +- .../{moe_a2a_driver.py => cake_a2a_driver.py} | 6 +++--- ... => test_cake_moe_alltoall_fused_adapter.py} | 6 +++--- ...l.py => test_cake_moe_alltoall_multi_gpu.py} | 17 +++++++++++++---- ...> test_cake_moe_alltoall_recv_view_cache.py} | 2 +- 5 files changed, 21 insertions(+), 12 deletions(-) rename tests/comm/{moe_a2a_driver.py => cake_a2a_driver.py} (97%) rename tests/comm/{test_mnnvl_moe_alltoall_fused_adapter.py => test_cake_moe_alltoall_fused_adapter.py} (99%) rename tests/comm/{test_mnnvl_moe_alltoall.py => test_cake_moe_alltoall_multi_gpu.py} (98%) rename tests/comm/{test_moe_a2a_recv_view_cache.py => test_cake_moe_alltoall_recv_view_cache.py} (98%) diff --git a/scripts/task_test_multi_node_comm_kernels.sh b/scripts/task_test_multi_node_comm_kernels.sh index 856e693a0ba..4798ac331eb 100755 --- a/scripts/task_test_multi_node_comm_kernels.sh +++ b/scripts/task_test_multi_node_comm_kernels.sh @@ -21,7 +21,7 @@ DISABLE_SANITY_TEST=true source "${SCRIPT_DIR}/test_utils.sh" # Define the specific test files for multi-node comm tests -TEST_FILES="tests/comm/test_mnnvl_memory.py tests/comm/test_trtllm_mnnvl_allreduce.py tests/comm/test_mnnvl_moe_alltoall.py tests/comm/test_trtllm_allreduce_checkpoint.py" +TEST_FILES="tests/comm/test_mnnvl_memory.py tests/comm/test_trtllm_mnnvl_allreduce.py tests/comm/test_cake_moe_alltoall_multi_gpu.py tests/comm/test_trtllm_allreduce_checkpoint.py" # Main execution main() { diff --git a/tests/comm/moe_a2a_driver.py b/tests/comm/cake_a2a_driver.py similarity index 97% rename from tests/comm/moe_a2a_driver.py rename to tests/comm/cake_a2a_driver.py index 015ac16a607..6f3653996f3 100644 --- a/tests/comm/moe_a2a_driver.py +++ b/tests/comm/cake_a2a_driver.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Test driver for the kernel-level operations of the Cake MoE all-to-all. -``MoeA2ADriver`` allocates and initializes an MNNVL workspace, shared by +``CakeA2ADriver`` allocates and initializes an MNNVL workspace, shared by drivers with the same geometry, and runs the ``dispatch``, ``sanitize_expert_ids`` and ``combine`` operations of :mod:`flashinfer.moe_ep.backends.split.comm.cake.kernels` on it, carrying the @@ -19,7 +19,7 @@ from flashinfer.moe_ep.backends.split.comm.cake.kernels import metainfo_index from flashinfer.tllm_enums import SfLayout -__all__ = ["MoeA2ADriver", "metainfo_index", "workspace_bytes_per_rank"] +__all__ = ["CakeA2ADriver", "metainfo_index", "workspace_bytes_per_rank"] def workspace_bytes_per_rank( @@ -41,7 +41,7 @@ def workspace_bytes_per_rank( ) -class MoeA2ADriver: +class CakeA2ADriver: # Workspaces keyed by (bytes per rank, ep_rank, ep_size, max tokens, # EPLB experts). _WORKSPACES: dict = {} diff --git a/tests/comm/test_mnnvl_moe_alltoall_fused_adapter.py b/tests/comm/test_cake_moe_alltoall_fused_adapter.py similarity index 99% rename from tests/comm/test_mnnvl_moe_alltoall_fused_adapter.py rename to tests/comm/test_cake_moe_alltoall_fused_adapter.py index 598fdd4a9a2..49ea8195f96 100644 --- a/tests/comm/test_mnnvl_moe_alltoall_fused_adapter.py +++ b/tests/comm/test_cake_moe_alltoall_fused_adapter.py @@ -16,7 +16,7 @@ from tests.utils_fp8 import mxfp8_quantize_reference from .conftest import mnnvl_available -from .moe_a2a_driver import MoeA2ADriver, workspace_bytes_per_rank +from .cake_a2a_driver import CakeA2ADriver, workspace_bytes_per_rank def test_cake_ops_are_not_public_comm_api(): @@ -879,7 +879,7 @@ def _run_public_mpi2_cycle(): eplb_stats_num_experts=5, ) mapping = Mapping(rank=rank, moe_ep_size=2, tp_size=2, world_size=2) - collective = MoeA2ADriver( + collective = CakeA2ADriver( mapping, max_num_tokens=max_tokens, top_k=top_k, @@ -922,7 +922,7 @@ def _run_public_mpi2_cycle(): _HIDDEN_SIZE, extra_payload_bytes_per_token=extra_payload_bytes, ) - topk8_collective = MoeA2ADriver( + topk8_collective = CakeA2ADriver( mapping, max_num_tokens=max_tokens, top_k=8, diff --git a/tests/comm/test_mnnvl_moe_alltoall.py b/tests/comm/test_cake_moe_alltoall_multi_gpu.py similarity index 98% rename from tests/comm/test_mnnvl_moe_alltoall.py rename to tests/comm/test_cake_moe_alltoall_multi_gpu.py index c303efc6119..489e8b6dad2 100644 --- a/tests/comm/test_mnnvl_moe_alltoall.py +++ b/tests/comm/test_cake_moe_alltoall_multi_gpu.py @@ -12,6 +12,15 @@ # 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. +"""Multi-GPU tests of the Cake MoE all-to-all kernels over MNNVL memory. + +One MPI rank per GPU; the workspace is shared through ``MnnvlMemory`` and the +received payloads and combined outputs are checked against a reference. +``test_cake_moe_alltoall.py`` covers the same kernels with every rank simulated +on one GPU. + +Run with: mpirun -n 4 python -m pytest tests/comm/test_cake_moe_alltoall_multi_gpu.py +""" import traceback @@ -23,7 +32,7 @@ from flashinfer.comm.mnnvl import MnnvlMemory from .conftest import mnnvl_available -from .moe_a2a_driver import MoeA2ADriver, metainfo_index, workspace_bytes_per_rank +from .cake_a2a_driver import CakeA2ADriver, metainfo_index, workspace_bytes_per_rank class MPIExit(Exception): @@ -339,7 +348,7 @@ def run_moe_a2a_dispatch_single_rank( hidden_size, extra_payload_bytes_per_token=extra_payload_bytes, ) - moe_a2a = MoeA2ADriver( + moe_a2a = CakeA2ADriver( mapping, max_num_tokens, top_k, @@ -852,7 +861,7 @@ def moe_a2a_dispatch_moe_combine_test_impl(distribution, top_k, use_lora=False): hidden_size, extra_payload_bytes_per_token=extra_payload_bytes, ) - moe_a2a = MoeA2ADriver( + moe_a2a = CakeA2ADriver( mapping=mapping, max_num_tokens=max_num_tokens, top_k=top_k, @@ -955,5 +964,5 @@ def test_moe_a2a_dispatch_moe_combine(distribution, top_k, use_lora): if __name__ == "__main__": - # Run with: mpirun -n 4 python -m pytest tests/comm/test_mnnvl_moe_alltoall.py -v + # Run with: mpirun -n 4 python -m pytest tests/comm/test_cake_moe_alltoall_multi_gpu.py -v pytest.main([__file__, "-v", "-s"]) diff --git a/tests/comm/test_moe_a2a_recv_view_cache.py b/tests/comm/test_cake_moe_alltoall_recv_view_cache.py similarity index 98% rename from tests/comm/test_moe_a2a_recv_view_cache.py rename to tests/comm/test_cake_moe_alltoall_recv_view_cache.py index f9230dfea21..2e0b8c88132 100644 --- a/tests/comm/test_moe_a2a_recv_view_cache.py +++ b/tests/comm/test_cake_moe_alltoall_recv_view_cache.py @@ -17,7 +17,7 @@ The cache is a pure host-side optimization: it skips rebuilding the workspace-backed views on every dispatch. These tests pin the three properties that make that safe, and they run without a GPU or an MPI harness by stubbing the FFI module -- the real -``tests/comm/test_cake_moe_alltoall.py`` path needs multi-rank NVLink hardware. +``tests/comm/test_cake_moe_alltoall_multi_gpu.py`` path needs multi-rank NVLink hardware. """ import torch From bd0391b4b6d59a887cd23d7bce7ec8e367c746ef Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:08:27 +0000 Subject: [PATCH 07/17] docs(bench): say MoE communication has no flashinfer_benchmark routine yet Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- benchmarks/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index aaf48d1f502..401dbda1655 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -54,7 +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: not a routine of this script. `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. +- 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: From aa7ad45a32a6a5bf25b77803ff5afa3c66204cc1 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:12:18 +0000 Subject: [PATCH 08/17] refactor(comm): give the generated Cake MoE all-to-all sources the cake_ prefix Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- ..._alltoall_sm100.cu => cake_moe_alltoall_sm100.cu} | 0 ..._alltoall_sm103.cu => cake_moe_alltoall_sm103.cu} | 0 flashinfer/jit/comm.py | 4 ++-- tests/comm/test_cake_moe_alltoall_fused_adapter.py | 12 ++++++------ 4 files changed, 8 insertions(+), 8 deletions(-) rename csrc/generated/moe_alltoall/{mnnvl_moe_alltoall_sm100.cu => cake_moe_alltoall_sm100.cu} (100%) rename csrc/generated/moe_alltoall/{mnnvl_moe_alltoall_sm103.cu => cake_moe_alltoall_sm103.cu} (100%) diff --git a/csrc/generated/moe_alltoall/mnnvl_moe_alltoall_sm100.cu b/csrc/generated/moe_alltoall/cake_moe_alltoall_sm100.cu similarity index 100% rename from csrc/generated/moe_alltoall/mnnvl_moe_alltoall_sm100.cu rename to csrc/generated/moe_alltoall/cake_moe_alltoall_sm100.cu diff --git a/csrc/generated/moe_alltoall/mnnvl_moe_alltoall_sm103.cu b/csrc/generated/moe_alltoall/cake_moe_alltoall_sm103.cu similarity index 100% rename from csrc/generated/moe_alltoall/mnnvl_moe_alltoall_sm103.cu rename to csrc/generated/moe_alltoall/cake_moe_alltoall_sm103.cu diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index fa999d7b119..ed5f414e8a0 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -308,8 +308,8 @@ def gen_ulysses_a2a_module() -> JitSpec: CakeAlltoAllTarget = Literal["sm100a", "sm103a"] _CAKE_ALLTOALL_GENERATED_SOURCE = { - "sm100a": "mnnvl_moe_alltoall_sm100.cu", - "sm103a": "mnnvl_moe_alltoall_sm103.cu", + "sm100a": "cake_moe_alltoall_sm100.cu", + "sm103a": "cake_moe_alltoall_sm103.cu", } _CAKE_ALLTOALL_NVCC_FLAGS = { "sm100a": sm100a_nvcc_flags, diff --git a/tests/comm/test_cake_moe_alltoall_fused_adapter.py b/tests/comm/test_cake_moe_alltoall_fused_adapter.py index 49ea8195f96..2573426f937 100644 --- a/tests/comm/test_cake_moe_alltoall_fused_adapter.py +++ b/tests/comm/test_cake_moe_alltoall_fused_adapter.py @@ -67,13 +67,13 @@ def build_and_load(self): ( "sm100a", "moe_ep_cake_alltoall_sm100a", - "mnnvl_moe_alltoall_sm100.cu", + "cake_moe_alltoall_sm100.cu", "compute_100a,code=sm_100a", ), ( "sm103a", "moe_ep_cake_alltoall_sm103a", - "mnnvl_moe_alltoall_sm103.cu", + "cake_moe_alltoall_sm103.cu", "compute_103a,code=sm_103a", ), ], @@ -286,8 +286,8 @@ def test_bf16_vector_routes_and_stage_grid_keep_exact_boundaries(): generated_root = Path(__file__).resolve().parents[2] / "csrc/generated/moe_alltoall" for generated_name in ( - "mnnvl_moe_alltoall_sm100.cu", - "mnnvl_moe_alltoall_sm103.cu", + "cake_moe_alltoall_sm100.cu", + "cake_moe_alltoall_sm103.cu", ): generated_source = (generated_root / generated_name).read_text() normalized_generated_source = " ".join(generated_source.split()) @@ -484,8 +484,8 @@ def test_fused_quantized_combine_routes_mirror_the_cake_host_rule(): generated_root = repo_root / "csrc/generated/moe_alltoall" for generated_name in ( - "mnnvl_moe_alltoall_sm100.cu", - "mnnvl_moe_alltoall_sm103.cu", + "cake_moe_alltoall_sm100.cu", + "cake_moe_alltoall_sm103.cu", ): generated = " ".join((generated_root / generated_name).read_text().split()) for top_k in _FUSED_QUANTIZED_COMBINE_TOP_K: From 126a2700d56fa097229d8326a75a842b690b623e Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:31:46 +0000 Subject: [PATCH 09/17] refactor(moe_ep): name the NVLink all-to-all sources and modules after their backends csrc/moe_ep_nvlink_{one,two}_sided_* -> csrc/nvlink_{one,two}_sided_*, the JIT modules moe_ep_nvlink_{one,two}_sided_alltoall -> nvlink_{one,two}_sided_alltoall and their generators and loaders to match, aligning them with NVLinkOneSidedAlltoAll and NVLinkTwoSidedAlltoAll. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- ...ltoall.cu => nvlink_one_sided_alltoall.cu} | 0 ...ltoall.cu => nvlink_two_sided_alltoall.cu} | 34 +++++++++---------- ...oall.cuh => nvlink_two_sided_alltoall.cuh} | 4 +-- ...prepare.cu => nvlink_two_sided_prepare.cu} | 8 ++--- ...epare.cuh => nvlink_two_sided_prepare.cuh} | 4 +-- docs/design_docs/moe_ep_architecture.md | 2 +- flashinfer/aot.py | 8 ++--- flashinfer/jit/__init__.py | 4 +-- flashinfer/jit/comm.py | 14 ++++---- .../comm/nvlink_one_sided/communication.py | 4 +-- .../split/comm/nvlink_one_sided/kernels.py | 10 +++--- .../split/comm/nvlink_two_sided/kernels.py | 22 +++++++----- .../test_cake_moe_alltoall_fused_adapter.py | 2 +- tests/moe_ep/test_communication.py | 4 ++- tests/moe_ep/test_nvlink_one_sided.py | 4 +-- 15 files changed, 66 insertions(+), 58 deletions(-) rename csrc/{moe_ep_nvlink_one_sided_alltoall.cu => nvlink_one_sided_alltoall.cu} (100%) rename csrc/{moe_ep_nvlink_two_sided_alltoall.cu => nvlink_two_sided_alltoall.cu} (89%) rename csrc/{moe_ep_nvlink_two_sided_alltoall.cuh => nvlink_two_sided_alltoall.cuh} (99%) rename csrc/{moe_ep_nvlink_two_sided_prepare.cu => nvlink_two_sided_prepare.cu} (99%) rename csrc/{moe_ep_nvlink_two_sided_prepare.cuh => nvlink_two_sided_prepare.cuh} (98%) diff --git a/csrc/moe_ep_nvlink_one_sided_alltoall.cu b/csrc/nvlink_one_sided_alltoall.cu similarity index 100% rename from csrc/moe_ep_nvlink_one_sided_alltoall.cu rename to csrc/nvlink_one_sided_alltoall.cu diff --git a/csrc/moe_ep_nvlink_two_sided_alltoall.cu b/csrc/nvlink_two_sided_alltoall.cu similarity index 89% rename from csrc/moe_ep_nvlink_two_sided_alltoall.cu rename to csrc/nvlink_two_sided_alltoall.cu index 7e8e136e854..4a347d6aa86 100644 --- a/csrc/moe_ep_nvlink_two_sided_alltoall.cu +++ b/csrc/nvlink_two_sided_alltoall.cu @@ -17,11 +17,11 @@ #include -#include "moe_ep_nvlink_two_sided_alltoall.cuh" -#include "moe_ep_nvlink_two_sided_prepare.cuh" +#include "nvlink_two_sided_alltoall.cuh" +#include "nvlink_two_sided_prepare.cuh" #include "tvm_ffi_utils.h" -using namespace flashinfer::moe_ep_nvlink_two_sided; +using namespace flashinfer::nvlink_two_sided; using tvm::ffi::Optional; using tvm::ffi::Tuple; @@ -52,9 +52,9 @@ void moeCommOp(TensorView input, TensorView sendRankCumSum, TensorView sendIndic TVM_FFI_ICHECK(epRank >= 0 && epRank < epSize) << "epRank must be in the range [0, epSize)"; - flashinfer::moe_ep_nvlink_two_sided::MoeEpWorldInfo worldInfo = {static_cast(epSize), - static_cast(epRank)}; - flashinfer::moe_ep_nvlink_two_sided::SendRecvDataInfo sendRecvDataInfo; + flashinfer::nvlink_two_sided::MoeEpWorldInfo worldInfo = {static_cast(epSize), + static_cast(epRank)}; + flashinfer::nvlink_two_sided::SendRecvDataInfo sendRecvDataInfo; size_t eltSize = get_element_size(input); size_t eltCountPerU64 = sizeof(uint64_t) / eltSize; @@ -63,7 +63,7 @@ void moeCommOp(TensorView input, TensorView sendRankCumSum, TensorView sendIndic sendRecvDataInfo.vectorSizeInU64 = input.size(1) / eltCountPerU64; sendRecvDataInfo.DoPreCompute(); - flashinfer::moe_ep_nvlink_two_sided::SendRecvDispls sendDispls, recvDispls; + flashinfer::nvlink_two_sided::SendRecvDispls sendDispls, recvDispls; sendDispls.dataPtr = static_cast(input.data_ptr()); sendDispls.rankCountCumSum = static_cast(sendRankCumSum.data_ptr()); sendDispls.rankLocalIndices = static_cast(sendIndices.data_ptr()); @@ -74,24 +74,24 @@ void moeCommOp(TensorView input, TensorView sendRankCumSum, TensorView sendIndic recvDispls.rankLocalIndices = static_cast(recvIndices.data_ptr()); recvDispls.vectorStrideInU64 = output.stride(0) / eltCountPerU64; - flashinfer::moe_ep_nvlink_two_sided::MoeCommWorkspace workspace; + flashinfer::nvlink_two_sided::MoeCommWorkspace workspace; workspace.workspacePtr = static_cast(allWorkspaces.data_ptr()); workspace.rankStrideInU64 = allWorkspaces.stride(0); auto stream = get_current_stream(); - flashinfer::moe_ep_nvlink_two_sided::moeAllToAll(worldInfo, sendRecvDataInfo, sendDispls, - recvDispls, workspace, stream); + flashinfer::nvlink_two_sided::moeAllToAll(worldInfo, sendRecvDataInfo, sendDispls, recvDispls, + workspace, stream); } int64_t getWorkspaceSizePerRank(int64_t epSize) { int epSize32 = static_cast(epSize); - return flashinfer::moe_ep_nvlink_two_sided::getMoeCommWorkspaceSize(epSize32); + return flashinfer::nvlink_two_sided::getMoeCommWorkspaceSize(epSize32); } int64_t getPrepareWorkspaceSizePerRank(int64_t epSize) { int epSize32 = static_cast(epSize); - return flashinfer::moe_ep_nvlink_two_sided::moe_prepare::getMoePrepareWorkspaceSize(epSize32); + return flashinfer::nvlink_two_sided::moe_prepare::getMoePrepareWorkspaceSize(epSize32); } void moePrepareOp(TensorView expertsIds, Optional scales, @@ -183,13 +183,13 @@ void moePrepareOp(TensorView expertsIds, Optional scales, gatheredExpertStaticsPtr = static_cast(gatheredExpertStatics.value().data_ptr()); } - flashinfer::moe_ep_nvlink_two_sided::moe_prepare::MoeCommWorkspace workspace; + flashinfer::nvlink_two_sided::moe_prepare::MoeCommWorkspace workspace; workspace.workspacePtr = static_cast(allWorkspaces.data_ptr()); workspace.rankStrideInU64 = allWorkspaces.stride(0); auto stream = get_current_stream(); - flashinfer::moe_ep_nvlink_two_sided::moe_prepare::computeCountAndIndice( + flashinfer::nvlink_two_sided::moe_prepare::computeCountAndIndice( static_cast(expertsIds.data_ptr()), static_cast(sendRankCountCumSum.data_ptr()), static_cast(recvRankCountCumSum.data_ptr()), static_cast(sendRankIndices.data_ptr()), @@ -197,11 +197,11 @@ void moePrepareOp(TensorView expertsIds, Optional scales, static_cast(recvRankIndices.data_ptr()), workspace, tokenCount, maxTokenCountPerRank, topK, slotCount, epRank, epSize, stream); - flashinfer::moe_ep_nvlink_two_sided::moe_prepare::computeCumsum( + flashinfer::nvlink_two_sided::moe_prepare::computeCumsum( static_cast(sendRankCountCumSum.data_ptr()), static_cast(recvRankCountCumSum.data_ptr()), epRank, epSize, stream); - flashinfer::moe_ep_nvlink_two_sided::moe_prepare::moveIndice( + flashinfer::nvlink_two_sided::moe_prepare::moveIndice( static_cast(sendRankCountCumSum.data_ptr()), static_cast(recvRankCountCumSum.data_ptr()), static_cast(sendRankIndices.data_ptr()), @@ -212,7 +212,7 @@ void moePrepareOp(TensorView expertsIds, Optional scales, static_cast(gatherRecvRankIndices.data_ptr()), epRank, epSize, maxTokenCountPerRank, stream); - flashinfer::moe_ep_nvlink_two_sided::moe_prepare::allToAllMetadata( + flashinfer::nvlink_two_sided::moe_prepare::allToAllMetadata( static_cast(expertsIds.data_ptr()), static_cast(preparedLocalExpertIds.data_ptr()), scalesPtr, preparedLocalScalesPtr, localExpertStaticsPtr, gatheredExpertStaticsPtr, workspace, diff --git a/csrc/moe_ep_nvlink_two_sided_alltoall.cuh b/csrc/nvlink_two_sided_alltoall.cuh similarity index 99% rename from csrc/moe_ep_nvlink_two_sided_alltoall.cuh rename to csrc/nvlink_two_sided_alltoall.cuh index 412ac4fece2..8e177dea698 100644 --- a/csrc/moe_ep_nvlink_two_sided_alltoall.cuh +++ b/csrc/nvlink_two_sided_alltoall.cuh @@ -26,7 +26,7 @@ namespace flashinfer { -namespace moe_ep_nvlink_two_sided { +namespace nvlink_two_sided { #ifdef __CUDACC__ #define ALIGN_256 __align__(256) @@ -533,5 +533,5 @@ void setMaxUsableSmCount(int smCount) { AllToAllChannelCommunicatorBase::setMaxUsableSmCount(smCount); } -} // namespace moe_ep_nvlink_two_sided +} // namespace nvlink_two_sided } // namespace flashinfer diff --git a/csrc/moe_ep_nvlink_two_sided_prepare.cu b/csrc/nvlink_two_sided_prepare.cu similarity index 99% rename from csrc/moe_ep_nvlink_two_sided_prepare.cu rename to csrc/nvlink_two_sided_prepare.cu index ccfc98a411b..7daf515a0e3 100644 --- a/csrc/moe_ep_nvlink_two_sided_prepare.cu +++ b/csrc/nvlink_two_sided_prepare.cu @@ -22,9 +22,9 @@ #include "flashinfer/exception.h" #include "flashinfer/utils.cuh" -#include "moe_ep_nvlink_two_sided_prepare.cuh" +#include "nvlink_two_sided_prepare.cuh" -// Local definition to avoid multiple definition issues from moe_ep_nvlink_two_sided_alltoall.cuh +// Local definition to avoid multiple definition issues from nvlink_two_sided_alltoall.cuh static int getMultiProcessorCount() { int device_id; int multi_processor_count; @@ -36,7 +36,7 @@ static int getMultiProcessorCount() { namespace cg = cooperative_groups; -namespace flashinfer::moe_ep_nvlink_two_sided { +namespace flashinfer::nvlink_two_sided { namespace moe_prepare { @@ -620,4 +620,4 @@ size_t getMoePrepareWorkspaceSize(int epSize) { } // namespace moe_prepare -} // namespace flashinfer::moe_ep_nvlink_two_sided +} // namespace flashinfer::nvlink_two_sided diff --git a/csrc/moe_ep_nvlink_two_sided_prepare.cuh b/csrc/nvlink_two_sided_prepare.cuh similarity index 98% rename from csrc/moe_ep_nvlink_two_sided_prepare.cuh rename to csrc/nvlink_two_sided_prepare.cuh index ee0f9613198..346b1cd01ca 100644 --- a/csrc/moe_ep_nvlink_two_sided_prepare.cuh +++ b/csrc/nvlink_two_sided_prepare.cuh @@ -21,7 +21,7 @@ #define DEBUG_PIPELINE 0 -namespace flashinfer::moe_ep_nvlink_two_sided { +namespace flashinfer::nvlink_two_sided { namespace moe_prepare { @@ -125,4 +125,4 @@ size_t getMoePrepareWorkspaceSize(int epSize); } // namespace moe_prepare -} // namespace flashinfer::moe_ep_nvlink_two_sided +} // namespace flashinfer::nvlink_two_sided diff --git a/docs/design_docs/moe_ep_architecture.md b/docs/design_docs/moe_ep_architecture.md index c27f5c3369b..a6bd3202964 100644 --- a/docs/design_docs/moe_ep_architecture.md +++ b/docs/design_docs/moe_ep_architecture.md @@ -93,7 +93,7 @@ kind, and both paths hand the inner kernel the same `SplitKernelContext`: |---|---|---|---| | `nvlink_one_sided` | `NVLinkOneSidedConfig` | `MoEEpCommunication` | MNNVL symmetric memory; dispatch puts tokens into peers' receive buffers, combine gets results back; optional CFT counted writes over the NVLink fabric (SM100+, CUDA 13.4+, driver 615+); up to 256 ranks. Kernels follow TensorRT-LLM's (`kernels/moe/communication/`) | | `cake` | `CakeAlltoAllConfig` | `MoEEpCommunication` | One-sided put/get like `nvlink_one_sided`, running generated Cake kernels (`cake/kernels.py`, `csrc/cake_moe_alltoall*`); SM100/SM103 only | -| `nvlink_two_sided` | `NVLinkTwoSidedConfig` | `MoEEpCommunication` | MNNVL FIFO channels, all-to-all-v (`nvlink_two_sided/kernels.py`, `csrc/moe_ep_nvlink_two_sided_*`); `num_experts % 4 == 0` | +| `nvlink_two_sided` | `NVLinkTwoSidedConfig` | `MoEEpCommunication` | MNNVL FIFO channels, all-to-all-v (`nvlink_two_sided/kernels.py`, `csrc/nvlink_two_sided_*`); `num_experts % 4 == 0` | | `nccl_ep` | `NcclEpConfig` | Fleet/Handle | see below | | `nixl_ep` | `NvepConfig` | Fleet/Handle | see below | diff --git a/flashinfer/aot.py b/flashinfer/aot.py index 2b7f549faaf..0f40ac787db 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -1037,8 +1037,8 @@ def gen_all_modules( gen_dcp_alltoall_module, gen_dcp_lse_reduce_module, gen_moe_ep_cake_module, - gen_moe_ep_nvlink_one_sided_module, - gen_moe_ep_nvlink_two_sided_module, + gen_nvlink_one_sided_alltoall_module, + gen_nvlink_two_sided_alltoall_module, gen_pcie_ipc_ag_rs_module, gen_pcie_ipc_comm_module, gen_trtllm_comm_module, @@ -1046,7 +1046,7 @@ def gen_all_modules( gen_vllm_comm_module, ) - jit_specs.append(gen_moe_ep_nvlink_two_sided_module()) + jit_specs.append(gen_nvlink_two_sided_alltoall_module()) if ( has_sm90 or has_sm100 @@ -1092,7 +1092,7 @@ def gen_all_modules( # NVLink one-sided MoE all-to-all behind flashinfer.moe_ep. The fence path # needs SM90+; the CFT path compiles only for SM100+ and traps elsewhere. if has_sm90 or has_sm100 or has_sm100f or has_sm103 or has_sm107: - jit_specs.append(gen_moe_ep_nvlink_one_sided_module()) + jit_specs.append(gen_nvlink_one_sided_alltoall_module()) jit_specs.append(gen_vllm_comm_module()) # No architecture gate: the kernels use only plain PTX loads/stores # and CUDA IPC, and target PCIe machines without NVLink, which is diff --git a/flashinfer/jit/__init__.py b/flashinfer/jit/__init__.py index 088b2147d9e..383e78169ac 100644 --- a/flashinfer/jit/__init__.py +++ b/flashinfer/jit/__init__.py @@ -84,10 +84,10 @@ from .comm import gen_pcie_ipc_comm_module as gen_pcie_ipc_comm_module from .comm import gen_moe_ep_cake_module as gen_moe_ep_cake_module from .comm import ( - gen_moe_ep_nvlink_one_sided_module as gen_moe_ep_nvlink_one_sided_module, + gen_nvlink_one_sided_alltoall_module as gen_nvlink_one_sided_alltoall_module, ) from .comm import ( - gen_moe_ep_nvlink_two_sided_module as gen_moe_ep_nvlink_two_sided_module, + gen_nvlink_two_sided_alltoall_module as gen_nvlink_two_sided_alltoall_module, ) from .comm import gen_dcp_alltoall_module as gen_dcp_alltoall_module from .dsv3_optimizations import ( diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index ed5f414e8a0..c4e75845067 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -30,13 +30,13 @@ from typing import Any, Dict, List, Literal, Tuple -def gen_moe_ep_nvlink_two_sided_module() -> JitSpec: +def gen_nvlink_two_sided_alltoall_module() -> JitSpec: """NVLink two-sided MoE all-to-all behind ``flashinfer.moe_ep.NVLinkTwoSidedAlltoAll``.""" return gen_jit_spec( - "moe_ep_nvlink_two_sided_alltoall", + "nvlink_two_sided_alltoall", [ - jit_env.FLASHINFER_CSRC_DIR / "moe_ep_nvlink_two_sided_alltoall.cu", - jit_env.FLASHINFER_CSRC_DIR / "moe_ep_nvlink_two_sided_prepare.cu", + jit_env.FLASHINFER_CSRC_DIR / "nvlink_two_sided_alltoall.cu", + jit_env.FLASHINFER_CSRC_DIR / "nvlink_two_sided_prepare.cu", ], ) @@ -367,7 +367,7 @@ def gen_moe_ep_cake_module(target: CakeAlltoAllTarget) -> JitSpec: ) -def gen_moe_ep_nvlink_one_sided_module() -> JitSpec: +def gen_nvlink_one_sided_alltoall_module() -> JitSpec: """NVLink one-sided MoE all-to-all behind ``flashinfer.moe_ep.NVLinkOneSidedAlltoAll``. Includes the CFT counted-write path, which is compiled in when the CUDA @@ -380,9 +380,9 @@ def gen_moe_ep_nvlink_one_sided_module() -> JitSpec: # sm100f fallback; the driver rejects that fallback for these kernels on SM107. nvcc_flags = current_compilation_context.get_nvcc_flags_list() return gen_jit_spec( - "moe_ep_nvlink_one_sided_alltoall", + "nvlink_one_sided_alltoall", [ - jit_env.FLASHINFER_CSRC_DIR / "moe_ep_nvlink_one_sided_alltoall.cu", + jit_env.FLASHINFER_CSRC_DIR / "nvlink_one_sided_alltoall.cu", nv_internal / "tensorrt_llm" / "kernels" diff --git a/flashinfer/moe_ep/backends/split/comm/nvlink_one_sided/communication.py b/flashinfer/moe_ep/backends/split/comm/nvlink_one_sided/communication.py index dfd236e4e93..ee1f8315cab 100644 --- a/flashinfer/moe_ep/backends/split/comm/nvlink_one_sided/communication.py +++ b/flashinfer/moe_ep/backends/split/comm/nvlink_one_sided/communication.py @@ -33,7 +33,7 @@ unmap_mnnvl_memory, ) from .config import NVLinkOneSidedConfig -from .kernels import get_nvlink_one_sided_module, layout_constants +from .kernels import get_nvlink_one_sided_alltoall_module, layout_constants if TYPE_CHECKING: import torch @@ -166,7 +166,7 @@ def __init__( f"got {params.top_k}" ) - self._module = get_nvlink_one_sided_module() + self._module = get_nvlink_one_sided_alltoall_module() self._index = index = layout_constants() if self.ep_size > index.MAX_RANKS: raise ValueError( diff --git a/flashinfer/moe_ep/backends/split/comm/nvlink_one_sided/kernels.py b/flashinfer/moe_ep/backends/split/comm/nvlink_one_sided/kernels.py index 32fd805e83d..2d21c8ed137 100644 --- a/flashinfer/moe_ep/backends/split/comm/nvlink_one_sided/kernels.py +++ b/flashinfer/moe_ep/backends/split/comm/nvlink_one_sided/kernels.py @@ -8,11 +8,11 @@ @functools.cache -def get_nvlink_one_sided_module() -> Any: +def get_nvlink_one_sided_alltoall_module() -> Any: """Build (on first use) and load the JIT module backing ``NVLinkOneSidedAlltoAll``.""" - from ......jit.comm import gen_moe_ep_nvlink_one_sided_module + from ......jit.comm import gen_nvlink_one_sided_alltoall_module - return gen_moe_ep_nvlink_one_sided_module().build_and_load() + return gen_nvlink_one_sided_alltoall_module().build_and_load() @functools.cache @@ -22,7 +22,9 @@ def layout_constants() -> SimpleNamespace: Attribute names drop the ``MOE_A2A_`` prefix of the C++ constants, e.g. ``layout_constants().COMBINE_INPUT_OFFSET_INDEX`` or ``.MAX_RANKS``. """ - names, values = get_nvlink_one_sided_module().moe_a2a_get_metainfo_index_pairs() + names, values = ( + get_nvlink_one_sided_alltoall_module().moe_a2a_get_metainfo_index_pairs() + ) return SimpleNamespace( **{ str(name).removeprefix("MOE_A2A_"): int(value) diff --git a/flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/kernels.py b/flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/kernels.py index 4d2fd29503c..f783bb2c44c 100644 --- a/flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/kernels.py +++ b/flashinfer/moe_ep/backends/split/comm/nvlink_two_sided/kernels.py @@ -16,11 +16,11 @@ @functools.cache -def get_nvlink_two_sided_module() -> Any: +def get_nvlink_two_sided_alltoall_module() -> Any: """Build (on first use) and load the JIT module backing ``NVLinkTwoSidedAlltoAll``.""" - from ......jit.comm import gen_moe_ep_nvlink_two_sided_module + from ......jit.comm import gen_nvlink_two_sided_alltoall_module - return gen_moe_ep_nvlink_two_sided_module().build_and_load() + return gen_nvlink_two_sided_alltoall_module().build_and_load() def set_max_usable_sm_count(max_sm_count: int) -> None: @@ -28,18 +28,22 @@ def set_max_usable_sm_count(max_sm_count: int) -> None: that several ranks can share one GPU. Must precede the first :func:`comm_workspace_size_per_rank` or :func:`alltoallv`, which fix the channel count.""" - get_nvlink_two_sided_module().set_moe_max_usable_sm_count(max_sm_count) + get_nvlink_two_sided_alltoall_module().set_moe_max_usable_sm_count(max_sm_count) def comm_workspace_size_per_rank(ep_size: int) -> int: """Bytes per rank of the FIFO workspace used by :func:`alltoallv`.""" - return get_nvlink_two_sided_module().get_moe_commworkspace_size_per_rank(ep_size) + return get_nvlink_two_sided_alltoall_module().get_moe_commworkspace_size_per_rank( + ep_size + ) def prepare_workspace_size_per_rank(ep_size: int) -> int: """Bytes per rank of the workspace used by :func:`prepare`.""" - return get_nvlink_two_sided_module().get_moe_prepare_workspace_size_per_rank( - ep_size + return ( + get_nvlink_two_sided_alltoall_module().get_moe_prepare_workspace_size_per_rank( + ep_size + ) ) @@ -62,7 +66,7 @@ def alltoallv( ``workspace`` is the ``[ep_size, words_per_rank]`` ``uint64`` FIFO workspace of :func:`comm_workspace_size_per_rank` bytes per rank. """ - get_nvlink_two_sided_module().moe_comm( + get_nvlink_two_sided_alltoall_module().moe_comm( input, send_rank_cum_sum, send_indices, @@ -140,7 +144,7 @@ def prepare( else torch.empty((ep_size, expert_count), **attrs) ) - get_nvlink_two_sided_module().moe_prepare( + get_nvlink_two_sided_alltoall_module().moe_prepare( expert_ids, scales, expert_statistics, diff --git a/tests/comm/test_cake_moe_alltoall_fused_adapter.py b/tests/comm/test_cake_moe_alltoall_fused_adapter.py index 2573426f937..df01ed84195 100644 --- a/tests/comm/test_cake_moe_alltoall_fused_adapter.py +++ b/tests/comm/test_cake_moe_alltoall_fused_adapter.py @@ -158,7 +158,7 @@ def spec(name): monkeypatch.setattr(aot, "gen_attention", lambda *args: ()) monkeypatch.setattr(aot, "gen_cudnn_fmha_module", lambda: spec("cudnn")) monkeypatch.setattr( - jit_comm, "gen_moe_ep_nvlink_two_sided_module", lambda: spec("two_sided") + jit_comm, "gen_nvlink_two_sided_alltoall_module", lambda: spec("two_sided") ) monkeypatch.setattr(jit_comm, "gen_vllm_comm_module", lambda: spec("vllm")) monkeypatch.setattr(jit_comm, "gen_pcie_ipc_comm_module", lambda: spec("pcie")) diff --git a/tests/moe_ep/test_communication.py b/tests/moe_ep/test_communication.py index f37d3a91e49..f492fc3af5a 100644 --- a/tests/moe_ep/test_communication.py +++ b/tests/moe_ep/test_communication.py @@ -583,7 +583,9 @@ def unsupported_reason(device_index): probe["probed"] += 1 return probe["reason"] - monkeypatch.setattr(one_sided, "get_nvlink_one_sided_module", lambda: module) + monkeypatch.setattr( + one_sided, "get_nvlink_one_sided_alltoall_module", lambda: module + ) monkeypatch.setattr( one_sided, "layout_constants", lambda: SimpleNamespace(**_ONE_SIDED_LAYOUT) ) diff --git a/tests/moe_ep/test_nvlink_one_sided.py b/tests/moe_ep/test_nvlink_one_sided.py index bf1a99afe68..a42a00d3106 100644 --- a/tests/moe_ep/test_nvlink_one_sided.py +++ b/tests/moe_ep/test_nvlink_one_sided.py @@ -784,7 +784,7 @@ def test_nvlink_one_sided(case: Case) -> None: def test_workspace_layout(cft: bool) -> None: """Regions are disjoint, aligned and fixed by the allocation-time capacities.""" from flashinfer.moe_ep.backends.split.comm.nvlink_one_sided.kernels import ( - get_nvlink_one_sided_module, + get_nvlink_one_sided_alltoall_module, layout_constants, ) @@ -794,7 +794,7 @@ def test_workspace_layout(cft: bool) -> None: combine_input_bytes = ep_size * capacity * hidden * 2 combine_recv_bytes = ep_size * capacity * hidden if cft else 0 layout = torch.empty(index.NUM_METAINFO_FIELDS, dtype=torch.int64) - get_nvlink_one_sided_module().moe_a2a_get_workspace_layout( + get_nvlink_one_sided_alltoall_module().moe_a2a_get_workspace_layout( ep_size, capacity, top_k, From 0747f5165f1cfc2569daa8eb9ff1b7af3f6f30a8 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:32:39 +0000 Subject: [PATCH 10/17] refactor(moe_ep): name the Cake all-to-all module after CakeAlltoAll gen_moe_ep_cake_module -> gen_cake_alltoall_module, the JIT modules moe_ep_cake_alltoall_{sm100a,sm103a} -> cake_alltoall_{sm100a,sm103a} and the loader get_cake_module -> get_cake_alltoall_module. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- flashinfer/aot.py | 6 +++--- flashinfer/jit/__init__.py | 2 +- flashinfer/jit/comm.py | 4 ++-- .../backends/split/comm/cake/kernels.py | 8 +++---- .../test_cake_moe_alltoall_fused_adapter.py | 21 ++++++++++--------- 5 files changed, 21 insertions(+), 20 deletions(-) diff --git a/flashinfer/aot.py b/flashinfer/aot.py index 0f40ac787db..4a244134c1f 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -1036,7 +1036,7 @@ def gen_all_modules( from .jit.comm import ( gen_dcp_alltoall_module, gen_dcp_lse_reduce_module, - gen_moe_ep_cake_module, + gen_cake_alltoall_module, gen_nvlink_one_sided_alltoall_module, gen_nvlink_two_sided_alltoall_module, gen_pcie_ipc_ag_rs_module, @@ -1077,9 +1077,9 @@ def gen_all_modules( ): jit_specs.append(gen_dcp_lse_reduce_module()) if has_sm100a_exact: - jit_specs.append(gen_moe_ep_cake_module("sm100a")) + jit_specs.append(gen_cake_alltoall_module("sm100a")) if has_sm103a_exact: - jit_specs.append(gen_moe_ep_cake_module("sm103a")) + jit_specs.append(gen_cake_alltoall_module("sm103a")) # Cake NVFP4-transport tensor-parallel all-reduce: exact-architecture # programs, one module per (program, exact target); never a family # target. diff --git a/flashinfer/jit/__init__.py b/flashinfer/jit/__init__.py index 383e78169ac..409dd23ee06 100644 --- a/flashinfer/jit/__init__.py +++ b/flashinfer/jit/__init__.py @@ -82,7 +82,7 @@ from .comm import gen_trtllm_comm_module as gen_trtllm_comm_module from .comm import gen_vllm_comm_module as gen_vllm_comm_module from .comm import gen_pcie_ipc_comm_module as gen_pcie_ipc_comm_module -from .comm import gen_moe_ep_cake_module as gen_moe_ep_cake_module +from .comm import gen_cake_alltoall_module as gen_cake_alltoall_module from .comm import ( gen_nvlink_one_sided_alltoall_module as gen_nvlink_one_sided_alltoall_module, ) diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index c4e75845067..17a7439d73a 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -317,11 +317,11 @@ def gen_ulysses_a2a_module() -> JitSpec: } -def gen_moe_ep_cake_module(target: CakeAlltoAllTarget) -> JitSpec: +def gen_cake_alltoall_module(target: CakeAlltoAllTarget) -> JitSpec: """Cake MoE all-to-all of one architecture behind ``flashinfer.moe_ep.CakeAlltoAll``.""" if target not in _CAKE_ALLTOALL_GENERATED_SOURCE: raise ValueError(f"unsupported Cake MoE all-to-all target: {target}") - module_name = f"moe_ep_cake_alltoall_{target}" + module_name = f"cake_alltoall_{target}" kernel_sources = [ jit_env.FLASHINFER_CSRC_DIR / "cake_moe_alltoall_dispatch.cu", jit_env.FLASHINFER_CSRC_DIR diff --git a/flashinfer/moe_ep/backends/split/comm/cake/kernels.py b/flashinfer/moe_ep/backends/split/comm/cake/kernels.py index 7b01898d59c..65f836dd5d7 100644 --- a/flashinfer/moe_ep/backends/split/comm/cake/kernels.py +++ b/flashinfer/moe_ep/backends/split/comm/cake/kernels.py @@ -25,11 +25,11 @@ @functools.cache -def get_cake_module(target: str) -> Any: +def get_cake_alltoall_module(target: str) -> Any: """Build (on first use) and load the module of one architecture.""" - from ......jit.comm import gen_moe_ep_cake_module + from ......jit.comm import gen_cake_alltoall_module - return gen_moe_ep_cake_module(target).build_and_load() # type: ignore[arg-type] + return gen_cake_alltoall_module(target).build_and_load() # type: ignore[arg-type] @functools.cache @@ -47,7 +47,7 @@ def _device_target(device_index: int) -> str: def current_module() -> Any: """The module of the current CUDA device's architecture.""" - return get_cake_module(_device_target(torch.cuda.current_device())) + return get_cake_alltoall_module(_device_target(torch.cuda.current_device())) def metainfo_index() -> dict[str, int]: diff --git a/tests/comm/test_cake_moe_alltoall_fused_adapter.py b/tests/comm/test_cake_moe_alltoall_fused_adapter.py index df01ed84195..415717414d3 100644 --- a/tests/comm/test_cake_moe_alltoall_fused_adapter.py +++ b/tests/comm/test_cake_moe_alltoall_fused_adapter.py @@ -47,18 +47,19 @@ def build_and_load(self): built.append(self.target) return SimpleNamespace(target=self.target) - kernels.get_cake_module.cache_clear() - monkeypatch.setattr(jit_comm, "gen_moe_ep_cake_module", FakeSpec) + kernels.get_cake_alltoall_module.cache_clear() + monkeypatch.setattr(jit_comm, "gen_cake_alltoall_module", FakeSpec) try: modules = { - target: kernels.get_cake_module(target) for target in ("sm100a", "sm103a") + target: kernels.get_cake_alltoall_module(target) + for target in ("sm100a", "sm103a") } for target, module in modules.items(): assert module.target == target - assert kernels.get_cake_module(target) is module + assert kernels.get_cake_alltoall_module(target) is module assert built == ["sm100a", "sm103a"] finally: - kernels.get_cake_module.cache_clear() + kernels.get_cake_alltoall_module.cache_clear() @pytest.mark.parametrize( @@ -66,13 +67,13 @@ def build_and_load(self): [ ( "sm100a", - "moe_ep_cake_alltoall_sm100a", + "cake_alltoall_sm100a", "cake_moe_alltoall_sm100.cu", "compute_100a,code=sm_100a", ), ( "sm103a", - "moe_ep_cake_alltoall_sm103a", + "cake_alltoall_sm103a", "cake_moe_alltoall_sm103.cu", "compute_103a,code=sm_103a", ), @@ -90,7 +91,7 @@ def capture(name, sources, **kwargs): return object() monkeypatch.setattr(jit_comm, "gen_jit_spec", capture) - jit_comm.gen_moe_ep_cake_module(target) + jit_comm.gen_cake_alltoall_module(target) names = {path.name for path in captured["sources"]} assert captured["name"] == module_name assert generated_name in names @@ -127,7 +128,7 @@ def get_device_capability(device): capability_queries.append(device) return capability - monkeypatch.setattr(kernels, "get_cake_module", lambda target: target) + monkeypatch.setattr(kernels, "get_cake_alltoall_module", lambda target: target) monkeypatch.setattr( kernels.torch.cuda, "get_device_capability", get_device_capability ) @@ -167,7 +168,7 @@ def spec(name): ) monkeypatch.setattr( jit_comm, - "gen_moe_ep_cake_module", + "gen_cake_alltoall_module", lambda target: selected.append(target) or spec(f"cake_{target}"), ) From 08a9f4ba11adef48f939696b7a71757fca1048d4 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:41:38 +0000 Subject: [PATCH 11/17] docs(moe_ep): list BootstrapConfig in the communication API Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- docs/api/moe_ep.rst | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/api/moe_ep.rst b/docs/api/moe_ep.rst index 11a98184d8f..9070d193775 100644 --- a/docs/api/moe_ep.rst +++ b/docs/api/moe_ep.rst @@ -10,12 +10,14 @@ Expert-parallel communication A :class:`MoEEpCommunication` dispatches each token to the ranks that own its experts and combines the expert outputs back on the token's source rank. -:func:`create_communication` builds one from a backend config; the backends -below run over NVLink (MNNVL) symmetric memory. +:func:`create_communication` builds one from a :class:`BootstrapConfig`, which +describes the expert-parallel group, and a backend config; the backends below +run over NVLink (MNNVL) symmetric memory. .. autosummary:: :toctree: ../generated + BootstrapConfig MoEEpCommunication MoEEpCommParams MoEEpDispatchResult From b23416d4dc5ce4d5097ff88dd643e531bd6514ca Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:46:41 +0000 Subject: [PATCH 12/17] docs(moe_ep): scope the API page to expert-parallel communication docs/api/moe_ep.rst -> docs/api/moe_ep_communication.rst, titled "flashinfer.moe_ep: Expert-Parallel Communication", since it documents only the communication API of flashinfer.moe_ep. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- docs/api/{moe_ep.rst => moe_ep_communication.rst} | 9 +++------ docs/index.rst | 2 +- 2 files changed, 4 insertions(+), 7 deletions(-) rename docs/api/{moe_ep.rst => moe_ep_communication.rst} (92%) diff --git a/docs/api/moe_ep.rst b/docs/api/moe_ep_communication.rst similarity index 92% rename from docs/api/moe_ep.rst rename to docs/api/moe_ep_communication.rst index 9070d193775..cb7a2e50467 100644 --- a/docs/api/moe_ep.rst +++ b/docs/api/moe_ep_communication.rst @@ -1,13 +1,10 @@ -.. _apimoe_ep: +.. _apimoe_ep_communication: -flashinfer.moe_ep -================= +flashinfer.moe_ep: Expert-Parallel Communication +================================================ .. currentmodule:: flashinfer.moe_ep -Expert-parallel communication ------------------------------ - A :class:`MoEEpCommunication` dispatches each token to the ranks that own its experts and combines the expert outputs back on the token's source rank. :func:`create_communication` builds one from a :class:`BootstrapConfig`, which diff --git a/docs/index.rst b/docs/index.rst index 23b9d0e705e..37a83fa458d 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -44,7 +44,7 @@ FlashInfer is a library and kernel generator for Large Language Models that prov api/grouped_mm api/fused_moe api/megamoe - api/moe_ep + api/moe_ep_communication api/cascade api/comm api/sparse From 6f10fc12cd56700271b59138016e2ddbaf671cfb Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:59:31 +0000 Subject: [PATCH 13/17] fix(moe_ep): include what the Cake dispatch header uses cake_moe_alltoall_dispatch.cuh uses nvinfer1::DataType and the fixed-width integer types but relied on its includers to provide them. After the move to csrc/ it sorts before tensorrt_llm/common/dataType.h in cake_moe_alltoall_dispatch.cu, so the Cake module stopped compiling. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- csrc/cake_moe_alltoall_dispatch.cuh | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/csrc/cake_moe_alltoall_dispatch.cuh b/csrc/cake_moe_alltoall_dispatch.cuh index b06304ce6d2..f4062d8cf08 100644 --- a/csrc/cake_moe_alltoall_dispatch.cuh +++ b/csrc/cake_moe_alltoall_dispatch.cuh @@ -18,6 +18,10 @@ #include #include +#include + +#include "tensorrt_llm/common/NvInferRuntime.h" + namespace tensorrt_llm::kernels::moe_alltoall { // Configuration constants From a61702e19f12c073889b2911605ff7710ece9d11 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:59:35 +0000 Subject: [PATCH 14/17] ci(comm): cover the MoE expert-parallel communication backends The comm-kernel lanes ran only the Cake kernel tests. The single-node lane now also runs the two-sided kernel test and tests/moe_ep/run_tests.sh comm (the MoEEpCommunication multirank test and the one-sided test, on up to four GPUs). The multi-node lane runs the multirank test across nodes, which can now be launched by an MPI launcher as well as by torchrun. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- scripts/task_test_multi_node_comm_kernels.sh | 2 +- scripts/task_test_single_node_comm_kernels.sh | 7 ++++ .../test_moe_ep_communication_multirank.py | 36 +++++++++++++++++++ 3 files changed, 44 insertions(+), 1 deletion(-) diff --git a/scripts/task_test_multi_node_comm_kernels.sh b/scripts/task_test_multi_node_comm_kernels.sh index 4798ac331eb..aacd4e3b11f 100755 --- a/scripts/task_test_multi_node_comm_kernels.sh +++ b/scripts/task_test_multi_node_comm_kernels.sh @@ -21,7 +21,7 @@ DISABLE_SANITY_TEST=true source "${SCRIPT_DIR}/test_utils.sh" # Define the specific test files for multi-node comm tests -TEST_FILES="tests/comm/test_mnnvl_memory.py tests/comm/test_trtllm_mnnvl_allreduce.py tests/comm/test_cake_moe_alltoall_multi_gpu.py tests/comm/test_trtllm_allreduce_checkpoint.py" +TEST_FILES="tests/comm/test_mnnvl_memory.py tests/comm/test_trtllm_mnnvl_allreduce.py tests/comm/test_cake_moe_alltoall_multi_gpu.py tests/comm/test_trtllm_allreduce_checkpoint.py tests/moe_ep/test_moe_ep_communication_multirank.py" # Main execution main() { diff --git a/scripts/task_test_single_node_comm_kernels.sh b/scripts/task_test_single_node_comm_kernels.sh index bd64649ef33..f586895fe70 100644 --- a/scripts/task_test_single_node_comm_kernels.sh +++ b/scripts/task_test_single_node_comm_kernels.sh @@ -59,6 +59,13 @@ pytest -s tests/moe/test_trtllm_cutlass_fused_moe.py pytest -s tests/comm/test_trtllm_moe_allreduce_fusion.py pytest -s tests/comm/test_trtllm_moe_allreduce_fusion_finalize.py pytest -s tests/comm/test_cake_moe_alltoall.py +# MoE expert-parallel communication backends (NVLink one-sided, Cake, two-sided) +pytest -s tests/comm/test_nvlink_two_sided_alltoall.py +if [ "$pcie_ipc_gpus" -lt 2 ]; then + echo "SKIPPING MoEEpCommunication multirank: $pcie_ipc_gpus GPU visible, needs >=2" +else + NPROC_MULTIRANK=$(( pcie_ipc_gpus < 4 ? pcie_ipc_gpus : 4 )) bash tests/moe_ep/run_tests.sh comm +fi # nvshmem ar pytest -s tests/comm/test_nvshmem.py pytest -s tests/comm/test_nvshmem_allreduce.py diff --git a/tests/moe_ep/test_moe_ep_communication_multirank.py b/tests/moe_ep/test_moe_ep_communication_multirank.py index 61007d2c203..83cf3abfdd9 100644 --- a/tests/moe_ep/test_moe_ep_communication_multirank.py +++ b/tests/moe_ep/test_moe_ep_communication_multirank.py @@ -2,6 +2,8 @@ Launched via torchrun: torchrun --nproc_per_node=4 -m pytest tests/moe_ep/test_moe_ep_communication_multirank.py -v +or via an MPI launcher, including across nodes: + mpirun -np 8 pytest tests/moe_ep/test_moe_ep_communication_multirank.py -v Each rank applies a synthetic expert computation to the rows it receives: a token routed to expert ``e`` with weight ``w`` contributes ``w * (e + 1) * x``. @@ -19,6 +21,40 @@ _PG_TIMEOUT = timedelta(minutes=60) _BACKENDS = ["nvlink_one_sided", "cake", "nvlink_two_sided"] +# Set by Open MPI, MPICH/Hydra and PMIx launchers respectively. +_MPI_LAUNCH_VARS = ("OMPI_COMM_WORLD_SIZE", "PMI_SIZE", "PMIX_RANK") + + +def _adopt_mpi_launch(): + """Set torch.distributed's ``env://`` variables from MPI when an MPI + launcher started this process instead of torchrun.""" + if "WORLD_SIZE" in os.environ or not any( + var in os.environ for var in _MPI_LAUNCH_VARS + ): + return + import socket + + from mpi4py import MPI + + comm = MPI.COMM_WORLD + master = None + if comm.Get_rank() == 0: + with socket.socket() as sock: + sock.bind(("", 0)) + master = (socket.gethostname(), sock.getsockname()[1]) + master_addr, master_port = comm.bcast(master, root=0) + local_rank = comm.Split_type(MPI.COMM_TYPE_SHARED).Get_rank() + os.environ.update( + RANK=str(comm.Get_rank()), + WORLD_SIZE=str(comm.Get_size()), + LOCAL_RANK=str(local_rank), + MASTER_ADDR=master_addr, + MASTER_PORT=str(master_port), + ) + + +# At import, so that collection's multi-rank checks see WORLD_SIZE. +_adopt_mpi_launch() def _init_dist(): From 5bef6e9bcb7de746c343cd2b728b477b0d05bbfa Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:10:45 +0000 Subject: [PATCH 15/17] test(comm): pass the EPLB argument to the Cake aux-data size moe_a2a_get_aux_data_size of the Cake module takes the EPLB expert count; the removed trtllm backend did not. Compare against the padded size, as the workspace lays out its regions. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- tests/comm/test_cake_moe_alltoall.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/comm/test_cake_moe_alltoall.py b/tests/comm/test_cake_moe_alltoall.py index 14c41493c11..f1cd071af56 100644 --- a/tests/comm/test_cake_moe_alltoall.py +++ b/tests/comm/test_cake_moe_alltoall.py @@ -1758,10 +1758,11 @@ def test_moe_workspace_size_per_rank(): ) assert empty_workspace_size > 0 - assert ( - empty_workspace_size - == cake_kernels.current_module().moe_a2a_get_aux_data_size(ep_size, num_tokens) + aux_data_size = cake_kernels.current_module().moe_a2a_get_aux_data_size( + ep_size, num_tokens, 0 ) + # Each workspace region is padded to 128 bytes. + assert empty_workspace_size == (aux_data_size + 127) // 128 * 128 non_empty_workspace_size = cake_kernels.workspace_size_per_rank( ep_size, num_tokens, hidden_size, hidden_size From 78f2137478cd9c31b48635a6163d895dcef09ab9 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:10:15 +0000 Subject: [PATCH 16/17] test: skip the Cake all-to-all suites on GPUs other than SM100/SM103 Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- tests/comm/test_cake_moe_alltoall.py | 7 +++++++ tests/comm/test_cake_moe_alltoall_multi_gpu.py | 8 ++++++++ 2 files changed, 15 insertions(+) diff --git a/tests/comm/test_cake_moe_alltoall.py b/tests/comm/test_cake_moe_alltoall.py index f1cd071af56..606a0087b0e 100644 --- a/tests/comm/test_cake_moe_alltoall.py +++ b/tests/comm/test_cake_moe_alltoall.py @@ -33,6 +33,13 @@ pynvml.nvmlInit() +# The Cake kernels are built only for SM100 and SM103. +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or get_compute_capability(torch.device("cuda:0")) not in ((10, 0), (10, 3)), + reason="Cake MoE all-to-all requires compute capability 10.0 or 10.3", +) + @pytest.fixture(autouse=True, scope="session") def setup_test_environment(): diff --git a/tests/comm/test_cake_moe_alltoall_multi_gpu.py b/tests/comm/test_cake_moe_alltoall_multi_gpu.py index 489e8b6dad2..852c5e29d4d 100644 --- a/tests/comm/test_cake_moe_alltoall_multi_gpu.py +++ b/tests/comm/test_cake_moe_alltoall_multi_gpu.py @@ -30,10 +30,18 @@ from flashinfer.comm.mapping import Mapping from flashinfer.comm.mnnvl import MnnvlMemory +from flashinfer.utils import get_compute_capability from .conftest import mnnvl_available from .cake_a2a_driver import CakeA2ADriver, metainfo_index, workspace_bytes_per_rank +# The Cake kernels are built only for SM100 and SM103. +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or get_compute_capability(torch.device("cuda:0")) not in ((10, 0), (10, 3)), + reason="Cake MoE all-to-all requires compute capability 10.0 or 10.3", +) + class MPIExit(Exception): pass From 13cdebbbf8c60b6ba0d52dbd7e71f4b6a5097627 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:10:19 +0000 Subject: [PATCH 17/17] ci: fail the multi-node comm lane when multi-rank tests lack a launcher The multi-node script sets FAIL_ON_DISTRIBUTED_SKIP=1, which makes the collection hook raise instead of skipping gpu_N tests when WORLD_SIZE is unset, so a misconfigured launch cannot pass with nothing run. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- scripts/task_test_multi_node_comm_kernels.sh | 4 ++++ tests/conftest.py | 7 +++++++ 2 files changed, 11 insertions(+) diff --git a/scripts/task_test_multi_node_comm_kernels.sh b/scripts/task_test_multi_node_comm_kernels.sh index aacd4e3b11f..3e0ebd3639b 100755 --- a/scripts/task_test_multi_node_comm_kernels.sh +++ b/scripts/task_test_multi_node_comm_kernels.sh @@ -23,6 +23,10 @@ source "${SCRIPT_DIR}/test_utils.sh" # Define the specific test files for multi-node comm tests TEST_FILES="tests/comm/test_mnnvl_memory.py tests/comm/test_trtllm_mnnvl_allreduce.py tests/comm/test_cake_moe_alltoall_multi_gpu.py tests/comm/test_trtllm_allreduce_checkpoint.py tests/moe_ep/test_moe_ep_communication_multirank.py" +# Every run of this script is launched across ranks, so multi-rank tests must +# not be skipped for a missing WORLD_SIZE (see tests/conftest.py). +export FAIL_ON_DISTRIBUTED_SKIP=1 + # Main execution main() { parse_args "$@" diff --git a/tests/conftest.py b/tests/conftest.py index ab2e63fc9f6..f4c3b1dfecc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -252,6 +252,13 @@ def pytest_collection_modifyitems(config, items): if mk not in item.keywords: continue if "WORLD_SIZE" not in os.environ: + # Runners that exist only to run multi-rank tests set this so + # that a missing launcher fails instead of skipping silently. + if os.environ.get("FAIL_ON_DISTRIBUTED_SKIP") == "1": + raise pytest.UsageError( + f"{item.nodeid} needs a torchrun or MPI launch " + "(WORLD_SIZE unset)" + ) # Multi-rank tests must be launched via torchrun (see # tests/moe_ep/run_tests.sh); under plain pytest auto-discovery # (e.g. CI unit-test sweeps) they would hang on dist init.