From 08d3bfda00724369febbeca4a107b41f7f31a730 Mon Sep 17 00:00:00 2001 From: djns99 <40156487+djns99@users.noreply.github.com> Date: Tue, 26 May 2026 13:16:35 +1200 Subject: [PATCH] [None][fix] Remove unused WarpPolicy in MOE OneSided A2A to reduce maintenance and prevent silent regressions Signed-off-by: djns99 <40156487+djns99@users.noreply.github.com> --- cpp/tensorrt_llm/common/envUtils.cpp | 17 +++-- cpp/tensorrt_llm/common/envUtils.h | 4 +- .../moeAlltoAllKernels.cu | 66 +++---------------- 3 files changed, 22 insertions(+), 65 deletions(-) diff --git a/cpp/tensorrt_llm/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp index bf3142a160dd..c5e22e65288d 100644 --- a/cpp/tensorrt_llm/common/envUtils.cpp +++ b/cpp/tensorrt_llm/common/envUtils.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -532,13 +532,18 @@ bool getEnvDisableChunkedAttentionInGenPhase() bool getEnvMoeA2AOneBlockPerToken() { - // Default true; return false only if env set to "0" - static std::optional const val = getIntEnv("TLLM_MOE_A2A_ONE_BLOCK_PER_TOKEN"); - if (!val.has_value()) + static bool const enabled = []() { + std::optional const val = getIntEnv("TLLM_MOE_A2A_ONE_BLOCK_PER_TOKEN"); + if (val.has_value()) + { + TLLM_LOG_WARNING( + "Changing A2A block scheduling mode via TLLM_MOE_A2A_ONE_BLOCK_PER_TOKEN is no longer " + "supported. One-block-per-token mode is now always enabled."); + } return true; - } - return val.value() != 0; + }(); + return enabled; } static int sanitizeBlockSize(std::optional const& val) diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index 82c56f300362..04fc592dbba3 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -153,7 +153,7 @@ bool getEnvDisaggBenchmarkGenOnly(); // Whether to disable the chunked-attention in the generation phase. bool getEnvDisableChunkedAttentionInGenPhase(); -// Whether to use one block per token for MoE A2A kernels (default true). +// Deprecated. Always returns true. bool getEnvMoeA2AOneBlockPerToken(); // TODO: For DEV purpose temporarily. diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index e1fd9bc7f08a..8e69c42bb27f 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -167,8 +167,7 @@ using tensorrt_llm::common::launchWithPdlWhenEnabled; } \ else \ { \ - using POLICY = WarpPolicy; \ - __VA_ARGS__ \ + TLLM_CHECK_WITH_INFO(false, "WarpPolicy is no longer supported for moe A2A"); \ } #if DISABLE_TIMEOUT @@ -201,29 +200,6 @@ __device__ int compute_target_rank_id(int expert_id, int num_experts_per_rank) // Helper Functions for Vectorized Memory Operations // ============================================================================ -struct WarpPolicy -{ - __device__ static int stride() - { - return warpSize; - } - - __device__ static int offset() - { - return (threadIdx.x % warpSize); - } - - __device__ static int token_idx() - { - return (blockIdx.x * blockDim.x + threadIdx.x) / warpSize; - } - - __device__ static void sync() - { - __syncwarp(); - } -}; - struct BlockPolicy { __device__ static int stride() @@ -423,20 +399,8 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ // Prepare per-policy shared-memory tiles for this token extern __shared__ int smem[]; - int* smem_topk_target_ranks; - int* smem_topk_send_indices; - int warps_per_block = blockDim.x / warpSize; - if constexpr (std::is_same::value) - { - int lane_id = threadIdx.x / warpSize; - smem_topk_target_ranks = smem + lane_id * TOP_K; - smem_topk_send_indices = smem + warps_per_block * TOP_K + lane_id * TOP_K; - } - else - { - smem_topk_target_ranks = smem; - smem_topk_send_indices = smem + TOP_K; - } + int* smem_topk_target_ranks = smem; + int* smem_topk_send_indices = smem + TOP_K; uint64_t already_copied = 0; int num_experts_per_rank = num_experts / ep_size; @@ -660,8 +624,6 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) kernel_ptrs.eplb_local_stats = params.eplb_local_stats; int const kBlockSize = tensorrt_llm::common::getEnvMoeA2ADispatchBlockSize(); - constexpr int kWarpSize = 32; - int const kWarpsPerBlock = kBlockSize / kWarpSize; // Configure kernel launch if (params.one_block_per_token) @@ -683,20 +645,7 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) } else { - int grid_size = ceilDiv(params.local_num_tokens, kWarpsPerBlock); - // 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; - } - int shared_bytes = 2 * kWarpsPerBlock * params.top_k * (int) sizeof(int); - SWITCH_BOOL(params.enable_eplb, EPLB_STATS, SWITCH_TOP_K(params.top_k, TOP_K, { - auto kernel_fn = moeA2ADispatchKernel; - launchWithPdlWhenEnabled("moeA2ADispatchKernel", kernel_fn, grid_size, kBlockSize, 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); - })) + TLLM_THROW("WarpPolicy is no longer supported for moe A2A"); } } @@ -1297,8 +1246,11 @@ void moe_a2a_prepare_combine_launch(MoeA2ACombineParams const& params) int const stride_per_token = low_precision_staged ? params.elements_per_token : params.elements_per_token * static_cast(sizeof(SrcT)); - auto kernel_fn = params.one_block_per_token ? moeA2APrepareCombineKernel - : moeA2APrepareCombineKernel; + if (!params.one_block_per_token) + { + TLLM_THROW("WarpPolicy is no longer supported for moe A2A"); + } + auto kernel_fn = moeA2APrepareCombineKernel; launchWithPdlWhenEnabled("moeA2APrepareCombineKernel", 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);