@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ecd1ea1363
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 07:15 UTC

[Kernel] Porting the TRTLLM minimax_allreduce_rms kernels (#37045)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Apr 11, 2026 at 00:20 UTC ecd1ea13634e6f9f6048dc8aecb94ce64a06b69e
14 files changed +1861 -4
.buildkite/test_areas/kernels.yaml
+14 -1
@@ -20,7 +20,20 @@ steps:
20 - tests/kernels/core
21 - tests/kernels/test_concat_mla_q.py
22 commands:
23 - - pytest -v -s kernels/core kernels/test_concat_mla_q.py
23 + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py
24 +
25 +- label: Kernels MiniMax Reduce RMS Test (2 GPUs)
26 + timeout_in_minutes: 15
27 + num_devices: 2
28 + device: h100
29 + source_file_dependencies:
30 + - csrc/minimax_reduce_rms_kernel.cu
31 + - csrc/minimax_reduce_rms_kernel.h
32 + - vllm/model_executor/layers/mamba/linear_attn.py
33 + - vllm/model_executor/layers/mamba/lamport_workspace.py
34 + - tests/kernels/core/test_minimax_reduce_rms.py
35 + commands:
36 + - pytest -v -s kernels/core/test_minimax_reduce_rms.py
37
38 - label: Kernels Attention Test %N
39 timeout_in_minutes: 35
CMakeLists.txt
+2
@@ -307,6 +307,8 @@ set(VLLM_EXT_SRC
307 "csrc/torch_bindings.cpp")
308
309 if(VLLM_GPU_LANG STREQUAL "CUDA")
310 + list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu")
311 +
312 SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
313
314 # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
csrc/minimax_reduce_rms_kernel.cu new
+879
@@ -0,0 +1,879 @@
1 +
2 +/*
3 + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
4 + *
5 + * Licensed under the Apache License, Version 2.0 (the "License");
6 + * you may not use this file except in compliance with the License.
7 + * You may obtain a copy of the License at
8 + *
9 + * http://www.apache.org/licenses/LICENSE-2.0
10 + *
11 + * Unless required by applicable law or agreed to in writing, software
12 + * distributed under the License is distributed on an "AS IS" BASIS,
13 + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 + * See the License for the specific language governing permissions and
15 + * limitations under the License.
16 + */
17 +
18 +#include <cooperative_groups.h>
19 +#include <cuda_runtime.h>
20 +
21 +#include <torch/cuda.h>
22 +#include <ATen/cuda/CUDAContext.h>
23 +#include <c10/cuda/CUDAGuard.h>
24 +
25 +#include "cuda_compat.h"
26 +#include "cuda_utils.h"
27 +#include "core/registration.h"
28 +#include "minimax_reduce_rms_kernel.h"
29 +
30 +#include <algorithm>
31 +
32 +#define FINAL_MASK 0xffffffff
33 +#define MINIMAX_REDUCE_RMS_WARP_SIZE 32
34 +
35 +namespace vllm {
36 +namespace tensorrt_llm {
37 +
38 +template <int NRanks>
39 +struct LamportComm {
40 + __device__ __forceinline__ LamportComm(void** workspace, int rank) {
41 + counter_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[0];
42 + flag_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[2];
43 + clear_ptr = &reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[0];
44 + flag_value = *flag_ptr;
45 + auto comm_size = reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[1];
46 + clear_size = *clear_ptr;
47 + int data_offset = flag_value % 3;
48 + int clear_offset = (flag_value + 2) % 3;
49 + for (int r = 0; r < NRanks; ++r) {
50 + data_bufs[r] = reinterpret_cast<uint8_t*>(workspace[2 * NRanks + r]) +
51 + data_offset * comm_size;
52 + }
53 + clear_buf = reinterpret_cast<uint8_t*>(workspace[2 * NRanks + rank]) +
54 + clear_offset * comm_size;
55 + __syncthreads();
56 + if (threadIdx.x == 0) {
57 + atomicAdd(counter_ptr, 1);
58 + }
59 + }
60 +
61 + __device__ __forceinline__ void update(int64_t new_clear_size) {
62 + if (blockIdx.x == 0 && threadIdx.x == 0) {
63 + while (*reinterpret_cast<int volatile*>(counter_ptr) != gridDim.x) {
64 + }
65 + *flag_ptr = (flag_value + 1) % 3;
66 + *clear_ptr = new_clear_size;
67 + *counter_ptr = 0;
68 + }
69 + }
70 +
71 + int* counter_ptr;
72 + int* flag_ptr;
73 + int64_t* clear_ptr;
74 + uint8_t* data_bufs[NRanks];
75 + uint8_t* clear_buf;
76 + int64_t clear_size;
77 + int flag_value;
78 +};
79 +
80 +__device__ __forceinline__ bool is_neg_zero(float v) {
81 + return *reinterpret_cast<uint32_t*>(&v) == 0x80000000;
82 +}
83 +
84 +__device__ __forceinline__ bool is_neg_zero(float4 v) {
85 + return is_neg_zero(v.x) || is_neg_zero(v.y) || is_neg_zero(v.z) ||
86 + is_neg_zero(v.w);
87 +}
88 +
89 +__device__ __forceinline__ float4 get_neg_zero() {
90 + float4 vec;
91 +#pragma unroll
92 + for (int i = 0; i < 4; ++i) {
93 + reinterpret_cast<uint32_t*>(&vec)[i] = 0x80000000;
94 + }
95 + return vec;
96 +}
97 +
98 +template <int Dim>
99 +__device__ __forceinline__ float rms_rsqrt(float& v, float eps) {
100 + constexpr float kInvDim = 1.0F / static_cast<float>(Dim);
101 + v = rsqrtf((v * kInvDim) + eps);
102 + return v;
103 +}
104 +
105 +template <int Dim>
106 +__device__ __forceinline__ float4 rms_rsqrt(float4& v, float eps) {
107 + constexpr float kInvDim = 1.0F / static_cast<float>(Dim);
108 + v.x = rsqrtf((v.x * kInvDim) + eps);
109 + v.y = rsqrtf((v.y * kInvDim) + eps);
110 + v.z = rsqrtf((v.z * kInvDim) + eps);
111 + v.w = rsqrtf((v.w * kInvDim) + eps);
112 + return v;
113 +}
114 +__device__ __forceinline__ float4 ld_global_volatile(float4* addr) {
115 + float4 val;
116 + asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];"
117 + : "=f"(val.x), "=f"(val.y), "=f"(val.z), "=f"(val.w)
118 + : "l"(addr));
119 + return val;
120 +}
121 +
122 +__device__ __forceinline__ float ld_global_volatile(float* addr) {
123 + float val;
124 + asm volatile("ld.volatile.global.f32 %0, [%1];" : "=f"(val) : "l"(addr));
125 + return val;
126 +}
127 +
128 +// Used by the scalar (non-float4) kernel only
129 +template <typename T, int NUM>
130 +__inline__ __device__ T warpReduceSumV2(T* val) {
131 +#pragma unroll
132 + for (int i = 0; i < NUM; i++) {
133 +#pragma unroll
134 + for (int mask = 16; mask > 0; mask >>= 1)
135 + val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32);
136 + }
137 + return (T)(0.0f);
138 +}
139 +
140 +template <typename T, int NUM>
141 +__inline__ __device__ T blockReduceSumV2(T* val) {
142 + static __shared__ T shared[NUM][33];
143 + int lane = threadIdx.x & 0x1f;
144 + int wid = threadIdx.x >> 5;
145 +
146 + warpReduceSumV2<T, NUM>(val);
147 +
148 + if (lane == 0) {
149 +#pragma unroll
150 + for (int i = 0; i < NUM; i++) {
151 + shared[i][wid] = val[i];
152 + }
153 + }
154 +
155 + __syncthreads();
156 +
157 + bool is_mask = threadIdx.x < (blockDim.x / 32.f);
158 +#pragma unroll
159 + for (int i = 0; i < NUM; i++) {
160 + val[i] = is_mask ? shared[i][lane] : (T)(0.0f);
161 + }
162 + warpReduceSumV2<T, NUM>(val);
163 + return (T)0.0f;
164 +}
165 +
166 +// for float4 version
167 +template <uint32_t kNumThreads, typename T, int ArraySize = 4>
168 +__device__ __forceinline__ void local_warp_reduce_sum_array(
169 + T* value_ptr, uint32_t active_mask = 0xffffffffu) {
170 + static_assert(kNumThreads >= 1 &&
171 + kNumThreads <= MINIMAX_REDUCE_RMS_WARP_SIZE);
172 +#pragma unroll
173 + for (int i = 0; i < ArraySize; ++i) {
174 +#pragma unroll
175 + for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) {
176 + value_ptr[i] += __shfl_xor_sync(active_mask, value_ptr[i], mask,
177 + MINIMAX_REDUCE_RMS_WARP_SIZE);
178 + }
179 + }
180 +}
181 +
182 +constexpr int next_pow2(int val) {
183 + int result = 1;
184 + while (result < val) {
185 + result <<= 1;
186 + }
187 + return result;
188 +}
189 +
190 +// ---------------------------------------------------------------------------
191 +
192 +template <typename DType>
193 +class IndexHelper {
194 + public:
195 + __device__ __forceinline__ IndexHelper(MiniMaxReduceRMSParams const& params) {
196 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
197 + namespace cg = cooperative_groups;
198 + cg::cluster_group cluster = cg::this_cluster();
199 + cg::grid_group grid = cg::this_grid();
200 + token_id = grid.cluster_rank();
201 + access_id_in_token = cluster.thread_rank();
202 + token_stride = grid.num_clusters();
203 +#else
204 + token_id = blockIdx.x;
205 + access_id_in_token = threadIdx.x;
206 + token_stride = gridDim.x;
207 +#endif
208 + access_id = token_id * params.hidden_dim / kElemsPerAccess<DType> +
209 + access_id_in_token;
210 + access_stride = token_stride * params.hidden_dim / kElemsPerAccess<DType>;
211 + tot_access = params.size_q / kElemsPerAccess<DType>;
212 + }
213 +
214 + int token_id;
215 + int access_id_in_token;
216 + int token_stride;
217 + int access_id;
218 + int access_stride;
219 + int tot_access;
220 +};
221 +
222 +/**
223 +* this kernel is used to for minimax attention module
224 +* input tensor [total_tokens, hidden_dim / tp_size], fp32
225 +* rms weight [hidden_dim / tp_size], bf16
226 +step 1: reduce from single rank to get the variance sum (reduce(input^2,
227 +dim=-1)) step 2: reduce from all ranks to get the variance sum
228 +(all_reduce(variance_sum)) step 3: calculate the rms norm (input *
229 +rsqrt(variance + eps)) in this case, max hidden_dim is 6144 (float data), for
230 +each token, we only need 6144 / 4 / tp_size = (1536 / tp_size) threads so we can
231 +assume cluster size is 1 (tp_size >= 2)
232 + */
233 +template <typename DType, int NRanks>
234 +__global__ void __launch_bounds__(1024)
235 + minimax_reduce_rms_kernel_lamport(MiniMaxReduceRMSParams params) {
236 + IndexHelper<DType> index_helper(params);
237 + int token_id = index_helper.token_id;
238 + int access_id_in_token = index_helper.access_id_in_token;
239 + int token_stride = index_helper.token_stride;
240 + int access_id = index_helper.access_id;
241 + int access_stride = index_helper.access_stride;
242 + int tot_access = index_helper.tot_access;
243 + int tot_tokens = params.size_q / params.hidden_dim;
244 + float4 clear_vec = get_neg_zero();
245 +
246 + LamportComm<NRanks> comm(params.workspace, params.rank);
247 + int clear_access = comm.clear_size / kElemsPerAccess<DType>;
248 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
249 + asm volatile("griddepcontrol.wait;");
250 +#endif
251 + for (int idx = access_id; idx < tot_access;
252 + idx += access_stride, token_id += token_stride) {
253 + alignas(16) DType vals[kElemsPerAccess<DType>];
254 + float sum_variance = 0.F;
255 + *reinterpret_cast<float4*>(vals) =
256 + reinterpret_cast<float4*>(params.allreduce_in)[idx];
257 +#pragma unroll
258 + for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
259 + sum_variance += static_cast<float>(vals[i]) * static_cast<float>(vals[i]);
260 + }
261 + blockReduceSumV2<float, 1>(&sum_variance);
262 + if (is_neg_zero(sum_variance)) {
263 + sum_variance = 0.F;
264 + }
265 + if (threadIdx.x == 0) {
266 + for (int r = 0; r < NRanks; ++r) {
267 + reinterpret_cast<float*>(
268 + comm.data_bufs[r])[(params.rank * tot_tokens) + token_id] =
269 + (sum_variance);
270 + }
271 + }
272 +
273 + bool done = false;
274 + float vars_all_ranks[NRanks];
275 + while (!done) {
276 + done = true;
277 +#pragma unroll
278 + for (int r = 0; r < NRanks; ++r) {
279 + vars_all_ranks[r] = ld_global_volatile(&reinterpret_cast<float*>(
280 + comm.data_bufs[params.rank])[(r * tot_tokens) + token_id]);
281 + done &= !is_neg_zero(vars_all_ranks[r]);
282 + }
283 + }
284 + sum_variance = 0.F;
285 +#pragma unroll
286 + for (int r = 0; r < NRanks; ++r) {
287 + sum_variance += vars_all_ranks[r];
288 + }
289 +
290 + DType norm_weight[kElemsPerAccess<DType>];
291 + *reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(norm_weight) =
292 + reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(
293 + params.rms_gamma)[access_id_in_token];
294 +
295 +#pragma unroll
296 + for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
297 + vals[i] = static_cast<DType>(
298 + static_cast<float>(vals[i]) *
299 + rsqrtf(
300 + (sum_variance / static_cast<float>(params.hidden_dim) / NRanks) +
301 + params.rms_eps) *
302 + static_cast<float>(norm_weight[i]));
303 + }
304 +
305 + reinterpret_cast<float4*>(params.rms_norm_out)[idx] =
306 + *reinterpret_cast<float4*>(vals);
307 + }
308 + for (int idx = access_id; idx < clear_access; idx += access_stride) {
309 + reinterpret_cast<float4*>(comm.clear_buf)[idx] = clear_vec;
310 + }
311 + comm.update(params.size_q * NRanks);
312 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
313 + asm volatile("griddepcontrol.launch_dependents;");
314 +#endif
315 +}
316 +
317 +/**
318 + * Float4 variant: process 4 rows at once, allreduce variance sums as float4 for
319 + * better memory coalescing. sum_variance is always float; applies to all DTypes
320 + * (half, bf16, float). When tot_tokens % 4 != 0, the last group pads rows with
321 + * zeros; padded rows are not written to rms_norm_out. IsQK: when true, process
322 + * Q+K in one loop with doubled comm buffer; when false, single-matrix (Q only).
323 + */
324 +template <typename DType, int NRanks, int OriginQDim, int OriginKDim>
325 +__global__ void __launch_bounds__(1024)
326 + minimax_reduce_qk_rms_kernel_lamport_float4(MiniMaxReduceRMSParams params) {
327 + // Compile-time per-rank dimensions
328 + constexpr int RankQDim = OriginQDim / NRanks;
329 + constexpr int RankKDim = OriginKDim / NRanks;
330 + // Threads needed to cover one row of Q / K with float4 accesses
331 + constexpr int ThreadsPerRowQ = RankQDim / kElemsPerAccess<DType>;
332 + constexpr int ThreadsPerRowK = RankKDim / kElemsPerAccess<DType>;
333 + // Number of warps dedicated to Q / K
334 + constexpr int NumWarpQ = (ThreadsPerRowQ + MINIMAX_REDUCE_RMS_WARP_SIZE - 1) /
335 + MINIMAX_REDUCE_RMS_WARP_SIZE;
336 + constexpr int NumWarpK = (ThreadsPerRowK + MINIMAX_REDUCE_RMS_WARP_SIZE - 1) /
337 + MINIMAX_REDUCE_RMS_WARP_SIZE;
338 +
339 + int tot_tokens = params.size_q / RankQDim;
340 + int tot_groups = (tot_tokens + 3) / 4; // ceiling; last group may be partial
341 +
342 + // Memory strides for strided qkv tensors (elements -> float4-access units)
343 + int access_stride_q = (params.stride_q > 0 ? params.stride_q : RankQDim) /
344 + kElemsPerAccess<DType>;
345 + int access_stride_k = (params.stride_k > 0 ? params.stride_k : RankKDim) /
346 + kElemsPerAccess<DType>;
347 + // Output strides: default to contiguous (hidden_dim / hidden_dim_k)
348 + int access_stride_q_out =
349 + (params.stride_q_out > 0 ? params.stride_q_out : params.hidden_dim) /
350 + kElemsPerAccess<DType>;
351 + int access_stride_k_out =
352 + (params.stride_k_out > 0 ? params.stride_k_out : params.hidden_dim_k) /
353 + kElemsPerAccess<DType>;
354 +
355 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
356 + namespace cg = cooperative_groups;
357 + cg::cluster_group cluster = cg::this_cluster();
358 + cg::grid_group grid = cg::this_grid();
359 + int group_id = grid.cluster_rank();
360 + int access_id_in_token = cluster.thread_rank();
361 + int group_stride = grid.num_clusters();
362 +#else
363 + int group_id = blockIdx.x;
364 + int access_id_in_token = threadIdx.x;
365 + int group_stride = gridDim.x;
366 +#endif
367 +
368 + bool is_q = (access_id_in_token < NumWarpQ * MINIMAX_REDUCE_RMS_WARP_SIZE);
369 + int k_thread_idx =
370 + access_id_in_token - (NumWarpQ * MINIMAX_REDUCE_RMS_WARP_SIZE);
371 + bool is_valid_q = (access_id_in_token < ThreadsPerRowQ);
372 + bool is_valid_k = (k_thread_idx >= 0 && k_thread_idx < ThreadsPerRowK);
373 + float4 clear_vec = get_neg_zero();
374 +
375 + // Shared memory for two-level block reduction and scale broadcast
376 + __shared__ float block_reduce_sum[4][MINIMAX_REDUCE_RMS_WARP_SIZE + 1];
377 + __shared__ float global_scale_q[4];
378 + __shared__ float global_scale_k[4];
379 +
380 + LamportComm<NRanks> comm(params.workspace, params.rank);
381 +
382 + DType norm_weight[kElemsPerAccess<DType>]{};
383 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
384 + asm volatile("griddepcontrol.wait;");
385 +#endif
386 + if (is_q) {
387 + if (is_valid_q) {
388 + *reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(
389 + norm_weight) =
390 + reinterpret_cast<typename ElemsPerAccess<DType>::vec_type const*>(
391 + params.rms_gamma)[access_id_in_token];
392 + }
393 + } else {
394 + if (is_valid_k) {
395 + *reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(
396 + norm_weight) =
397 + reinterpret_cast<typename ElemsPerAccess<DType>::vec_type const*>(
398 + params.rms_gamma_k)[k_thread_idx];
399 + }
400 + }
401 +
402 + // Main loop: process one group of 4 tokens per iteration.
403 + for (int g = group_id; g < tot_groups; g += group_stride) {
404 + alignas(16) DType vals[4][kElemsPerAccess<DType>]{};
405 + float warp_sum_variance[4]{0.F, 0.F, 0.F, 0.F};
406 +
407 + if (is_q) {
408 +#pragma unroll
409 + for (int row = 0; row < 4; ++row) {
410 + int token_r = g * 4 + row;
411 + if (token_r >= tot_tokens || !is_valid_q) {
412 + continue;
413 + }
414 + int idx_r = token_r * access_stride_q + access_id_in_token;
415 + *reinterpret_cast<float4*>(&vals[row][0]) =
416 + reinterpret_cast<float4 const*>(params.allreduce_in)[idx_r];
417 +#pragma unroll
418 + for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
419 + float x = static_cast<float>(vals[row][i]);
420 + warp_sum_variance[row] += x * x;
421 + }
422 + }
423 + } else {
424 +#pragma unroll
425 + for (int row = 0; row < 4; ++row) {
426 + int token_r = g * 4 + row;
427 + if (token_r >= tot_tokens || !is_valid_k) {
428 + continue;
429 + }
430 + int idx_r = token_r * access_stride_k + k_thread_idx;
431 + *reinterpret_cast<float4*>(&vals[row][0]) =
432 + reinterpret_cast<float4 const*>(params.allreduce_in_k)[idx_r];
433 +#pragma unroll
434 + for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
435 + float x = static_cast<float>(vals[row][i]);
436 + warp_sum_variance[row] += x * x;
437 + }
438 + }
439 + }
440 +
441 + local_warp_reduce_sum_array<MINIMAX_REDUCE_RMS_WARP_SIZE, float, 4>(
442 + warp_sum_variance);
443 + // Warp lane 0 writes its warp's partial sum to shared memory
444 + int lane = threadIdx.x & (MINIMAX_REDUCE_RMS_WARP_SIZE - 1);
445 + if (lane == 0) {
446 +#pragma unroll
447 + for (int t = 0; t < 4; ++t) {
448 + block_reduce_sum[t][threadIdx.x / MINIMAX_REDUCE_RMS_WARP_SIZE] =
449 + warp_sum_variance[t];
450 + }
451 + }
452 + __syncthreads();
453 +
454 + int tid = threadIdx.x;
455 +
456 + if (tid < MINIMAX_REDUCE_RMS_WARP_SIZE) {
457 + constexpr int kNumWarpQPow2 =
458 + (next_pow2(NumWarpQ) > NRanks) ? next_pow2(NumWarpQ) : NRanks;
459 + float local_sum[4];
460 +#pragma unroll
461 + for (int t = 0; t < 4; ++t) {
462 + local_sum[t] = (tid < NumWarpQ) ? block_reduce_sum[t][tid] : 0.F;
463 + }
464 + // After this, all kNumWarpQPow2 lanes (including tid 0..NRanks-1) have
465 + // the total Q sum-of-squares for all 4 tokens.
466 + local_warp_reduce_sum_array<kNumWarpQPow2, float, 4>(local_sum);
467 +
468 + if (tid < NRanks) {
469 +#pragma unroll
470 + for (int t = 0; t < 4; ++t) {
471 + if (is_neg_zero(local_sum[t])) {
472 + local_sum[t] = 0.F;
473 + }
474 + }
475 + // Parallel push: thread tid writes this rank's Q sum to rank tid's buf
476 + reinterpret_cast<float4*>(
477 + comm.data_bufs[tid])[(params.rank * tot_groups * 2) + (2 * g)] =
478 + *reinterpret_cast<float4*>(local_sum);
479 +
480 + // Parallel pull: thread tid reads rank tid's contribution from
481 + // this rank's (params.rank's) buffer
482 + bool done = false;
483 + float4 var_all_ranks;
484 + while (!done) {
485 + done = true;
486 + var_all_ranks = ld_global_volatile(&reinterpret_cast<float4*>(
487 + comm.data_bufs[params.rank])[(tid * tot_groups * 2) + (2 * g)]);
488 + done &= !is_neg_zero(var_all_ranks);
489 + }
490 +
491 + // Warp-level allreduce: each of the NRanks threads holds one rank's
492 + // partial sum; after this all NRanks threads have the global total.
493 + constexpr uint32_t kQActiveMask = (1u << NRanks) - 1u;
494 + local_warp_reduce_sum_array<NRanks, float, 4>(
495 + reinterpret_cast<float*>(&var_all_ranks), kQActiveMask);
496 +
497 + // Thread 0 computes rsqrt with compile-time Dim and writes to smem
498 + if (tid == 0) {
499 + *reinterpret_cast<float4*>(global_scale_q) =
500 + rms_rsqrt<OriginQDim>(var_all_ranks, params.rms_eps);
501 + }
502 + }
503 + } else if (tid >= MINIMAX_REDUCE_RMS_WARP_SIZE * NumWarpQ &&
504 + tid < MINIMAX_REDUCE_RMS_WARP_SIZE * (NumWarpQ + 1)) {
505 + // --- K leader warp ---
506 + constexpr int kNumWarpKPow2 =
507 + (next_pow2(NumWarpK) > NRanks) ? next_pow2(NumWarpK) : NRanks;
508 + float local_sum[4];
509 +#pragma unroll
510 + for (int t = 0; t < 4; ++t) {
511 + local_sum[t] = (k_thread_idx < NumWarpK)
512 + ? block_reduce_sum[t][NumWarpQ + k_thread_idx]
513 + : 0.F;
514 + }
515 + local_warp_reduce_sum_array<kNumWarpKPow2, float, 4>(local_sum);
516 +
517 + if (k_thread_idx < NRanks) {
518 +#pragma unroll
519 + for (int t = 0; t < 4; ++t) {
520 + if (is_neg_zero(local_sum[t])) {
521 + local_sum[t] = 0.F;
522 + }
523 + }
524 + reinterpret_cast<float4*>(
525 + comm.data_bufs[k_thread_idx])[(params.rank * tot_groups * 2) +
526 + (2 * g + 1)] =
527 + *reinterpret_cast<float4*>(local_sum);
528 +
529 + bool done = false;
530 + float4 var_all_ranks;
531 + while (!done) {
532 + done = true;
533 + var_all_ranks = ld_global_volatile(&reinterpret_cast<float4*>(
534 + comm.data_bufs[params.rank])[(k_thread_idx * tot_groups * 2) +
535 + (2 * g + 1)]);
536 + done &= !is_neg_zero(var_all_ranks);
537 + }
538 +
539 + constexpr uint32_t kKActiveMask = (1u << NRanks) - 1u;
540 + local_warp_reduce_sum_array<NRanks, float, 4>(
541 + reinterpret_cast<float*>(&var_all_ranks), kKActiveMask);
542 +
543 + if (k_thread_idx == 0) {
544 + *reinterpret_cast<float4*>(global_scale_k) =
545 + rms_rsqrt<OriginKDim>(var_all_ranks, params.rms_eps);
546 + }
547 + }
548 + }
549 + __syncthreads();
550 +
551 + if (is_q) {
552 +#pragma unroll
553 + for (int t = 0; t < 4; ++t) {
554 + warp_sum_variance[t] = global_scale_q[t];
555 + }
556 +#pragma unroll
557 + for (int r = 0; r < 4; ++r) {
558 +#pragma unroll
559 + for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
560 + vals[r][i] = static_cast<DType>(static_cast<float>(vals[r][i]) *
561 + warp_sum_variance[r] *
562 + static_cast<float>(norm_weight[i]));
563 + }
564 + int token_r = g * 4 + r;
565 + if (token_r >= tot_tokens || !is_valid_q) {
566 + continue;
567 + }
568 + int idx_out = token_r * access_stride_q_out + access_id_in_token;
569 + reinterpret_cast<float4*>(params.rms_norm_out)[idx_out] =
570 + *reinterpret_cast<float4*>(&vals[r][0]);
571 + }
572 + } else {
573 +#pragma unroll
574 + for (int t = 0; t < 4; ++t) {
575 + warp_sum_variance[t] = global_scale_k[t];
576 + }
577 +#pragma unroll
578 + for (int r = 0; r < 4; ++r) {
579 +#pragma unroll
580 + for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
581 + vals[r][i] = static_cast<DType>(static_cast<float>(vals[r][i]) *
582 + warp_sum_variance[r] *
583 + static_cast<float>(norm_weight[i]));
584 + }
585 + int token_r = g * 4 + r;
586 + if (token_r >= tot_tokens || !is_valid_k) {
587 + continue;
588 + }
589 + int idx_out = token_r * access_stride_k_out + k_thread_idx;
590 + reinterpret_cast<float4*>(params.rms_norm_out_k)[idx_out] =
591 + *reinterpret_cast<float4*>(&vals[r][0]);
592 + }
593 + }
594 + } // end group loop
595 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
596 + asm volatile("griddepcontrol.launch_dependents;");
597 +#endif
598 +
599 + int clear_access = static_cast<int>(comm.clear_size / kElemsPerAccess<DType>);
600 + int clear_stride = group_stride * blockDim.x;
601 + for (int idx = group_id * blockDim.x + threadIdx.x; idx < clear_access;
602 + idx += clear_stride) {
603 + reinterpret_cast<float4*>(comm.clear_buf)[idx] = clear_vec;
604 + }
605 +
606 + comm.update(static_cast<int64_t>(2) * tot_groups * kElemsPerAccess<DType> *
607 + NRanks);
608 +}
609 +
610 +int get_sm_count() {
611 + static int sm_count = 0;
612 + if (sm_count == 0) {
613 + int device_id;
614 + CUDA_CHECK(cudaGetDevice(&device_id));
615 + cudaDeviceProp device_prop;
616 + cudaGetDeviceProperties(&device_prop, device_id);
617 + sm_count = device_prop.multiProcessorCount;
618 + }
619 + return sm_count;
620 +}
621 +
622 +inline int getSMVersion(bool queryRealSmArch = false) {
623 + int device{-1};
624 + CUDA_CHECK(cudaGetDevice(&device));
625 + int sm_major = 0;
626 + int sm_minor = 0;
627 + CUDA_CHECK(cudaDeviceGetAttribute(&sm_major,
628 + cudaDevAttrComputeCapabilityMajor, device));
629 + CUDA_CHECK(cudaDeviceGetAttribute(&sm_minor,
630 + cudaDevAttrComputeCapabilityMinor, device));
631 + int sm = sm_major * 10 + sm_minor;
632 + if (sm == 121 && !queryRealSmArch) {
633 + return 120;
634 + }
635 + return sm;
636 +}
637 +
638 +template <typename KernelFunc>
639 +int get_max_active_blocks(KernelFunc kernel, int block_size,
640 + int dynamic_smem = 0) {
641 + int max_active = 0;
642 + CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
643 + &max_active, kernel, block_size, dynamic_smem));
644 + return std::max(max_active, 1);
645 +}
646 +
647 +template <typename DType, int NRanks>
648 +void minimax_reduce_rms_kernel_launcher(MiniMaxReduceRMSParams const& params) {
649 + static int SM = getSMVersion();
650 + int token_num = params.size_q / params.hidden_dim;
651 + int sm_count = get_sm_count();
652 + int cluster_size = 1;
653 + int cluster_num = token_num;
654 + int threads_per_token = params.hidden_dim / kElemsPerAccess<DType>;
655 + int block_size = threads_per_token;
656 +
657 + int max_blocks_per_sm = get_max_active_blocks(
658 + minimax_reduce_rms_kernel_lamport<DType, NRanks>, block_size);
659 + int max_grid = max_blocks_per_sm * sm_count;
660 +
661 + int grid_size =
662 + (std::min(max_grid, cluster_num * cluster_size) / cluster_size) *
663 + cluster_size;
664 +
665 + cudaLaunchConfig_t cfg;
666 + cfg.gridDim = grid_size;
667 + cfg.blockDim = block_size;
668 + cfg.dynamicSmemBytes = 0;
669 + cfg.stream = params.stream;
670 +
671 + cudaLaunchAttribute attribute[2];
672 + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
673 + attribute[0].val.programmaticStreamSerializationAllowed = 1;
674 + attribute[1].id = cudaLaunchAttributeClusterDimension;
675 + attribute[1].val.clusterDim.x = cluster_size;
676 + attribute[1].val.clusterDim.y = 1;
677 + attribute[1].val.clusterDim.z = 1;
678 + cfg.attrs = attribute;
679 + cfg.numAttrs = SM >= 90 ? 2 : 0;
680 +
681 + CUDA_CHECK(cudaLaunchKernelEx(
682 + &cfg, minimax_reduce_rms_kernel_lamport<DType, NRanks>, params));
683 +}
684 +
685 +template <typename DType, int NRanks, int OriginQDim, int OriginKDim>
686 +void minimax_reduce_rms_kernel_launcher_float4(
687 + MiniMaxReduceRMSParams const& params) {
688 + TORCH_CHECK(params.size_q % params.hidden_dim == 0);
689 + TORCH_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
690 + if (params.stride_q > 0) {
691 + TORCH_CHECK(params.stride_q % kElemsPerAccess<DType> == 0);
692 + }
693 + TORCH_CHECK(params.allreduce_in_k != nullptr,
694 + "float4 QK kernel requires K input");
695 + TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k);
696 + TORCH_CHECK(params.size_k % params.hidden_dim_k == 0);
697 + TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess<DType> == 0);
698 + TORCH_CHECK(params.size_q / params.hidden_dim ==
699 + params.size_k / params.hidden_dim_k);
700 + if (params.stride_k > 0) {
701 + TORCH_CHECK(params.stride_k % kElemsPerAccess<DType> == 0);
702 + }
703 +
704 + int token_num = params.size_q / params.hidden_dim;
705 + int tot_groups = (token_num + 3) / 4;
706 + if (tot_groups == 0) {
707 + return;
708 + }
709 +
710 + static int SM = getSMVersion();
711 + int sm_count = get_sm_count();
712 + int cluster_size = 1;
713 + int cluster_num = tot_groups;
714 +
715 + int access_per_row_q = params.hidden_dim / kElemsPerAccess<DType>;
716 + int access_per_row_k = params.hidden_dim_k / kElemsPerAccess<DType>;
717 +
718 + // Round each section up to a warp boundary
719 + auto divUp = [](int a, int b) { return (a + b - 1) / b * b; };
720 + int block_size = divUp(access_per_row_q, MINIMAX_REDUCE_RMS_WARP_SIZE) +
721 + divUp(access_per_row_k, MINIMAX_REDUCE_RMS_WARP_SIZE);
722 +
723 + auto kfn =
724 + minimax_reduce_qk_rms_kernel_lamport_float4<DType, NRanks, OriginQDim,
725 + OriginKDim>;
726 +
727 + int max_blocks_per_sm = get_max_active_blocks(kfn, block_size);
728 + int max_grid = max_blocks_per_sm * sm_count;
729 + int grid_size =
730 + (std::min(max_grid, cluster_num * cluster_size) / cluster_size) *
731 + cluster_size;
732 +
733 + cudaLaunchConfig_t cfg;
734 + cfg.gridDim = grid_size;
735 + cfg.blockDim = block_size;
736 + cfg.dynamicSmemBytes = 0;
737 + cfg.stream = params.stream;
738 +
739 + cudaLaunchAttribute attribute[2];
740 + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
741 + attribute[0].val.programmaticStreamSerializationAllowed = 1;
742 + attribute[1].id = cudaLaunchAttributeClusterDimension;
743 + attribute[1].val.clusterDim.x = cluster_size;
744 + attribute[1].val.clusterDim.y = 1;
745 + attribute[1].val.clusterDim.z = 1;
746 + cfg.attrs = attribute;
747 + cfg.numAttrs = SM >= 90 ? 2 : 0;
748 +
749 + CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params));
750 +}
751 +
752 +template <int NRanks>
753 +void dispatch_dtype(MiniMaxReduceRMSParams const& params) {
754 + // Use the optimized QK float4 kernel when:
755 + // - K input is present, AND
756 + // - the full (NRanks * per-rank) dimensions match the MiniMax M2 shape.
757 + // Otherwise fall back to the scalar kernel.
758 + bool use_float4 = (params.allreduce_in_k != nullptr) &&
759 + (params.hidden_dim * params.nranks == 6144) &&
760 + (params.hidden_dim_k * params.nranks == 1024);
761 +
762 + if (params.dtype == at::ScalarType::Half) {
763 + if (use_float4) {
764 + minimax_reduce_rms_kernel_launcher_float4<half, NRanks, 6144, 1024>(
765 + params);
766 + } else {
767 + minimax_reduce_rms_kernel_launcher<half, NRanks>(params);
768 + }
769 + } else if (params.dtype == at::ScalarType::BFloat16) {
770 + if (use_float4) {
771 + minimax_reduce_rms_kernel_launcher_float4<__nv_bfloat16, NRanks, 6144,
772 + 1024>(params);
773 + } else {
774 + minimax_reduce_rms_kernel_launcher<__nv_bfloat16, NRanks>(params);
775 + }
776 + } else if (params.dtype == at::ScalarType::Float) {
777 + if (use_float4) {
778 + minimax_reduce_rms_kernel_launcher_float4<float, NRanks, 6144, 1024>(
779 + params);
780 + } else {
781 + minimax_reduce_rms_kernel_launcher<float, NRanks>(params);
782 + }
783 + } else {
784 + TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op");
785 + }
786 +}
787 +
788 +void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) {
789 + if (params.nranks == 2) {
790 + dispatch_dtype<2>(params);
791 + } else if (params.nranks == 4) {
792 + dispatch_dtype<4>(params);
793 + } else if (params.nranks == 8) {
794 + dispatch_dtype<8>(params);
795 + } else if (params.nranks == 16) {
796 + dispatch_dtype<16>(params);
797 + } else {
798 + TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!");
799 + }
800 +}
801 +} // namespace tensorrt_llm
802 +} // namespace vllm
803 +
804 +torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
805 + torch::Tensor const& norm_weight,
806 + torch::Tensor workspace, int64_t const rank,
807 + int64_t const nranks, double const eps) {
808 + auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
809 +
810 + allreduce_params.nranks = static_cast<int>(nranks);
811 + allreduce_params.rank = static_cast<int>(rank);
812 + allreduce_params.dtype = input.scalar_type();
813 + allreduce_params.size_q = static_cast<int>(input.numel());
814 + allreduce_params.hidden_dim = static_cast<int>(input.size(-1));
815 + allreduce_params.stride_q = allreduce_params.hidden_dim;
816 + allreduce_params.workspace =
817 + reinterpret_cast<void**>(workspace.mutable_data_ptr());
818 + allreduce_params.allreduce_in = input.data_ptr();
819 + allreduce_params.rms_gamma = norm_weight.data_ptr();
820 + allreduce_params.rms_eps = static_cast<float>(eps);
821 + allreduce_params.stream = at::cuda::getCurrentCUDAStream(input.get_device());
822 +
823 + torch::Tensor rms_norm_out = torch::empty_like(input);
824 + allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr();
825 +
826 + vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params);
827 +
828 + return rms_norm_out;
829 +}
830 +
831 +std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
832 + torch::Tensor qkv, torch::Tensor const& norm_weight_q,
833 + torch::Tensor const& norm_weight_k, torch::Tensor workspace,
834 + int64_t const q_size, int64_t const kv_size, int64_t const rank,
835 + int64_t const nranks, double const eps) {
836 + TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D");
837 + TORCH_CHECK(qkv.is_contiguous(),
838 + "minimax_allreduce_rms_qk: qkv must be contiguous");
839 + int64_t qkv_dim = qkv.size(-1);
840 + TORCH_CHECK(qkv_dim == q_size + 2 * kv_size,
841 + "minimax_allreduce_rms_qk: qkv last dim must equal "
842 + "q_size + 2 * kv_size");
843 + TORCH_CHECK(rank < nranks,
844 + "minimax_allreduce_rms_qk: rank must be less than nranks");
845 +
846 + int64_t num_tokens = qkv.size(0);
847 + int elem_bytes = qkv.element_size();
848 +
849 + torch::Tensor q_out = torch::empty({num_tokens, q_size}, qkv.options());
850 + torch::Tensor k_out = torch::empty({num_tokens, kv_size}, qkv.options());
851 +
852 + auto params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
853 + params.nranks = static_cast<int>(nranks);
854 + params.rank = static_cast<int>(rank);
855 + params.dtype = qkv.scalar_type();
856 + params.size_q = static_cast<int>(num_tokens * q_size);
857 + params.hidden_dim = static_cast<int>(q_size);
858 + params.size_k = static_cast<int>(num_tokens * kv_size);
859 + params.hidden_dim_k = static_cast<int>(kv_size);
860 + params.stride_q = static_cast<int>(qkv_dim);
861 + params.stride_k = static_cast<int>(qkv_dim);
862 + params.stride_q_out = 0; // q_out is contiguous; kernel uses hidden_dim
863 + params.stride_k_out = 0; // k_out is contiguous; kernel uses hidden_dim_k
864 + params.workspace = reinterpret_cast<void**>(workspace.mutable_data_ptr());
865 +
866 + uint8_t* base = static_cast<uint8_t*>(qkv.data_ptr());
867 + params.allreduce_in = base;
868 + params.allreduce_in_k = base + q_size * elem_bytes;
869 + params.rms_gamma = norm_weight_q.data_ptr();
870 + params.rms_gamma_k = norm_weight_k.data_ptr();
871 + params.rms_eps = static_cast<float>(eps);
872 + params.stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
873 +
874 + params.rms_norm_out = q_out.mutable_data_ptr();
875 + params.rms_norm_out_k = k_out.mutable_data_ptr();
876 +
877 + vllm::tensorrt_llm::minimax_reduce_rms_op(params);
878 + return {q_out, k_out};
879 +}
csrc/minimax_reduce_rms_kernel.h new
+79
@@ -0,0 +1,79 @@
1 +/*
2 + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
3 + *
4 + * Licensed under the Apache License, Version 2.0 (the "License");
5 + * you may not use this file except in compliance with the License.
6 + * You may obtain a copy of the License at
7 + *
8 + * http://www.apache.org/licenses/LICENSE-2.0
9 + *
10 + * Unless required by applicable law or agreed to in writing, software
11 + * distributed under the License is distributed on an "AS IS" BASIS,
12 + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 + * See the License for the specific language governing permissions and
14 + * limitations under the License.
15 + */
16 +
17 +#pragma once
18 +
19 +#include <cuda_bf16.h>
20 +#include <cuda_fp16.h>
21 +
22 +#include <torch/types.h>
23 +
24 +namespace vllm {
25 +namespace tensorrt_llm {
26 +
27 +template <typename DType>
28 +struct ElemsPerAccess;
29 +
30 +template <>
31 +struct ElemsPerAccess<half> {
32 + static constexpr int value = 8;
33 + using vec_type = float4;
34 +};
35 +
36 +template <>
37 +struct ElemsPerAccess<nv_bfloat16> {
38 + static constexpr int value = 8;
39 + using vec_type = float4;
40 +};
41 +
42 +template <>
43 +struct ElemsPerAccess<float> {
44 + static constexpr int value = 4;
45 + using vec_type = float4;
46 +};
47 +
48 +template <typename DType>
49 +static constexpr int kElemsPerAccess = ElemsPerAccess<DType>::value;
50 +
51 +struct MiniMaxReduceRMSParams {
52 + int nranks{};
53 + int rank{};
54 + at::ScalarType dtype{at::ScalarType::Undefined};
55 + int size_q{};
56 + int hidden_dim{};
57 + int size_k{};
58 + int hidden_dim_k{};
59 + int stride_q{}; // row stride for q input (elements); when > hidden_dim,
60 + // q is part of a wider qkv tensor
61 + int stride_k{}; // row stride for k input (elements); when > hidden_dim_k,
62 + // k is part of a wider qkv tensor
63 + int stride_q_out{}; // row stride for q output (elements); 0 = contiguous
64 + int stride_k_out{}; // row stride for k output (elements); 0 = contiguous
65 + void** workspace{};
66 + void* allreduce_in{};
67 + void* rms_norm_out{};
68 + void* rms_gamma{};
69 + void* allreduce_in_k{};
70 + void* rms_norm_out_k{};
71 + void* rms_gamma_k{};
72 + float rms_eps{};
73 + cudaStream_t stream{};
74 +};
75 +
76 +void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params);
77 +
78 +} // namespace tensorrt_llm
79 +} // namespace vllm
csrc/ops.h
+12
@@ -308,4 +308,16 @@ int64_t qr_max_size();
308 #ifndef USE_ROCM
309 void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
310 torch::Tensor const& mat_b);
311 +#endif
312 +
313 +#ifndef USE_ROCM
314 +torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
315 + torch::Tensor const& norm_weight,
316 + torch::Tensor workspace, int64_t const rank,
317 + int64_t const nranks, double const eps);
318 +std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
319 + torch::Tensor qkv, torch::Tensor const& norm_weight_q,
320 + torch::Tensor const& norm_weight_k, torch::Tensor workspace,
321 + int64_t const q_size, int64_t const kv_size, int64_t const rank,
322 + int64_t const nranks, double const eps);
323 #endif
\ No newline at end of file
csrc/torch_bindings.cpp
+23
@@ -496,6 +496,29 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
496 "Tensor? b_qzeros, "
497 "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
498 "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
499 +
500 + ops.def(
501 + "minimax_allreduce_rms("
502 + "Tensor input,"
503 + "Tensor norm_weight,"
504 + "Tensor workspace,"
505 + "int rank,"
506 + "int nranks,"
507 + "float eps) -> Tensor");
508 + ops.impl("minimax_allreduce_rms", torch::kCUDA, &minimax_allreduce_rms);
509 + ops.def(
510 + "minimax_allreduce_rms_qk("
511 + "Tensor qkv,"
512 + "Tensor norm_weight_q,"
513 + "Tensor norm_weight_k,"
514 + "Tensor workspace,"
515 + "int q_size,"
516 + "int kv_size,"
517 + "int rank,"
518 + "int nranks,"
519 + "float eps) -> (Tensor, Tensor)");
520 + ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk);
521 +
522 // conditionally compiled so impl in source file
523 #endif
524 }
tests/kernels/core/test_minimax_reduce_rms.py new
+152
@@ -0,0 +1,152 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for MiniMax QK RMS-norm: NCCL reference vs Lamport fused kernel."""
4 +
5 +import pytest
6 +import torch
7 +import torch.nn as nn
8 +from torch.multiprocessing import spawn
9 +
10 +from tests.kernels.utils import opcheck
11 +from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
12 +from vllm.distributed import cleanup_dist_env_and_memory
13 +from vllm.model_executor.layers.mamba.linear_attn import MiniMaxText01RMSNormTP
14 +from vllm.platforms import current_platform
15 +from vllm.utils.network_utils import get_open_port
16 +from vllm.utils.torch_utils import set_random_seed
17 +
18 +
19 +@ensure_current_vllm_config()
20 +def _worker_forward_qk(
21 + local_rank,
22 + world_size,
23 + port,
24 + num_tokens,
25 + hidden_q_full,
26 + hidden_k_full,
27 + dtype,
28 + seed,
29 + eps,
30 +):
31 + """Per-rank worker: compare NCCL allreduce path vs Lamport fused kernel."""
32 +
33 + if not hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
34 + cleanup_dist_env_and_memory()
35 + return
36 + device = torch.device(f"cuda:{local_rank}")
37 + torch.accelerator.set_device_index(device)
38 + init_test_distributed_environment(
39 + world_size, 1, local_rank, port, local_rank=local_rank
40 + )
41 +
42 + hq = hidden_q_full // world_size
43 + hk = hidden_k_full // world_size
44 +
45 + q_norm = MiniMaxText01RMSNormTP(hidden_q_full, eps=eps).cuda()
46 + k_norm = MiniMaxText01RMSNormTP(hidden_k_full, eps=eps).cuda()
47 +
48 + set_random_seed(seed)
49 + qw = torch.randn(hidden_q_full, dtype=dtype, device="cuda")
50 + kw = torch.randn(hidden_k_full, dtype=dtype, device="cuda")
51 + q_norm.weight = nn.Parameter(qw[local_rank * hq : (local_rank + 1) * hq])
52 + k_norm.weight = nn.Parameter(kw[local_rank * hk : (local_rank + 1) * hk])
53 +
54 + torch.manual_seed(seed + 1000 + local_rank)
55 + qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda")
56 +
57 + q_ref, k_ref, v_ref = qkv.clone().split([hq, hk, hk], dim=-1)
58 + ref_q, ref_k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q_ref, k_ref)
59 +
60 + # Set up Lamport workspace.
61 + from vllm.distributed.parallel_state import get_tp_group
62 + from vllm.model_executor.layers.mamba.lamport_workspace import (
63 + get_allreduce_workspace,
64 + )
65 +
66 + workspace = get_allreduce_workspace(
67 + rank=local_rank,
68 + world_size=world_size,
69 + max_tokens=num_tokens,
70 + process_group=get_tp_group().cpu_group,
71 + )
72 +
73 + opcheck(
74 + torch.ops._C.minimax_allreduce_rms_qk,
75 + (
76 + qkv.clone(),
77 + q_norm.weight,
78 + k_norm.weight,
79 + workspace,
80 + hq,
81 + hk,
82 + local_rank,
83 + world_size,
84 + eps,
85 + ),
86 + )
87 + fused_q, fused_k = torch.ops._C.minimax_allreduce_rms_qk(
88 + qkv.clone(),
89 + q_norm.weight,
90 + k_norm.weight,
91 + workspace,
92 + hq,
93 + hk,
94 + local_rank,
95 + world_size,
96 + eps,
97 + )
98 + _, _, fused_v = qkv.split([hq, hk, hk], dim=-1)
99 + torch.accelerator.synchronize()
100 +
101 + torch.testing.assert_close(
102 + fused_q,
103 + ref_q,
104 + atol=3e-2,
105 + rtol=3e-2,
106 + )
107 + torch.testing.assert_close(fused_k, ref_k, atol=3e-2, rtol=3e-2)
108 +
109 + cleanup_dist_env_and_memory()
110 +
111 +
112 +@pytest.mark.skipif(
113 + not current_platform.is_cuda(),
114 + reason="CUDA required",
115 +)
116 +@pytest.mark.parametrize("world_size", [2, 4, 8])
117 +@pytest.mark.parametrize("num_tokens", [1, 128, 333])
118 +@pytest.mark.parametrize(
119 + "hidden_dims",
120 + [(6144, 1024)],
121 +)
122 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
123 +@pytest.mark.parametrize("eps", [1e-6])
124 +@pytest.mark.parametrize("seed", [42])
125 +def test_minimax_reduce_rms_qk(
126 + world_size,
127 + num_tokens,
128 + hidden_dims,
129 + dtype,
130 + eps,
131 + seed,
132 +):
133 + num_gpus = current_platform.device_count()
134 + if num_gpus < world_size:
135 + pytest.skip(f"Need >= {world_size} GPUs, have {num_gpus}")
136 + hidden_q_full, hidden_k_full = hidden_dims
137 + port = str(get_open_port())
138 + spawn(
139 + _worker_forward_qk,
140 + args=(
141 + world_size,
142 + port,
143 + num_tokens,
144 + hidden_q_full,
145 + hidden_k_full,
146 + dtype,
147 + seed,
148 + eps,
149 + ),
150 + nprocs=world_size,
151 + join=True,
152 + )
vllm/_custom_ops.py
+35
@@ -3491,3 +3491,38 @@ if hasattr(torch.ops._C, "hadacore_transform"):
3491 @register_fake("_C::hadacore_transform")
3492 def _hadacore_transform_fake(x: torch.Tensor, inplace: bool) -> torch.Tensor:
3493 return torch.empty_like(x) if not inplace else x
3494 +
3495 +
3496 +if hasattr(torch.ops._C, "minimax_allreduce_rms"):
3497 +
3498 + @register_fake("_C::minimax_allreduce_rms")
3499 + def _minimax_allreduce_rms_fake(
3500 + input: torch.Tensor,
3501 + norm_weight: torch.Tensor,
3502 + workspace: torch.Tensor,
3503 + rank: int,
3504 + nranks: int,
3505 + eps: float,
3506 + ) -> torch.Tensor:
3507 + return torch.empty_like(input)
3508 +
3509 +
3510 +if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
3511 +
3512 + @register_fake("_C::minimax_allreduce_rms_qk")
3513 + def _minimax_allreduce_rms_qk_fake(
3514 + qkv: torch.Tensor,
3515 + norm_weight_q: torch.Tensor,
3516 + norm_weight_k: torch.Tensor,
3517 + workspace: torch.Tensor,
3518 + q_size: int,
3519 + kv_size: int,
3520 + rank: int,
3521 + nranks: int,
3522 + eps: float,
3523 + ) -> tuple[torch.Tensor, torch.Tensor]:
3524 + token_num = qkv.shape[0]
3525 + return (
3526 + torch.empty([token_num, q_size], dtype=qkv.dtype, device=qkv.device),
3527 + torch.empty([token_num, kv_size], dtype=qkv.dtype, device=qkv.device),
3528 + )
vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py new
+340
@@ -0,0 +1,340 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +Fusion pass: replace MiniMax QK allreduce + RMS norm with the Lamport
6 +fused kernel (minimax_allreduce_rms_qk) for decode-size batches.
7 +
8 +Pattern (inlined forward_qk in compiled graph):
9 + q, k, v = qkv.split([q_size, kv_size, kv_size], -1)
10 + q_fp32 = q.to(float32); k_fp32 = k.to(float32)
11 + q_var = q_fp32.pow(2).mean(-1, keepdim=True)
12 + k_var = k_fp32.pow(2).mean(-1, keepdim=True)
13 + qk_var = cat([q_var, k_var], -1)
14 + qk_var = allreduce(qk_var) / tp_world
15 + q_var, k_var = qk_var.chunk(2, -1)
16 + q_out = (q_fp32 * rsqrt(q_var + eps) * q_weight).to(orig_dtype)
17 + k_out = (k_fp32 * rsqrt(k_var + eps) * k_weight).to(orig_dtype)
18 + return q_out, k_out, v
19 +
20 +Replacement (pure, no in-place on qkv/q/k):
21 + q_out, k_out = minimax_qk_norm_fused(qkv, q_weight, k_weight, workspace, ...)
22 + v = qkv.split([q_size, kv_size, kv_size], -1)[2]
23 + return q_out, k_out, v
24 +
25 +is_applicable_for_range: only fires for compile_range.end <= max_decode_tokens
26 +so that large prefill batches fall through to the original forward_qk (= main).
27 +"""
28 +
29 +import torch
30 +import torch._inductor.pattern_matcher as pm
31 +import torch.fx as fx
32 +from torch._inductor.pattern_matcher import PatternMatcherPass
33 +
34 +from vllm.config import VllmConfig
35 +from vllm.config.utils import Range
36 +from vllm.distributed import tensor_model_parallel_all_reduce
37 +from vllm.distributed.parallel_state import (
38 + get_tensor_model_parallel_rank,
39 + get_tensor_model_parallel_world_size,
40 +)
41 +from vllm.logger import init_logger
42 +from vllm.utils.torch_utils import direct_register_custom_op
43 +
44 +from ..inductor_pass import enable_fake_mode
45 +from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
46 +
47 +logger = init_logger(__name__)
48 +
49 +MAX_TOKEN_NUM = 2048
50 +
51 +_MINIMAX_QK_NORM_FUSED_OP = None
52 +if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
53 +
54 + def _minimax_qk_norm_fused(
55 + qkv: torch.Tensor,
56 + norm_weight_q: torch.Tensor,
57 + norm_weight_k: torch.Tensor,
58 + q_size: int,
59 + kv_size: int,
60 + rank: int,
61 + nranks: int,
62 + eps: float,
63 + max_tokens: int,
64 + ) -> tuple[torch.Tensor, torch.Tensor]:
65 + from vllm.distributed.parallel_state import get_tp_group
66 + from vllm.model_executor.layers.mamba.lamport_workspace import (
67 + get_allreduce_workspace,
68 + )
69 +
70 + workspace = get_allreduce_workspace(
71 + rank=rank,
72 + world_size=nranks,
73 + max_tokens=max_tokens,
74 + process_group=get_tp_group().cpu_group,
75 + )
76 + return torch.ops._C.minimax_allreduce_rms_qk(
77 + qkv,
78 + norm_weight_q,
79 + norm_weight_k,
80 + workspace,
81 + q_size,
82 + kv_size,
83 + rank,
84 + nranks,
85 + eps,
86 + )
87 +
88 + def _minimax_qk_norm_fused_fake(
89 + qkv: torch.Tensor,
90 + norm_weight_q: torch.Tensor,
91 + norm_weight_k: torch.Tensor,
92 + q_size: int,
93 + kv_size: int,
94 + rank: int,
95 + nranks: int,
96 + eps: float,
97 + max_tokens: int,
98 + ) -> tuple[torch.Tensor, torch.Tensor]:
99 + T = qkv.shape[0]
100 + return (
101 + torch.empty([T, q_size], dtype=qkv.dtype, device=qkv.device),
102 + torch.empty([T, kv_size], dtype=qkv.dtype, device=qkv.device),
103 + )
104 +
105 + direct_register_custom_op(
106 + op_name="minimax_qk_norm_fused",
107 + op_func=_minimax_qk_norm_fused,
108 + fake_impl=_minimax_qk_norm_fused_fake,
109 + mutates_args=[],
110 + )
111 + _MINIMAX_QK_NORM_FUSED_OP = torch.ops.vllm.minimax_qk_norm_fused.default
112 +
113 +
114 +class MiniMaxQKNormPattern:
115 + """
116 + Match the forward_qk allreduce+rms pattern and replace with Lamport kernel.
117 + """
118 +
119 + def __init__(
120 + self,
121 + q_size: int,
122 + kv_size: int,
123 + eps: float,
124 + tp_world: int,
125 + tp_rank: int,
126 + max_tokens: int,
127 + dtype: torch.dtype,
128 + device: str | None,
129 + ) -> None:
130 + self.q_size = q_size
131 + self.kv_size = kv_size
132 + self.eps = eps
133 + self.tp_world = tp_world
134 + self.tp_rank = tp_rank
135 + self.max_tokens = max_tokens
136 + self.dtype = dtype
137 + self.device = device
138 +
139 + def get_inputs(self) -> list[torch.Tensor]:
140 + T = 4
141 + qkv = torch.empty(
142 + [T, self.q_size + 2 * self.kv_size],
143 + device=self.device,
144 + dtype=self.dtype,
145 + )
146 + q_weight = torch.empty([self.q_size], device=self.device, dtype=self.dtype)
147 + k_weight = torch.empty([self.kv_size], device=self.device, dtype=self.dtype)
148 + return [qkv, q_weight, k_weight]
149 +
150 + def register(self, pm_pass: PatternMatcherPass) -> None:
151 + q_size = self.q_size
152 + kv_size = self.kv_size
153 + eps = self.eps
154 + tp_world = self.tp_world
155 + max_tokens = self.max_tokens
156 + tp_rank = self.tp_rank
157 + dtype = self.dtype
158 +
159 + def pattern(
160 + qkv: torch.Tensor,
161 + q_weight: torch.Tensor,
162 + k_weight: torch.Tensor,
163 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
164 + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
165 + q_fp32 = q.to(torch.float32)
166 + k_fp32 = k.to(torch.float32)
167 + q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True)
168 + k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True)
169 + qk_var = torch.cat([q_var, k_var], dim=-1)
170 + qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world
171 + q_var, k_var = qk_var.chunk(2, dim=-1)
172 + q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype)
173 + k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype)
174 + return q_out, k_out, v
175 +
176 + def replacement(
177 + qkv: torch.Tensor,
178 + q_weight: torch.Tensor,
179 + k_weight: torch.Tensor,
180 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
181 + assert _MINIMAX_QK_NORM_FUSED_OP is not None
182 + q_out, k_out = torch.ops.vllm.minimax_qk_norm_fused(
183 + qkv,
184 + q_weight,
185 + k_weight,
186 + q_size,
187 + kv_size,
188 + tp_rank,
189 + tp_world,
190 + eps,
191 + max_tokens,
192 + )
193 + _, _, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
194 + return q_out, k_out, v
195 +
196 + pm.register_replacement(
197 + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
198 + )
199 +
200 + # Second pattern: three separate split_with_sizes nodes (one per output),
201 + # each with _users=1. This occurs when the QKV projection uses a
202 + # functional GEMM kernel (e.g. cutlass_scaled_mm via auto_functionalized),
203 + # which causes inductor to generate one split per consumer.
204 + def pattern_split3(
205 + qkv: torch.Tensor,
206 + q_weight: torch.Tensor,
207 + k_weight: torch.Tensor,
208 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
209 + q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0]
210 + k = qkv.split([q_size, kv_size, kv_size], dim=-1)[1]
211 + v = qkv.split([q_size, kv_size, kv_size], dim=-1)[2]
212 + q_fp32 = q.to(torch.float32)
213 + k_fp32 = k.to(torch.float32)
214 + q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True)
215 + k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True)
216 + qk_var = torch.cat([q_var, k_var], dim=-1)
217 + qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world
218 + q_var, k_var = qk_var.chunk(2, dim=-1)
219 + q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype)
220 + k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype)
221 + return q_out, k_out, v
222 +
223 + pm.register_replacement(
224 + pattern_split3, replacement, self.get_inputs(), pm.fwd_only, pm_pass
225 + )
226 +
227 +
228 +class MiniMaxQKNormPass(VllmPatternMatcherPass):
229 + """
230 + Replace forward_qk allreduce+norm with the Lamport fused kernel.
231 + Only applied for decode-size compile ranges (small token counts).
232 + """
233 +
234 + def __init__(self, config: VllmConfig) -> None:
235 + super().__init__(config)
236 + self.disabled = True
237 +
238 + if _MINIMAX_QK_NORM_FUSED_OP is None:
239 + logger.warning_once(
240 + "minimax_allreduce_rms_qk op not found, MiniMaxQKNormPass disabled."
241 + )
242 + return
243 +
244 + tp_world = get_tensor_model_parallel_world_size()
245 + if tp_world <= 1:
246 + logger.warning_once("MiniMaxQKNormPass disabled: tp_size <= 1.")
247 + return
248 +
249 + if config.model_config is None:
250 + logger.warning_once("MiniMaxQKNormPass disabled: no model_config.")
251 + return
252 +
253 + hf_cfg = config.model_config.hf_config
254 +
255 + model_name = getattr(hf_cfg, "architectures", "")[0]
256 + if model_name != "MiniMaxM2ForCausalLM":
257 + return
258 +
259 + num_attention_heads = getattr(hf_cfg, "num_attention_heads", 0)
260 + num_key_value_heads = getattr(hf_cfg, "num_key_value_heads", 0)
261 + hidden_size = getattr(hf_cfg, "hidden_size", 0)
262 + head_dim = getattr(hf_cfg, "head_dim", 0)
263 + eps: float = getattr(hf_cfg, "rms_norm_eps", 1e-6)
264 +
265 + if (
266 + num_attention_heads != 48
267 + or num_key_value_heads != 8
268 + or hidden_size != 3072
269 + or head_dim != 128
270 + ):
271 + logger.warning_once(
272 + "MiniMaxQKNormPass disabled: cannot infer model info from hf_config."
273 + )
274 + return
275 +
276 + num_heads_per_rank = num_attention_heads // tp_world
277 + num_kv_heads_per_rank = max(1, num_key_value_heads // tp_world)
278 + q_size = num_heads_per_rank * head_dim
279 + kv_size = num_kv_heads_per_rank * head_dim
280 +
281 + self.max_token_num = min(
282 + MAX_TOKEN_NUM, config.scheduler_config.max_num_batched_tokens
283 + )
284 +
285 + tp_rank = get_tensor_model_parallel_rank()
286 + # Allocate Lamport workspace first.
287 + from vllm.distributed.parallel_state import get_tp_group
288 + from vllm.model_executor.layers.mamba.lamport_workspace import (
289 + get_allreduce_workspace,
290 + )
291 +
292 + get_allreduce_workspace(
293 + rank=tp_rank,
294 + world_size=tp_world,
295 + max_tokens=self.max_token_num,
296 + process_group=get_tp_group().cpu_group,
297 + )
298 +
299 + self.patterns: PatternMatcherPass = PatternMatcherPass(
300 + pass_name="minimax_qk_norm_pass"
301 + )
302 + self._register_patterns(q_size, kv_size, eps, tp_world, tp_rank)
303 + self.dump_patterns(config, self.patterns)
304 + self.disabled = False
305 +
306 + @enable_fake_mode
307 + def _register_patterns(
308 + self,
309 + q_size: int,
310 + kv_size: int,
311 + eps: float,
312 + tp_world: int,
313 + tp_rank: int,
314 + ) -> None:
315 + MiniMaxQKNormPattern(
316 + q_size=q_size,
317 + kv_size=kv_size,
318 + eps=eps,
319 + tp_world=tp_world,
320 + tp_rank=tp_rank,
321 + max_tokens=self.max_token_num,
322 + dtype=self.model_dtype,
323 + device=self.device,
324 + ).register(self.patterns)
325 +
326 + def is_applicable_for_range(self, compile_range: Range) -> bool:
327 + if self.disabled:
328 + return False
329 +
330 + return bool(compile_range.end <= self.max_token_num)
331 +
332 + @VllmInductorPass.time_and_log
333 + def __call__(self, graph: fx.Graph) -> None:
334 + if self.disabled:
335 + return
336 + self.matched_count = self.patterns.apply(graph)
337 + logger.debug("MiniMaxQKNormPass replaced %s patterns", self.matched_count)
338 +
339 + def uuid(self) -> str:
340 + return VllmInductorPass.hash_source(self, MiniMaxQKNormPattern)
vllm/compilation/passes/pass_manager.py
+4
@@ -38,6 +38,7 @@ if current_platform.is_cuda_alike():
38 if current_platform.is_cuda():
39 from .fusion.allreduce_rms_fusion import AllReduceFusionPass
40 from .fusion.collective_fusion import AsyncTPPass
41 + from .fusion.minimax_qk_norm_fusion import MiniMaxQKNormPass
42
43 from .inductor_pass import (
44 CustomGraphPass,
@@ -137,6 +138,9 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
138 if self.pass_config.fuse_allreduce_rms:
139 self.passes += [AllReduceFusionPass(config)]
140
141 + if self.pass_config.fuse_minimax_qk_norm:
142 + self.passes += [MiniMaxQKNormPass(config)]
143 +
144 if self.pass_config.fuse_norm_quant:
145 self.passes += [RMSNormQuantFusionPass(config)]
146 if rocm_aiter_ops.is_enabled():
vllm/config/compilation.py
+2
@@ -134,6 +134,8 @@ class PassConfig:
134 """Enable async TP."""
135 fuse_allreduce_rms: bool = None # type: ignore[assignment]
136 """Enable flashinfer allreduce fusion."""
137 + fuse_minimax_qk_norm: bool = None # type: ignore[assignment]
138 + """Enable fused allreduce+RMSNorm for MiniMax QK norm."""
139 enable_qk_norm_rope_fusion: bool = False
140 """Enable fused Q/K RMSNorm + RoPE pass."""
141
vllm/config/vllm.py
+16
@@ -1627,6 +1627,22 @@ class VllmConfig:
1627 compile_range_end,
1628 )
1629
1630 + if compilation_config.pass_config.fuse_minimax_qk_norm:
1631 + from vllm.compilation.passes.fusion.minimax_qk_norm_fusion import (
1632 + MAX_TOKEN_NUM,
1633 + )
1634 +
1635 + max_token_num = min(
1636 + MAX_TOKEN_NUM, self.scheduler_config.max_num_batched_tokens
1637 + )
1638 + if compile_range_end is not None and max_token_num < compile_range_end:
1639 + computed_compile_ranges_endpoints.append(max_token_num)
1640 + else:
1641 + logger.debug(
1642 + "Max num batched tokens below MiniMax QK norm fusion threshold, "
1643 + "MiniMax QK norm fusion enabled for all num_tokens."
1644 + )
1645 +
1646 if compilation_config.compile_ranges_endpoints is not None:
1647 for x in compilation_config.compile_ranges_endpoints:
1648 assert isinstance(x, int)
vllm/model_executor/layers/mamba/lamport_workspace.py new
+302
@@ -0,0 +1,302 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import array
6 +import contextlib
7 +import struct
8 +import sys
9 +import threading
10 +
11 +import torch
12 +
13 +try:
14 + from cuda.bindings import runtime as cudart
15 +except ImportError:
16 + from cuda import cudart
17 +
18 +_ALIGN = 1 << 21 # 2 MiB — CUDA IPC allocation alignment
19 +
20 +
21 +# ---------------------------------------------------------------------------
22 +# CUDA helpers
23 +# ---------------------------------------------------------------------------
24 +
25 +
26 +def _check(error):
27 + """Raise on CUDA runtime error."""
28 + success = getattr(cudart.cudaError_t, "cudaSuccess", None) or cudart.cudaError_t(0)
29 + if error != success:
30 + raise RuntimeError(f"CUDA runtime error: {error}")
31 +
32 +
33 +def _cuda_malloc(size: int):
34 + aligned = ((size + _ALIGN - 1) >> 21) << 21
35 + err, ptr = cudart.cudaMalloc(aligned)
36 + _check(err)
37 + return ptr, aligned
38 +
39 +
40 +def _cuda_free(ptr: int):
41 + if ptr:
42 + _check(cudart.cudaFree(ptr)[0])
43 +
44 +
45 +def _cuda_memset_zero(ptr: int, size: int):
46 + _check(cudart.cudaMemset(ptr, 0, size)[0])
47 +
48 +
49 +def _cuda_memcpy_d2d(dst: int, src: int, size: int):
50 + _check(
51 + cudart.cudaMemcpy(
52 + dst, src, size, cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice
53 + )[0]
54 + )
55 +
56 +
57 +# ---------------------------------------------------------------------------
58 +# IPC buffer
59 +# ---------------------------------------------------------------------------
60 +
61 +
62 +class IpcBuffer:
63 + """
64 + Allocates CUDA device memory and exchanges IPC handles with all ranks
65 + so that every rank holds a valid device pointer to every other rank's buffer.
66 + """
67 +
68 + def __init__(self, rank: int, world_size: int, size: int, process_group=None):
69 + self.rank = rank
70 + self.world_size = world_size
71 + self.peer_ptrs: list[int] = [0] * world_size
72 + self.local_ptr: int = 0
73 + self._alive = False
74 +
75 + if size <= 0:
76 + return
77 +
78 + self.local_ptr, _ = _cuda_malloc(size)
79 + _cuda_memset_zero(self.local_ptr, size)
80 + self._alive = True
81 +
82 + # --- exchange IPC handles via torch.distributed ---
83 + err, local_handle = cudart.cudaIpcGetMemHandle(self.local_ptr)
84 + _check(err)
85 +
86 + all_handles: list[bytes | None] = [None] * world_size
87 + torch.distributed.all_gather_object(
88 + all_handles, bytes(local_handle.reserved), group=process_group
89 + )
90 +
91 + for r in range(world_size):
92 + if r == rank:
93 + self.peer_ptrs[r] = self.local_ptr
94 + else:
95 + handle = cudart.cudaIpcMemHandle_t()
96 + handle.reserved = all_handles[r]
97 + err, ptr = cudart.cudaIpcOpenMemHandle(
98 + handle, cudart.cudaIpcMemLazyEnablePeerAccess
99 + )
100 + _check(err)
101 + self.peer_ptrs[r] = ptr
102 +
103 + def serialize(self) -> list[int]:
104 + """Return peer pointers as a list of int64 values (one per rank)."""
105 + raw = b""
106 + for ptr in self.peer_ptrs:
107 + raw += struct.pack("P", ptr)
108 + return array.array("Q", raw).tolist()
109 +
110 + def cleanup(self):
111 + if not self._alive:
112 + return
113 + self._alive = False
114 + for r in range(self.world_size):
115 + if self.peer_ptrs[r] == 0:
116 + continue
117 + if r == self.rank:
118 + _cuda_free(self.peer_ptrs[r])
119 + else:
120 + with contextlib.suppress(RuntimeError):
121 + _check(cudart.cudaIpcCloseMemHandle(self.peer_ptrs[r])[0])
122 + self.peer_ptrs[r] = 0
123 + self.local_ptr = 0
124 +
125 + def __del__(self):
126 + if not sys.is_finalizing():
127 + self.cleanup()
128 +
129 +
130 +# ---------------------------------------------------------------------------
131 +# Lamport negative-zero initialization
132 +# ---------------------------------------------------------------------------
133 +
134 +
135 +def _lamport_fill_neg_zero(device_ptr: int, size_bytes: int):
136 + """
137 + Fill device memory with IEEE-754 negative zero (-0.0f = 0x80000000).
138 + This is the "slot empty" sentinel for the Lamport protocol: the kernel
139 + spin-waits until a value is *not* negative zero.
140 + """
141 + if size_bytes == 0 or device_ptr == 0:
142 + return
143 + n_floats = size_bytes // 4
144 + # torch preserves -0.0 in IEEE-754
145 + fill = torch.full((n_floats,), -0.0, dtype=torch.float32, device="cuda")
146 + _cuda_memcpy_d2d(device_ptr, fill.data_ptr(), size_bytes)
147 + del fill
148 +
149 +
150 +# ---------------------------------------------------------------------------
151 +# LamportWorkspace — the main class
152 +# ---------------------------------------------------------------------------
153 +
154 +
155 +class LamportWorkspace:
156 + """
157 + Self-contained workspace for Lamport-based cross-GPU AllReduce.
158 +
159 + Parameters
160 + ----------
161 + rank : int
162 + Local rank (0-based).
163 + world_size : int
164 + Total number of ranks in the TP group.
165 + comm_size : int
166 + Size in bytes of *one* Lamport buffer slot. The total IPC allocation
167 + per rank is ``3 * comm_size`` (triple-buffering). Must be large enough
168 + to hold the per-slot data written by the kernel. Use
169 + ``compute_comm_size_for_minimax()`` for a safe default.
170 + process_group : optional
171 + ``torch.distributed`` process group for IPC handle exchange.
172 + ``None`` uses the default group.
173 + """
174 +
175 + def __init__(self, rank: int, world_size: int, comm_size: int, process_group=None):
176 + assert world_size >= 2, "Lamport workspace requires at least 2 ranks"
177 + assert comm_size > 0, "comm_size must be positive"
178 +
179 + self.rank = rank
180 + self.world_size = world_size
181 + self.comm_size = comm_size
182 +
183 + # 1) Lamport triple-buffer (the only IPC memory the kernel reads/writes)
184 + lamport_total = 3 * comm_size
185 + self._lamport = IpcBuffer(rank, world_size, lamport_total, process_group)
186 + _lamport_fill_neg_zero(self._lamport.local_ptr, lamport_total)
187 +
188 + # 2) flag_buffer on device: int32[3] = {counter, unused, lamport_flag}
189 + # counter — used for block-level sync inside the kernel
190 + # unused — reserved (index 1)
191 + # lamport_flag — triple-buffer rotation index (0 → 1 → 2 → 0 …)
192 + self._flag_buf = torch.zeros(3, dtype=torch.int32, device="cuda")
193 +
194 + # 3) layout_buffer on device: int64[2] = {clear_size, comm_size}
195 + # clear_size — bytes to clear from *previous* slot (set by kernel)
196 + # comm_size — size of one triple-buffer slot
197 + self._layout_buf = torch.tensor(
198 + [0, comm_size], dtype=torch.int64, device="cuda"
199 + )
200 +
201 + # 4) Assemble device-side void* pointer array
202 + N = world_size
203 + ptrs: list[int] = []
204 + ptrs += [0] * N # [0 .. N-1] ipc_buffers (placeholder)
205 + ptrs += [0] * N # [N .. 2N-1] ipc_barriers (placeholder)
206 + ptrs += self._lamport.serialize() # [2N .. 3N-1] lamport peer ptrs
207 + ptrs.append(self._flag_buf.data_ptr()) # [3N] flag_buffer
208 + ptrs.append(self._layout_buf.data_ptr()) # [3N+1] layout_buffer
209 +
210 + self._workspace = torch.tensor(ptrs, dtype=torch.int64, device="cuda")
211 +
212 + @property
213 + def workspace(self) -> torch.Tensor:
214 + """Device tensor (int64) that can be passed to the kernel
215 + as ``void** workspace``."""
216 + return self._workspace
217 +
218 + # ------------------------------------------------------------------
219 + # Helpers
220 + # ------------------------------------------------------------------
221 +
222 + @staticmethod
223 + def compute_comm_size_for_minimax(
224 + max_tokens: int,
225 + world_size: int,
226 + fused_qk: bool = True,
227 + ) -> int:
228 + """
229 + Return a safe ``comm_size`` (in bytes) for MiniMaxReduceRMSKernel.
230 +
231 + The kernel stores per-token variance scalars in the Lamport buffer:
232 + - single-matrix path: ``world_size × max_tokens × 4`` bytes per slot
233 + - fused Q+K path: ``world_size × 2 × ceil(max_tokens/4) × 16`` bytes per slot
234 +
235 + The returned value is rounded up to 2 MiB alignment.
236 + """
237 + if fused_qk:
238 + groups = (max_tokens + 3) // 4
239 + slot_bytes = world_size * 2 * groups * 16 # 16 = sizeof(float4)
240 + else:
241 + slot_bytes = world_size * max_tokens * 4 # 4 = sizeof(float)
242 + return ((slot_bytes + _ALIGN - 1) >> 21) << 21
243 +
244 + def cleanup(self):
245 + if hasattr(self, "_lamport"):
246 + self._lamport.cleanup()
247 +
248 + def __del__(self):
249 + if not sys.is_finalizing():
250 + self.cleanup()
251 +
252 + def __repr__(self):
253 + return (
254 + f"LamportWorkspace(rank={self.rank}, world_size={self.world_size}, "
255 + f"comm_size={self.comm_size})"
256 + )
257 +
258 +
259 +# ---------------------------------------------------------------------------
260 +# Cached convenience function (mirrors TRT-LLM's get_allreduce_workspace)
261 +# ---------------------------------------------------------------------------
262 +
263 +_cache_lock = threading.Lock()
264 +_workspace_cache: dict = {}
265 +
266 +
267 +def get_allreduce_workspace(
268 + rank: int,
269 + world_size: int,
270 + comm_size: int | None = None,
271 + max_tokens: int = 16384,
272 + process_group=None,
273 +) -> torch.Tensor:
274 + """
275 + Return a cached workspace tensor for the given (rank, world_size) pair.
276 +
277 + On first call the workspace is allocated and IPC handles are exchanged;
278 + subsequent calls with the same arguments return the cached tensor.
279 +
280 + Parameters
281 + ----------
282 + rank, world_size : int
283 + TP rank and TP size.
284 + comm_size : int, optional
285 + Explicit slot size in bytes. If ``None``, computed automatically
286 + from ``max_tokens`` and ``world_size`` (fused Q+K path).
287 + max_tokens : int
288 + Maximum number of tokens per batch (used when ``comm_size is None``).
289 + process_group : optional
290 + ``torch.distributed`` process group.
291 + """
292 + if comm_size is None:
293 + comm_size = LamportWorkspace.compute_comm_size_for_minimax(
294 + max_tokens, world_size, fused_qk=True
295 + )
296 + pg_id = id(process_group) if process_group is not None else 0
297 + key = (rank, world_size, comm_size, pg_id)
298 + with _cache_lock:
299 + if key not in _workspace_cache:
300 + ws = LamportWorkspace(rank, world_size, comm_size, process_group)
301 + _workspace_cache[key] = ws
302 + return _workspace_cache[key].workspace
vllm/model_executor/models/minimax_m2.py
+1 -3
@@ -233,9 +233,7 @@ class MiniMaxM2Attention(nn.Module):
233 ) -> torch.Tensor:
234 qkv, _ = self.qkv_proj(hidden_states)
235 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
236 - q, k = MiniMaxText01RMSNormTP.forward_qk(
237 - self.q_norm, self.k_norm, q.contiguous(), k.contiguous()
238 - )
236 + q, k = MiniMaxText01RMSNormTP.forward_qk(self.q_norm, self.k_norm, q, k)
237 q, k = self.rotary_emb(positions, q, k)
238 attn_output = self.attn(q, k, v)
239 output, _ = self.o_proj(attn_output)