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 3, 16:30 UTC
[Compile] Fix compile warning with topk softplus sqrt (#41261)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
May 14, 2026 at 08:12 UTC
65485604960c3dc2ce9dcb6d7e02e65c5acc3321
1 file changed
+102
-101
csrc/moe/topk_softplus_sqrt_kernels.cu
+102
-101
@@ -298,130 +298,131 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
298
asm volatile("griddepcontrol.launch_dependents;");
299
#endif
300
return;
301
- }
302
-
301
+ } else {
302
#pragma unroll
304
- for (int ii = 0; ii < VPT; ++ii) {
305
- float val = row_chunk[ii];
306
- float val_b = val * beta;
307
- // Compute softplus: log(1 + exp(val)) with numerical stability
308
- // When val > threshold, softplus(x) ≈ x to avoid exp overflow
309
- val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
310
- val = sqrtf(val);
311
- if (correction_bias) {
312
- const int group_id = ii / ELTS_PER_LDG;
313
- const int local_id = ii % ELTS_PER_LDG;
314
- const int expert_idx = first_elt_read_by_thread +
315
- group_id * THREADS_PER_ROW * ELTS_PER_LDG +
316
- local_id;
317
- val = val + correction_bias[expert_idx];
303
+ for (int ii = 0; ii < VPT; ++ii) {
304
+ float val = row_chunk[ii];
305
+ float val_b = val * beta;
306
+ // Compute softplus: log(1 + exp(val)) with numerical stability
307
+ // When val > threshold, softplus(x) ≈ x to avoid exp overflow
308
+ val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
309
+ val = sqrtf(val);
310
+ if (correction_bias) {
311
+ const int group_id = ii / ELTS_PER_LDG;
312
+ const int local_id = ii % ELTS_PER_LDG;
313
+ const int expert_idx = first_elt_read_by_thread +
314
+ group_id * THREADS_PER_ROW * ELTS_PER_LDG +
315
+ local_id;
316
+ val = val + correction_bias[expert_idx];
317
+ }
318
+ row_chunk[ii] = val;
319
}
319
- row_chunk[ii] = val;
320
- }
320
322
- // Original TopK path: find top-k experts by score
323
- // Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to find
324
- // the topk elements in each row, along with the max index.
325
- int start_col = first_elt_read_by_thread;
326
- static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
327
-
328
- float selected_sum = 0.f;
329
- for (int k_idx = 0; k_idx < k; ++k_idx) {
330
- // First, each thread does the local argmax
331
- float max_val = row_chunk[0];
332
- int expert = start_col;
321
+ // Original TopK path: find top-k experts by score
322
+ // Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to
323
+ // find the topk elements in each row, along with the max index.
324
+ int start_col = first_elt_read_by_thread;
325
+ static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
326
+
327
+ float selected_sum = 0.f;
328
+ for (int k_idx = 0; k_idx < k; ++k_idx) {
329
+ // First, each thread does the local argmax
330
+ float max_val = row_chunk[0];
331
+ int expert = start_col;
332
#pragma unroll
334
- for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
335
- ++ldg, col += COLS_PER_GROUP_LDG) {
333
+ for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
334
+ ++ldg, col += COLS_PER_GROUP_LDG) {
335
#pragma unroll
337
- for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
338
- float val = row_chunk[ldg * ELTS_PER_LDG + ii];
339
-
340
- // No check on the experts here since columns with the smallest index
341
- // are processed first and only updated if > (not >=)
342
- if (val > max_val) {
343
- max_val = val;
344
- expert = col + ii;
336
+ for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
337
+ float val = row_chunk[ldg * ELTS_PER_LDG + ii];
338
+
339
+ // No check on the experts here since columns with the smallest index
340
+ // are processed first and only updated if > (not >=)
341
+ if (val > max_val) {
342
+ max_val = val;
343
+ expert = col + ii;
344
+ }
345
}
346
}
347
- }
347
348
// Now, we perform the argmax reduce. We use the butterfly pattern so threads
349
// reach consensus about the max. This will be useful for K > 1 so that the
350
// threads can agree on "who" had the max value. That thread can then blank out
351
// their max with -inf and the warp can run more iterations...
352
#pragma unroll
354
- for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
355
- float other_max =
356
- VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
357
- int other_expert =
358
- VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
359
-
360
- // We want lower indices to "win" in every thread so we break ties this
361
- // way
362
- if (other_max > max_val ||
363
- (other_max == max_val && other_expert < expert)) {
364
- max_val = other_max;
365
- expert = other_expert;
353
+ for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
354
+ float other_max =
355
+ VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
356
+ int other_expert =
357
+ VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
358
+
359
+ // We want lower indices to "win" in every thread so we break ties this
360
+ // way
361
+ if (other_max > max_val ||
362
+ (other_max == max_val && other_expert < expert)) {
363
+ max_val = other_max;
364
+ expert = other_expert;
365
+ }
366
}
367
- }
367
369
- // Write the max for this k iteration to global memory.
370
- if (thread_group_idx == 0) {
371
- // Add a guard to ignore experts not included by this node
372
- const bool node_uses_expert =
373
- expert >= start_expert && expert < end_expert;
374
- const bool should_process_row = row_is_active && node_uses_expert;
375
-
376
- // The lead thread from each sub-group will write out the final results to
377
- // global memory. (This will be a single) thread per row of the
378
- // input/output matrices.
379
- const int idx = k * thread_row + k_idx;
380
- if (correction_bias != nullptr) {
381
- max_val -= correction_bias[expert];
382
- }
383
- output[idx] = max_val;
384
- indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
385
- source_rows[idx] = k_idx * num_rows + thread_row;
386
- if (renormalize) {
387
- selected_sum += max_val;
368
+ // Write the max for this k iteration to global memory.
369
+ if (thread_group_idx == 0) {
370
+ // Add a guard to ignore experts not included by this node
371
+ const bool node_uses_expert =
372
+ expert >= start_expert && expert < end_expert;
373
+ const bool should_process_row = row_is_active && node_uses_expert;
374
+
375
+ // The lead thread from each sub-group will write out the final results
376
+ // to global memory. (This will be a single) thread per row of the
377
+ // input/output matrices.
378
+ const int idx = k * thread_row + k_idx;
379
+ if (correction_bias != nullptr) {
380
+ max_val -= correction_bias[expert];
381
+ }
382
+ output[idx] = max_val;
383
+ indices[idx] =
384
+ should_process_row ? (expert - start_expert) : NUM_EXPERTS;
385
+ source_rows[idx] = k_idx * num_rows + thread_row;
386
+ if (renormalize) {
387
+ selected_sum += max_val;
388
+ }
389
}
389
- }
390
391
- // Finally, we clear the value in the thread with the current max if there
392
- // is another iteration to run.
393
- if (k_idx + 1 < k) {
394
- const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
395
- const int thread_to_clear_in_group =
396
- (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
397
-
398
- // Only the thread in the group which produced the max will reset the
399
- // "winning" value to -inf.
400
- if (thread_group_idx == thread_to_clear_in_group) {
401
- const int offset_for_expert = expert % ELTS_PER_LDG;
402
- // Safe to set to any negative value since row_chunk values must be
403
- // between 0 and 1.
404
- row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
405
- -10000.f;
391
+ // Finally, we clear the value in the thread with the current max if there
392
+ // is another iteration to run.
393
+ if (k_idx + 1 < k) {
394
+ const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
395
+ const int thread_to_clear_in_group =
396
+ (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
397
+
398
+ // Only the thread in the group which produced the max will reset the
399
+ // "winning" value to -inf.
400
+ if (thread_group_idx == thread_to_clear_in_group) {
401
+ const int offset_for_expert = expert % ELTS_PER_LDG;
402
+ // Safe to set to any negative value since row_chunk values must be
403
+ // between 0 and 1.
404
+ row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
405
+ -10000.f;
406
+ }
407
}
408
}
408
- }
409
410
- // Apply renormalization and routed scaling factor to final weights.
411
- if (thread_group_idx == 0) {
412
- float scale = static_cast<float>(routed_scaling_factor);
413
- if (renormalize) {
414
- const float denom = selected_sum > 0.f ? selected_sum : 1.f;
415
- scale /= denom;
416
- }
417
- for (int k_idx = 0; k_idx < k; ++k_idx) {
418
- const int idx = k * thread_row + k_idx;
419
- output[idx] = output[idx] * scale;
410
+ // Apply renormalization and routed scaling factor to final weights.
411
+ if (thread_group_idx == 0) {
412
+ float scale = static_cast<float>(routed_scaling_factor);
413
+ if (renormalize) {
414
+ const float denom = selected_sum > 0.f ? selected_sum : 1.f;
415
+ scale /= denom;
416
+ }
417
+ for (int k_idx = 0; k_idx < k; ++k_idx) {
418
+ const int idx = k * thread_row + k_idx;
419
+ output[idx] = output[idx] * scale;
420
+ }
421
}
421
- }
422
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
423
- asm volatile("griddepcontrol.launch_dependents;");
423
+ asm volatile("griddepcontrol.launch_dependents;");
424
#endif
425
+ }
426
}
427
428
namespace detail {