[Refactor] Remove dead kernel code (#46405)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 25, 2026 at 21:09 UTC
cc7981599eac6d6ed6d08c07f9ec47d771969712
3 files changed
-431
csrc/custom_all_reduce_test.cu
deleted
-361
@@ -1,361 +0,0 @@
1
-/**
2
- * This is a standalone test for custom allreduce.
3
- * To compile, make sure you have MPI and NCCL installed in your system.
4
- * export MPI_HOME=XXX
5
- * nvcc -O2 -arch=native -std=c++17 custom_all_reduce_test.cu -o
6
- * custom_all_reduce_test -lnccl -I${MPI_HOME}/include -lmpi
7
- *
8
- * Warning: this C++ test is not designed to be very readable and was used
9
- * during the rapid prototyping process.
10
- *
11
- * To run:
12
- * mpirun --allow-run-as-root -np 8 ./custom_all_reduce_test
13
- */
14
-#include <cuda.h>
15
-#include <curand_kernel.h>
16
-#include <stdio.h>
17
-#include <stdlib.h>
18
-
19
-#include <limits>
20
-#include <vector>
21
-
22
-#include "cuda_profiler_api.h"
23
-#include "custom_all_reduce.cuh"
24
-#include "mpi.h"
25
-#ifdef USE_ROCM
26
- #include <hip/hip_bf16.h>
27
-typedef __hip_bfloat16 nv_bfloat16;
28
- #include "rccl/rccl.h"
29
- #include "custom_all_reduce_hip.cuh"
30
-#else
31
- #include "nccl.h"
32
- #include "custom_all_reduce.cuh"
33
-#endif
34
-
35
-#define MPICHECK(cmd) \
36
- do { \
37
- int e = cmd; \
38
- if (e != MPI_SUCCESS) { \
39
- printf("Failed: MPI error %s:%d '%d'\n", __FILE__, __LINE__, e); \
40
- exit(EXIT_FAILURE); \
41
- } \
42
- } while (0)
43
-
44
-#define NCCLCHECK(cmd) \
45
- do { \
46
- ncclResult_t r = cmd; \
47
- if (r != ncclSuccess) { \
48
- printf("Failed, NCCL error %s:%d '%s'\n", __FILE__, __LINE__, \
49
- ncclGetErrorString(r)); \
50
- exit(EXIT_FAILURE); \
51
- } \
52
- } while (0)
53
-
54
-#ifdef USE_ROCM
55
-__global__ void dummy_kernel() {
56
- for (int i = 0; i < 100; i++) {
57
- uint64_t start = wall_clock64();
58
- uint64_t cycles_elapsed;
59
- do {
60
- cycles_elapsed = wall_clock64() - start;
61
- } while (cycles_elapsed < 100);
62
- }
63
- for (int i = 0; i < 100; i++) __nanosleep(1000000); // 100ms
64
-}
65
-#else
66
-__global__ void dummy_kernel() {
67
- #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
68
- for (int i = 0; i < 100; i++) __nanosleep(1000000); // 100ms
69
- #else
70
- for (int i = 0; i < 100; i++) {
71
- long long int start = clock64();
72
- while (clock64() - start < 150000000); // approximately 98.4ms on P40
73
- }
74
- #endif
75
-}
76
-#endif
77
-
78
-template <typename T>
79
-__global__ void set_data(T* data, int size, int myRank) {
80
- for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
81
- idx += gridDim.x * blockDim.x) {
82
- data[idx] = myRank * 0.11f;
83
- }
84
-}
85
-
86
-template <typename T>
87
-__global__ void convert_data(const T* data1, const T* data2, double* fdata1,
88
- double* fdata2, int size) {
89
- for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
90
- idx += gridDim.x * blockDim.x) {
91
- fdata1[idx] = data1[idx];
92
- fdata2[idx] = data2[idx];
93
- }
94
-}
95
-
96
-__global__ void init_rand(curandState_t* state, int size, int nRanks) {
97
- for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
98
- idx += gridDim.x * blockDim.x) {
99
- for (int i = 0; i < nRanks; i++) {
100
- curand_init(i + 1, idx, 0, &state[idx * nRanks + i]);
101
- }
102
- }
103
-}
104
-
105
-template <typename T>
106
-__global__ void gen_data(curandState_t* state, T* data, double* ground_truth,
107
- int myRank, int nRanks, int size) {
108
- for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
109
- idx += gridDim.x * blockDim.x) {
110
- double sum = 0.0;
111
- for (int i = 0; i < nRanks; i++) {
112
- double val = curand_uniform_double(&state[idx * nRanks + i]) * 4;
113
- T hval = val; // downcast first
114
- sum += static_cast<double>(hval);
115
- if (i == myRank) data[idx] = hval;
116
- }
117
- ground_truth[idx] = sum;
118
- }
119
-}
120
-
121
-template <typename T>
122
-void run(int myRank, int nRanks, ncclComm_t& comm, int threads, int block_limit,
123
- int data_size, bool performance_test) {
124
- T* result;
125
- cudaStream_t stream;
126
- CUDACHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
127
- CUDACHECK(cudaMalloc(&result, data_size * sizeof(T)));
128
- CUDACHECK(cudaMemset(result, 0, data_size * sizeof(T)));
129
-
130
- cudaIpcMemHandle_t self_data_handle;
131
- cudaIpcMemHandle_t data_handles[8];
132
- vllm::Signal* buffer;
133
- T* self_data_copy;
134
- /**
135
- * Allocate IPC buffer
136
- *
137
- * The first section is a temporary buffer for storing intermediate allreduce
138
- * results, if a particular algorithm requires it. The second section is for
139
- * the input to the allreduce. The actual API takes the input pointer as an
140
- * argument (that is, they can and usually should be allocated separately).
141
- * But since the input pointers and the temporary buffer all require IPC
142
- * registration, they are allocated and registered together in the test for
143
- * convenience.
144
- */
145
-#ifdef USE_ROCM
146
- CUDACHECK(hipExtMallocWithFlags(
147
- (void**)&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal),
148
- hipDeviceMallocUncached));
149
-#else
150
- CUDACHECK(
151
- cudaMalloc(&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal)));
152
-#endif
153
- CUDACHECK(
154
- cudaMemset(buffer, 0, 2 * data_size * sizeof(T) + sizeof(vllm::Signal)));
155
- CUDACHECK(cudaMalloc(&self_data_copy, data_size * sizeof(T)));
156
- CUDACHECK(cudaIpcGetMemHandle(&self_data_handle, buffer));
157
-
158
- MPICHECK(MPI_Allgather(&self_data_handle, sizeof(cudaIpcMemHandle_t),
159
- MPI_BYTE, data_handles, sizeof(cudaIpcMemHandle_t),
160
- MPI_BYTE, MPI_COMM_WORLD));
161
-
162
- void* rank_data;
163
- size_t rank_data_sz = 16 * 1024 * 1024;
164
- CUDACHECK(cudaMalloc(&rank_data, rank_data_sz));
165
- vllm::Signal* ipc_ptrs[8];
166
- for (int i = 0; i < nRanks; i++) {
167
- if (i == myRank)
168
- ipc_ptrs[i] = buffer;
169
- else
170
- CUDACHECK(cudaIpcOpenMemHandle((void**)&ipc_ptrs[i], data_handles[i],
171
- cudaIpcMemLazyEnablePeerAccess));
172
- }
173
- vllm::CustomAllreduce fa(ipc_ptrs, rank_data, rank_data_sz, myRank, nRanks);
174
- auto* self_data =
175
- reinterpret_cast<T*>(reinterpret_cast<char*>(buffer) +
176
- sizeof(vllm::Signal) + data_size * sizeof(T));
177
- // hack buffer registration
178
- {
179
- void* data[8];
180
- for (int i = 0; i < nRanks; i++) {
181
- data[i] =
182
- ((char*)ipc_ptrs[i]) + sizeof(vllm::Signal) + data_size * sizeof(T);
183
- }
184
- fa.register_buffer(data);
185
- }
186
-
187
- double* ground_truth;
188
- CUDACHECK(cudaMallocHost(&ground_truth, data_size * sizeof(double)));
189
- curandState_t* states;
190
- CUDACHECK(cudaMalloc(&states, sizeof(curandState_t) * nRanks * data_size));
191
- init_rand<<<108, 1024, 0, stream>>>(states, data_size, nRanks);
192
- gen_data<T><<<108, 1024, 0, stream>>>(states, self_data, ground_truth, myRank,
193
- nRanks, data_size);
194
- CUDACHECK(cudaMemcpyAsync(self_data_copy, self_data, data_size * sizeof(T),
195
- cudaMemcpyDeviceToDevice, stream));
196
- cudaEvent_t start, stop;
197
- CUDACHECK(cudaEventCreate(&start));
198
- CUDACHECK(cudaEventCreate(&stop));
199
-
200
- ncclDataType_t ncclDtype;
201
- if (std::is_same<T, half>::value) {
202
- ncclDtype = ncclFloat16;
203
- } else if (std::is_same<T, nv_bfloat16>::value) {
204
- ncclDtype = ncclBfloat16;
205
- } else {
206
- ncclDtype = ncclFloat;
207
- }
208
- double *nccl_result, *my_result;
209
- CUDACHECK(cudaMallocHost(&nccl_result, data_size * sizeof(double)));
210
- CUDACHECK(cudaMallocHost(&my_result, data_size * sizeof(double)));
211
- if (performance_test) {
212
- dummy_kernel<<<1, 1, 0, stream>>>();
213
- constexpr int warmup_iters = 5;
214
- constexpr int num_iters = 100;
215
- // warmup
216
- for (int i = 0; i < warmup_iters; i++) {
217
- NCCLCHECK(ncclAllReduce(result, result, data_size, ncclDtype, ncclSum,
218
- comm, stream));
219
- }
220
- CUDACHECK(cudaEventRecord(start, stream));
221
- for (int i = 0; i < num_iters; i++) {
222
- NCCLCHECK(ncclAllReduce(result, result, data_size, ncclDtype, ncclSum,
223
- comm, stream));
224
- }
225
- CUDACHECK(cudaEventRecord(stop, stream));
226
- CUDACHECK(cudaStreamSynchronize(stream));
227
- float allreduce_ms = 0;
228
- cudaEventElapsedTime(&allreduce_ms, start, stop);
229
-
230
- dummy_kernel<<<1, 1, 0, stream>>>();
231
- // warm up
232
- for (int i = 0; i < warmup_iters; i++) {
233
- fa.allreduce<T>(stream, self_data, result, data_size, threads,
234
- block_limit);
235
- }
236
- CUDACHECK(cudaEventRecord(start, stream));
237
- for (int i = 0; i < num_iters; i++) {
238
- fa.allreduce<T>(stream, self_data, result, data_size, threads,
239
- block_limit);
240
- }
241
- CUDACHECK(cudaEventRecord(stop, stream));
242
- CUDACHECK(cudaStreamSynchronize(stream));
243
-
244
- float duration_ms = 0;
245
- cudaEventElapsedTime(&duration_ms, start, stop);
246
- if (myRank == 0)
247
- printf(
248
- "Rank %d done, nGPUs:%d, sz (kb): %d, %d, %d, my time:%.2fus, nccl "
249
- "time:%.2fus\n",
250
- myRank, nRanks, data_size * sizeof(T) / 1024, threads, block_limit,
251
- duration_ms * 1e3 / num_iters, allreduce_ms * 1e3 / num_iters);
252
-
253
- // And wait for all the queued up work to complete
254
- CUDACHECK(cudaStreamSynchronize(stream));
255
-
256
- NCCLCHECK(ncclAllReduce(self_data_copy, self_data, data_size, ncclDtype,
257
- ncclSum, comm, stream));
258
-
259
- convert_data<T><<<108, 1024, 0, stream>>>(self_data, result, nccl_result,
260
- my_result, data_size);
261
- CUDACHECK(cudaStreamSynchronize(stream));
262
-
263
- for (unsigned long j = 0; j < data_size; j++) {
264
- auto diff = abs(nccl_result[j] - my_result[j]);
265
- if (diff >= 4e-2) {
266
- printf("Rank %d: Verification mismatch at %lld: %f != (my) %f, gt=%f\n",
267
- myRank, j, nccl_result[j], my_result[j], ground_truth[j]);
268
- break;
269
- }
270
- }
271
- long double nccl_diffs = 0.0;
272
- long double my_diffs = 0.0;
273
- for (int j = 0; j < data_size; j++) {
274
- nccl_diffs += abs(nccl_result[j] - ground_truth[j]);
275
- my_diffs += abs(my_result[j] - ground_truth[j]);
276
- }
277
- if (myRank == 0)
278
- std::cout << "average abs diffs: nccl: " << nccl_diffs / data_size
279
- << " me: " << my_diffs / data_size << std::endl;
280
- } else {
281
- for (int i = 0; i < 100; i++) {
282
- fa.allreduce<T>(stream, self_data, result, data_size, threads,
283
- block_limit);
284
- CUDACHECK(cudaStreamSynchronize(stream));
285
- NCCLCHECK(ncclAllReduce(self_data, self_data_copy, data_size, ncclDtype,
286
- ncclSum, comm, stream));
287
- convert_data<T><<<108, 1024, 0, stream>>>(
288
- self_data_copy, result, nccl_result, my_result, data_size);
289
- CUDACHECK(cudaStreamSynchronize(stream));
290
-
291
- for (unsigned long j = 0; j < data_size; j++) {
292
- auto diff = abs(nccl_result[j] - my_result[j]);
293
- if (diff >= 4e-2) {
294
- printf(
295
- "Rank %d: Verification mismatch at %lld: %f != (my) %f, gt=%f\n",
296
- myRank, j, nccl_result[j], my_result[j], ground_truth[j]);
297
- break;
298
- }
299
- }
300
- }
301
- if (myRank == 0)
302
- printf("Test passed: nGPUs:%d, sz (kb): %d, %d, %d\n", nRanks,
303
- data_size * sizeof(T) / 1024, threads, block_limit);
304
- // long double nccl_diffs = 0.0;
305
- // long double my_diffs = 0.0;
306
- // for (int j = 0; j < data_size; j++) {
307
- // nccl_diffs += abs(nccl_result[j] - ground_truth[j]);
308
- // my_diffs += abs(my_result[j] - ground_truth[j]);
309
- // }
310
- // if (myRank == 0)
311
- // std::cout << "average abs diffs: nccl: " << nccl_diffs / data_size
312
- // << " me: " << my_diffs / data_size << std::endl;
313
- }
314
-
315
- CUDACHECK(cudaFree(result));
316
- CUDACHECK(cudaFree(self_data_copy));
317
- CUDACHECK(cudaFree(rank_data));
318
- CUDACHECK(cudaFree(buffer));
319
- CUDACHECK(cudaFree(states));
320
- CUDACHECK(cudaFreeHost(ground_truth));
321
- CUDACHECK(cudaFreeHost(nccl_result));
322
- CUDACHECK(cudaFreeHost(my_result));
323
- CUDACHECK(cudaStreamDestroy(stream));
324
-}
325
-
326
-int main(int argc, char** argv) {
327
- int nRanks, myRank;
328
- MPICHECK(MPI_Init(&argc, &argv));
329
- MPICHECK(MPI_Comm_rank(MPI_COMM_WORLD, &myRank));
330
- MPICHECK(MPI_Comm_size(MPI_COMM_WORLD, &nRanks));
331
- CUDACHECK(cudaSetDevice(myRank));
332
- ncclUniqueId id;
333
- ncclComm_t comm;
334
- if (myRank == 0) ncclGetUniqueId(&id);
335
- MPICHECK(MPI_Bcast(static_cast<void*>(&id), sizeof(id), MPI_BYTE, 0,
336
- MPI_COMM_WORLD));
337
- NCCLCHECK(ncclCommInitRank(&comm, nRanks, id, myRank));
338
-
339
- bool performance_test = true;
340
- cudaProfilerStart();
341
-// Uncomment to scan through different block size configs.
342
-// for (int threads : {256, 512, 1024}) {
343
-// for (int block_limit = 16; block_limit < 112; block_limit += 4) {
344
-// run<half>(myRank, nRanks, comm, threads, block_limit, 1024 * 1024,
345
-// performance_test);
346
-// }
347
-// }
348
-#ifdef USE_ROCM
349
- const int block_limit = 16;
350
-#else
351
- const int block_limit = 36;
352
-#endif
353
- // Scan through different sizes to test performance.
354
- for (int sz = 512; sz <= (8 << 20); sz *= 2) {
355
- run<half>(myRank, nRanks, comm, 512, 36, sz + 8 * 47, performance_test);
356
- }
357
-
358
- cudaProfilerStop();
359
- MPICHECK(MPI_Finalize());
360
- return EXIT_SUCCESS;
361
-}
\ No newline at end of file
csrc/ops.h
-6
@@ -41,12 +41,6 @@ void gelu_fast(torch::Tensor& out, torch::Tensor& input);
41
42
void gelu_quick(torch::Tensor& out, torch::Tensor& input);
43
44
-void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
45
- torch::Tensor const& q_pe,
46
- torch::Tensor const& kv_c_and_k_pe_cache,
47
- torch::Tensor const& seq_lens,
48
- torch::Tensor const& page_table, double scale);
49
-
44
void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
45
torch::Tensor const& scale,
46
std::optional<torch::Tensor> const& azp);
vllm/_custom_ops.py
-64
@@ -2962,70 +2962,6 @@ def qr_max_size() -> int:
2962
return torch.ops._C_custom_ar.qr_max_size()
2963
2964
2965
-def get_flash_mla_metadata(
2966
- cache_seqlens: torch.Tensor,
2967
- num_heads_per_head_k: int,
2968
- num_heads_k: int,
2969
-) -> tuple[torch.Tensor, torch.Tensor]:
2970
- """
2971
- Arguments:
2972
- cache_seqlens: (batch_size), dtype torch.int32.
2973
- num_heads_per_head_k: Equals to seq_len_q * num_heads_q // num_heads_k.
2974
- num_heads_k: num_heads_k.
2975
-
2976
- Return:
2977
- tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32.
2978
- num_splits: (batch_size + 1), dtype torch.int32.
2979
- """
2980
- return torch.ops._C.get_flash_mla_metadata(
2981
- cache_seqlens, num_heads_per_head_k, num_heads_k
2982
- )
2983
-
2984
-
2985
-def flash_mla_with_kvcache(
2986
- q: torch.Tensor,
2987
- k_cache: torch.Tensor,
2988
- block_table: torch.Tensor,
2989
- cache_seqlens: torch.Tensor,
2990
- head_dim_v: int,
2991
- tile_scheduler_metadata: torch.Tensor,
2992
- num_splits: torch.Tensor,
2993
- softmax_scale: float | None = None,
2994
- causal: bool = False,
2995
-) -> tuple[torch.Tensor, torch.Tensor]:
2996
- """
2997
- Arguments:
2998
- q: (batch_size, seq_len_q, num_heads_q, head_dim).
2999
- k_cache: (num_blocks, page_block_size, num_heads_k, head_dim).
3000
- block_table: (batch_size, max_num_blocks_per_seq), torch.int32.
3001
- cache_seqlens: (batch_size), torch.int32.
3002
- head_dim_v: Head_dim of v.
3003
- tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), torch.int32, return by get_mla_metadata.
3004
- num_splits: (batch_size + 1), torch.int32, return by get_mla_metadata.
3005
- softmax_scale: float. The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim).
3006
- causal: bool. Whether to apply causal attention mask.
3007
-
3008
- Return:
3009
- out: (batch_size, seq_len_q, num_heads_q, head_dim_v).
3010
- softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32.
3011
- """
3012
- if softmax_scale is None:
3013
- softmax_scale = q.shape[-1] ** (-0.5)
3014
- out, softmax_lse = torch.ops._C.flash_mla_fwd_kvcache(
3015
- q,
3016
- k_cache,
3017
- None,
3018
- head_dim_v,
3019
- cache_seqlens,
3020
- block_table,
3021
- softmax_scale,
3022
- causal,
3023
- tile_scheduler_metadata,
3024
- num_splits,
3025
- )
3026
- return out, softmax_lse
3027
-
3028
-
2965
def sm100_cutlass_mla_decode(
2966
out: torch.Tensor,
2967
lse: torch.Tensor,