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 14, 05:30 UTC
Silu v2 (#25074)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: elvircrn <elvircrn@gmail.com> Signed-off-by: Elvir Crnčević <elvircrn@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com> Co-authored-by: Varun Sundar Rabindranath <varunsundar08@gmail.com>
Elvir Crnčević committed
Oct 10, 2025 at 17:19 UTC
7b03584de8819a870644bc853cf24cd2ff8a9f64
6 files changed
+512
-398
benchmarks/kernels/benchmark_silu_mul_fp8_quant.py
+168
-123
@@ -1,5 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""
5
+Comprehensive 3-way SiLU Benchmark Suite
6
+
7
+This benchmark compares three SiLU implementations:
8
+1. SiLU V2 (CUDA) - Optimized CUDA kernel implementation
9
+2. Triton Kernel - Triton-based implementation
10
+
11
+The suite generates detailed performance comparisons including:
12
+- Memory bandwidth utilization
13
+- Speedup ratios (baseline vs optimized implementations)
14
+- Performance across different expert configurations and token distributions
15
+"""
16
+
17
from collections.abc import Callable
18
19
import matplotlib.pyplot as plt
@@ -7,7 +21,7 @@ import numpy as np
21
import torch
22
23
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
10
- silu_mul_fp8_quant_deep_gemm_cuda,
24
+ persistent_masked_m_silu_mul_quant,
25
)
26
from vllm.platforms import current_platform
27
from vllm.triton_utils import tl, triton
@@ -94,6 +108,7 @@ def silu_mul_fp8_quant_deep_gemm_triton(
108
num_parallel_tokens,
109
group_size: int = 128,
110
eps: float = 1e-10,
111
+ expert_offsets: torch.Tensor = None,
112
) -> tuple[torch.Tensor, torch.Tensor]:
113
"""Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales
114
@@ -174,7 +189,7 @@ def silu_mul_fp8_quant_deep_gemm_triton(
189
190
191
# Parse generation strategies
177
-strategies = ["uniform", "max_t", "first_t"]
192
+strategies = ["random_imbalanced", "uniform", "max_t"]
193
194
195
def benchmark(
@@ -195,15 +210,27 @@ def benchmark(
210
current_platform.seed_everything(42 + seed_offset)
211
y = torch.rand((E, T, 2 * H), dtype=torch.bfloat16, device="cuda").contiguous()
212
198
- if gen_strategy == "uniform":
199
- r = torch.rand(size=(E,), device="cuda")
213
+ if gen_strategy == "random_imbalanced":
214
+
215
+ def generate_expert_loads(n_e, total_tokens, ratio, device="cuda"):
216
+ mean = total_tokens // n_e
217
+ min_max = mean // ratio
218
+ e = torch.ones(size=(E,), dtype=torch.int64, device=device) * mean
219
+ e[0] = min_max
220
+ r = torch.rand(size=(E - 1,))
221
+ r /= r.sum()
222
+ r *= total_tokens - min_max
223
+ r = r.round().long()
224
+ e[1:] = r.to(device=device)
225
+ return e
226
+
227
+ tokens_per_expert = generate_expert_loads(E, total_tokens, 0.7, "cuda")
228
+ elif gen_strategy == "uniform":
229
+ r = torch.rand(size=(E,))
230
r /= r.sum()
231
r *= total_tokens
202
- tokens_per_expert = r.int()
203
- tokens_per_expert = torch.minimum(
204
- tokens_per_expert,
205
- torch.ones((E,), device=r.device, dtype=torch.int) * T,
206
- )
232
+ r = r.round().long()
233
+ tokens_per_expert = r
234
elif gen_strategy == "max_t":
235
tokens_per_expert = torch.empty(size=(E,), dtype=torch.int32, device="cuda")
236
tokens_per_expert.fill_(total_tokens / E)
@@ -281,40 +308,34 @@ def benchmark(
308
309
310
def create_comparison_plot(
284
- ratio, cuda_times, baseline_times, config_labels, strategy_name, id
311
+ ratios, silu_v2_times, triton_times, config_labels, strategy_name, id
312
):
286
- """Create a comparison plot for a specific generation strategy"""
287
- fig, ax = plt.subplots(1, 1, figsize=(16, 6))
313
+ fig, ax = plt.subplots(1, 1, figsize=(18, 6))
314
315
# Configure x-axis positions
316
x = np.arange(len(config_labels))
291
- width = 0.35
317
+ width = 0.25
318
319
# Execution Time plot (lower is better)
320
+ ax.bar(x, silu_v2_times, width, label="SiLU V2 (CUDA)", alpha=0.8, color="blue")
321
ax.bar(
295
- x - width / 2, cuda_times, width, label="CUDA Kernel", alpha=0.8, color="blue"
296
- )
297
- ax.bar(
298
- x + width / 2,
299
- baseline_times,
300
- width,
301
- label="Baseline",
302
- alpha=0.8,
303
- color="orange",
322
+ x + width, triton_times, width, label="Triton Kernel", alpha=0.8, color="green"
323
)
324
306
- # Add speedup labels over each bar pair
325
+ # Add speedup labels over each bar trio
326
for i in range(len(x)):
308
- speedup = ratio[i]
309
- max_height = max(cuda_times[i], baseline_times[i])
327
+ triton_v2_speedup = ratios[i][1] # triton/v2
328
+ max_height = max(silu_v2_times[i], triton_times[i])
329
+
330
+ # Triton/V2 speedup
331
ax.text(
311
- x[i],
332
+ x[i] + width / 2,
333
max_height + max_height * 0.02,
313
- f"{speedup:.2f}x",
334
+ f"{triton_v2_speedup:.2f}x",
335
ha="center",
336
va="bottom",
337
fontweight="bold",
317
- fontsize=9,
338
+ fontsize=8,
339
)
340
341
ax.set_xlabel("Configuration")
@@ -332,56 +353,75 @@ def create_comparison_plot(
353
354
355
def create_combined_plot(all_results):
335
- """Create a combined plot with all strategies in one PNG"""
356
num_strategies = len(all_results)
337
- fig, axes = plt.subplots(num_strategies, 1, figsize=(20, 6 * num_strategies))
357
+ fig, axes = plt.subplots(num_strategies, 1, figsize=(22, 7 * num_strategies))
358
359
if num_strategies == 1:
360
axes = [axes]
361
362
for idx, (
363
strategy_name,
344
- ratio,
345
- cuda_times,
346
- baseline_times,
364
+ all_ratios,
365
+ all_silu_v2_results,
366
+ all_triton_results,
367
config_labels,
368
+ config_x_axis,
369
) in enumerate(all_results):
370
ax = axes[idx]
371
372
+ # Flatten the nested results to get bandwidth percentages for plotting
373
+ silu_v2_bandwidths = []
374
+ triton_bandwidths = []
375
+ flat_ratios = []
376
+
377
+ for config_results in all_silu_v2_results:
378
+ for result in config_results:
379
+ silu_v2_bandwidths.append(result[3]) # bandwidth percentage
380
+
381
+ for config_results in all_triton_results:
382
+ for result in config_results:
383
+ triton_bandwidths.append(result[3]) # bandwidth percentage
384
+
385
+ for config_ratios in all_ratios:
386
+ for ratio in config_ratios:
387
+ flat_ratios.append(ratio)
388
+
389
# Configure x-axis positions
390
x = np.arange(len(config_labels))
353
- width = 0.35
391
+ width = 0.25
392
355
- # Execution Time plot (lower is better)
393
+ # Bandwidth utilization plot (higher is better)
394
ax.bar(
357
- x - width / 2,
358
- cuda_times,
395
+ x,
396
+ silu_v2_bandwidths,
397
width,
360
- label="CUDA Kernel",
398
+ label="SiLU V2 (CUDA)",
399
alpha=0.8,
400
color="blue",
401
)
402
ax.bar(
365
- x + width / 2,
366
- baseline_times,
403
+ x + width,
404
+ triton_bandwidths,
405
width,
368
- label="Baseline",
406
+ label="Triton Kernel",
407
alpha=0.8,
370
- color="orange",
408
+ color="green",
409
)
410
373
- # Add speedup labels over each bar pair
411
+ # Add speedup labels over each bar trio
412
for i in range(len(x)):
375
- speedup = ratio[i]
376
- max_height = max(cuda_times[i], baseline_times[i])
413
+ triton_v2_speedup = flat_ratios[i] # triton/v2
414
+ max_height = max(silu_v2_bandwidths[i], triton_bandwidths[i])
415
+
416
+ # Triton/V2 speedup
417
ax.text(
378
- x[i],
418
+ x[i] + width / 2,
419
max_height + max_height * 0.02,
380
- f"{speedup:.2f}x",
420
+ f"{triton_v2_speedup:.2f}x",
421
ha="center",
422
va="bottom",
423
fontweight="bold",
384
- fontsize=9,
424
+ fontsize=8,
425
)
426
427
ax.set_xlabel("Configuration")
@@ -395,7 +435,7 @@ def create_combined_plot(all_results):
435
ax.grid(True, alpha=0.3)
436
437
plt.tight_layout()
398
- filename = "../../silu_bench/silu_benchmark_combined.png"
438
+ filename = "silu_benchmark_combined_3way.png"
439
plt.savefig(filename, dpi=300, bbox_inches="tight")
440
plt.show()
441
@@ -405,7 +445,9 @@ def create_combined_plot(all_results):
445
outer_dim = 7168
446
configs = [
447
# DeepSeekV3 Configs
448
+ # (1, 56, 7168),
449
(8, 1024, 7168),
450
+ # (32, 56, 7168),
451
# DeepSeekV3 Configs
452
(32, 1024, 7168),
453
# DeepSeekV3 Configs
@@ -417,6 +459,7 @@ num_warmups = 20
459
460
strategy_descriptions = {
461
"uniform": "Uniform Random",
462
+ "random_imbalanced": "Imbalanced Random",
463
"max_t": "Even Assignment",
464
"first_t": "experts[0] = T, experts[1:] = 0",
465
}
@@ -433,28 +476,31 @@ for id, strategy in enumerate(strategies):
476
print(f"Testing strategy: {strategy_descriptions[strategy]}")
477
print(f"{'=' * 60}")
478
436
- # Collect benchmark data for both algorithms
479
+ # Collect benchmark data for all three algorithms
480
config_labels = []
481
config_x_axis = []
439
- all_cuda_results = []
440
- all_baseline_results = []
482
+ all_silu_v2_results = []
483
+ all_triton_results = []
484
all_ratios = []
485
486
for E, T, H in configs:
444
- total_tokens_config = [8 * E, 16 * E, 32 * E, 64 * E, 128 * E, 256 * E]
487
+ total_tokens_config = []
488
+ for i in [8, 16, 32, 64, 128, 256, 512]:
489
+ if i <= T:
490
+ total_tokens_config.append(i * E)
491
config_x_axis.append(total_tokens_config)
492
447
- cuda_results = []
448
- baseline_results = []
493
+ silu_v2_results = []
494
+ triton_results = []
495
ratios = []
496
497
for total_tokens in total_tokens_config:
498
config_label = f"E={E},T={T},H={H},TT={total_tokens}"
499
config_labels.append(config_label)
500
455
- # CUDA kernel results
456
- time_ms_cuda, gflops, gbps, perc = benchmark(
457
- silu_mul_fp8_quant_deep_gemm_cuda,
501
+ # SiLU V2 (CUDA kernel) results
502
+ time_ms_silu_v2, gflops, gbps, perc = benchmark(
503
+ persistent_masked_m_silu_mul_quant,
504
E,
505
T,
506
H,
@@ -463,9 +509,9 @@ for id, strategy in enumerate(strategies):
509
num_warmups=num_warmups,
510
gen_strategy=strategy,
511
)
466
- cuda_results.append((time_ms_cuda, gflops, gbps, perc))
512
+ silu_v2_results.append((time_ms_silu_v2, gflops, gbps, perc))
513
468
- # Baseline results
514
+ # Triton kernel results
515
time_ms_triton, gflops, gbps, perc = benchmark(
516
silu_mul_fp8_quant_deep_gemm_triton,
517
E,
@@ -476,12 +522,20 @@ for id, strategy in enumerate(strategies):
522
num_warmups=num_warmups,
523
gen_strategy=strategy,
524
)
479
- baseline_results.append((time_ms_triton, gflops, gbps, perc))
480
- ratios.append(time_ms_triton / time_ms_cuda)
525
+ triton_results.append((time_ms_triton, gflops, gbps, perc))
526
482
- print(f"Completed: {config_label}")
483
- all_cuda_results.append(cuda_results)
484
- all_baseline_results.append(baseline_results)
527
+ # Calculate speedup ratios (triton baseline / implementation)
528
+ triton_v2_ratio = time_ms_triton / time_ms_silu_v2
529
+ ratios.append(triton_v2_ratio)
530
+
531
+ print(
532
+ f"Completed: {config_label}:"
533
+ f" V2: {time_ms_silu_v2:.3f}ms,"
534
+ f" Triton: {time_ms_triton:.3f}ms"
535
+ )
536
+
537
+ all_silu_v2_results.append(silu_v2_results)
538
+ all_triton_results.append(triton_results)
539
all_ratios.append(ratios)
540
541
# Store results for combined plotting
@@ -489,8 +543,8 @@ for id, strategy in enumerate(strategies):
543
(
544
strategy_descriptions[strategy],
545
all_ratios,
492
- all_cuda_results,
493
- all_baseline_results,
546
+ all_silu_v2_results,
547
+ all_triton_results,
548
config_labels,
549
config_x_axis,
550
)
@@ -498,15 +552,18 @@ for id, strategy in enumerate(strategies):
552
553
# Print summary table for this strategy
554
print(f"\nSummary Table - {strategy_descriptions[strategy]}:")
501
- print(f"{'Config':<20} {'CUDA Time(ms)':<12} {'Base Time(ms)':<12} {'Speedup':<8}")
502
- print("-" * 60)
555
+ print(f" {'V2 Time(ms)':<12} {'Triton Time(ms)':<14} {'Triton/V2':<10}")
556
+ print("-" * 90)
557
558
for i, (E, T, H) in enumerate(configs):
505
- speedup = baseline_results[i][0] / cuda_results[i][0]
559
+ # Get the first result for each config (simplifying for summary)
560
+ v2_time = silu_v2_results[i][0]
561
+ triton_time = triton_results[i][0]
562
+ triton_v2_speedup = triton_time / v2_time
563
config_label = f"E={E:3d},T={T:4d},H={H:4d}"
564
print(
508
- f"{config_label:<20} {cuda_results[i][0]:8.5f} "
509
- f"{baseline_results[i][0]:8.5f} {speedup:6.2f}x"
565
+ f"{config_label:<20} {v2_time:8.5f} {triton_time:10.5f} "
566
+ f"{triton_v2_speedup:8.2f}x"
567
)
568
569
@@ -514,15 +571,14 @@ def create_total_tokens_plot(all_results):
571
num_strategies = len(all_results)
572
num_configs = len(configs)
573
517
- # Create side-by-side subplots: 2 columns for speedup and bandwidth percentage
574
fig, axs = plt.subplots(
519
- num_strategies, num_configs * 2, figsize=(28, 6 * num_strategies)
575
+ num_strategies, num_configs * 2, figsize=(32, 8 * num_strategies)
576
)
577
578
# Add main title to the entire figure
579
fig.suptitle(
524
- "Performance Analysis: Speedup vs Bandwidth Utilization (Triton & CUDA)",
525
- fontsize=16,
580
+ "Performance Analysis: Speedup vs Bandwidth Utilization (SiLU V2, and Triton)",
581
+ fontsize=18,
582
fontweight="bold",
583
y=0.98,
584
)
@@ -539,8 +595,8 @@ def create_total_tokens_plot(all_results):
595
(
596
strategy_name,
597
all_ratios,
542
- all_cuda_results,
543
- all_baseline_results,
598
+ all_silu_v2_results,
599
+ all_triton_results,
600
config_labels,
601
config_x_axis,
602
) = result
@@ -555,42 +611,54 @@ def create_total_tokens_plot(all_results):
611
ratios = all_ratios[config_idx]
612
total_tokens_values = config_x_axis[config_idx]
613
558
- # Extract CUDA and Triton bandwidth percentages
559
- cuda_bandwidth_percentages = [
560
- result[3] for result in all_cuda_results[config_idx]
614
+ # Extract speedup ratios
615
+ triton_v2_ratios = [ratio for ratio in ratios]
616
+
617
+ # Extract bandwidth percentages for all implementations
618
+ v2_bandwidth_percentages = [
619
+ result[3] for result in all_silu_v2_results[config_idx]
620
]
621
triton_bandwidth_percentages = [
563
- result[3] for result in all_baseline_results[config_idx]
622
+ result[3] for result in all_triton_results[config_idx]
623
]
624
625
# Plot speedup ratios vs total tokens (left plot)
626
ax_speedup.plot(
568
- total_tokens_values, ratios, "bo-", linewidth=3, markersize=8
627
+ total_tokens_values,
628
+ triton_v2_ratios,
629
+ "go-",
630
+ linewidth=3,
631
+ markersize=8,
632
+ label="Triton/V2 Speedup",
633
)
634
ax_speedup.set_title(
571
- f"{strategy_name}\nSpeedup (CUDA/Triton)\nE={E}, T={T}, H={H}",
635
+ f"{strategy_name}\nSpeedup vs Baseline (Triton)\nE={E}, T={T}, H={H}",
636
fontsize=12,
637
fontweight="bold",
638
)
639
ax_speedup.set_xlabel("Total Tokens", fontweight="bold", fontsize=11)
640
ax_speedup.set_ylabel("Speedup Ratio", fontweight="bold", fontsize=11)
641
+ ax_speedup.legend(prop={"weight": "bold"})
642
ax_speedup.grid(True, alpha=0.3)
643
644
+ # Plot bandwidth utilization (right plot)
645
ax_bandwidth.plot(
646
total_tokens_values,
581
- cuda_bandwidth_percentages,
582
- "ro-",
647
+ v2_bandwidth_percentages,
648
+ "o-",
649
linewidth=3,
650
markersize=8,
585
- label="CUDA",
651
+ label="SiLU V2",
652
+ color="blue",
653
)
654
ax_bandwidth.plot(
655
total_tokens_values,
656
triton_bandwidth_percentages,
590
- "go-",
657
+ "o-",
658
linewidth=3,
659
markersize=8,
660
label="Triton",
661
+ color="green",
662
)
663
ax_bandwidth.set_title(
664
f"{strategy_name}\nBandwidth Utilization (Hopper)\nE={E}, T={T}, H={H}",
@@ -618,38 +686,12 @@ def create_total_tokens_plot(all_results):
686
for label in ax.get_xticklabels() + ax.get_yticklabels():
687
label.set_fontweight("bold")
688
621
- # Add value labels on speedup points
622
- for x, y in zip(total_tokens_values, ratios):
689
+ # Add value labels on Triton/V2 speedup points
690
+ for x, y in zip(total_tokens_values, triton_v2_ratios):
691
ax_speedup.annotate(
692
f"{y:.2f}x",
693
(x, y),
694
textcoords="offset points",
627
- xytext=(0, 12),
628
- ha="center",
629
- fontsize=10,
630
- fontweight="bold",
631
- bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.7),
632
- )
633
-
634
- # Add value labels on CUDA bandwidth points
635
- for x, y in zip(total_tokens_values, cuda_bandwidth_percentages):
636
- ax_bandwidth.annotate(
637
- f"{y:.1f}%",
638
- (x, y),
639
- textcoords="offset points",
640
- xytext=(0, 12),
641
- ha="center",
642
- fontsize=9,
643
- fontweight="bold",
644
- bbox=dict(boxstyle="round,pad=0.2", facecolor="red", alpha=0.3),
645
- )
646
-
647
- # Add value labels on Triton bandwidth points
648
- for x, y in zip(total_tokens_values, triton_bandwidth_percentages):
649
- ax_bandwidth.annotate(
650
- f"{y:.1f}%",
651
- (x, y),
652
- textcoords="offset points",
695
xytext=(0, -15),
696
ha="center",
697
fontsize=9,
@@ -659,17 +701,20 @@ def create_total_tokens_plot(all_results):
701
702
plt.tight_layout()
703
plt.subplots_adjust(top=0.93) # Make room for main title
662
- filename = "silu_benchmark_total_tokens.png"
704
+ filename = "silu_benchmark_total_tokens_3way.png"
705
plt.savefig(filename, dpi=300, bbox_inches="tight")
706
plt.show()
707
708
return filename
709
710
669
-# Create combined plot with all strategies
670
-combined_plot_filename = create_total_tokens_plot(all_results)
711
+# Create comprehensive 3-way comparison plots
712
+combined_plot_filename = create_combined_plot(all_results)
713
+total_tokens_plot_filename = create_total_tokens_plot(all_results)
714
672
-print(f"\n{'=' * 60}")
673
-print("Benchmark Complete!")
674
-print(f"Generated combined plot: {combined_plot_filename}")
675
-print(f"{'=' * 60}")
715
+print(f"\n{'=' * 80}")
716
+print("3-Way Benchmark Suite Complete!")
717
+print(f"Generated combined comparison plot: {combined_plot_filename}")
718
+print(f"Generated total tokens analysis plot: {total_tokens_plot_filename}")
719
+print("Compared: SiLU V2 (CUDA), and Triton implementations")
720
+print(f"{'=' * 80}")
csrc/ops.h
+2
-2
@@ -138,12 +138,12 @@ void silu_and_mul_nvfp4_quant(torch::Tensor& out,
138
torch::Tensor& input,
139
torch::Tensor& input_global_scale);
140
#endif
141
-void silu_mul_fp8_quant_deep_gemm_cuda(
141
+void persistent_masked_m_silu_mul_quant(
142
const at::Tensor& input, // (E, T, 2*H)
143
const at::Tensor& counts, // (E)
144
at::Tensor& y_q, // (E, T, H) [OUT]
145
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
146
- int64_t group_size, bool use_ue8m0, int64_t num_parallel_tokens);
146
+ bool use_ue8m0);
147
148
void mul_and_silu(torch::Tensor& out, torch::Tensor& input);
149
csrc/quantization/activation_kernels.cu
+295
-233
@@ -114,13 +114,22 @@ __global__ void act_and_mul_quant_kernel(
114
}
115
116
__device__ __forceinline__ float silu(float x) {
117
- return (__fdividef(x, (1.f + expf(-x))));
117
+ return __fdividef(x, (1.f + expf(-x)));
118
}
119
120
__device__ __forceinline__ float2 silu2(float2 x) {
121
return make_float2(silu(x.x), silu(x.y));
122
}
123
124
+__device__ __forceinline__ __nv_bfloat162 silu2_v2(float2 x) {
125
+#ifndef USE_ROCM
126
+ return make_bfloat162(__float2bfloat16_rn(silu(x.x)),
127
+ __float2bfloat16_rn(silu(x.y)));
128
+#else
129
+ return __float22bfloat162_rn(make_float2(silu(x.x), silu(x.y)));
130
+#endif
131
+}
132
+
133
#ifndef USE_ROCM
134
__device__ __forceinline__ float warp_max(float v) {
135
static constexpr unsigned FULL_MASK = 0xffffffffu;
@@ -223,224 +232,308 @@ constexpr __nv_bfloat16 get_fp8_min() {
232
return __nv_bfloat16(__nv_bfloat16_raw{.x = 50032});
233
}
234
}
226
-#ifndef USE_ROCM
227
-template <typename fp8_type, int32_t NUM_WARPS, typename Idx_t,
228
- int NUM_PARALLEL_TOKENS, bool USE_UE8M0, int GROUP_SIZE = 128,
235
+
236
+template <typename Idx_t>
237
+__device__ __forceinline__ int warp_expert_search(
238
+ int idx, int n, const Idx_t* __restrict__ input, Idx_t val) {
239
+ const Idx_t* input_ptr = input + idx;
240
+ int base_offset = 0;
241
+
242
+ for (;;) {
243
+ bool move_on = (idx < n && *input_ptr <= val);
244
+
245
+ unsigned mask = __ballot_sync(0xffffffff, move_on);
246
+
247
+ if (mask != 0xffffffffu) {
248
+ int last_lane = 31 - __clz(mask);
249
+ return base_offset + last_lane;
250
+ }
251
+
252
+ input_ptr += 32;
253
+ base_offset += 32;
254
+ idx += 32;
255
+ }
256
+}
257
+
258
+template <int num_parallel_tokens>
259
+__device__ __forceinline__ void token_bounds(int32_t n_tokens,
260
+ int32_t worker_id,
261
+ int32_t& n_tokens_lower,
262
+ int32_t& n_tokens_upper) {
263
+ if (n_tokens < num_parallel_tokens && worker_id < n_tokens) {
264
+ if (worker_id >= num_parallel_tokens) return;
265
+ n_tokens_lower = worker_id;
266
+ n_tokens_upper = worker_id + 1;
267
+ } else {
268
+ int32_t chunk_size = n_tokens / num_parallel_tokens;
269
+ int32_t residual = n_tokens - chunk_size * num_parallel_tokens;
270
+ auto calc_id = [&](int32_t id) {
271
+ if (id < residual)
272
+ return min(n_tokens, id * (chunk_size + 1));
273
+ else
274
+ return min(n_tokens, id * chunk_size + residual);
275
+ };
276
+ n_tokens_lower = calc_id(worker_id);
277
+ n_tokens_upper = calc_id(worker_id + 1);
278
+ }
279
+}
280
+
281
+template <int BLOCK_COUNT, int SMEM_SIZE_BYTES_Y, typename fp8_type,
282
+ int THREADS, typename Idx_t, bool USE_UE8M0, int GROUP_SIZE = 128,
283
int NUM_STAGES = 3>
284
__global__ void silu_mul_fp8_quant_deep_gemm_kernel(
285
const __nv_bfloat16* __restrict__ _input, fp8_type* __restrict__ _y_q,
232
- float* __restrict__ _y_s, const int32_t* __restrict__ counts,
233
-
286
+ float* __restrict__ _y_s, const int32_t* __restrict__ tokens_per_expert,
287
// sizes
235
- int H, int G,
236
-
288
+ Idx_t E, Idx_t T, Idx_t H,
289
// strides (in elements)
290
Idx_t stride_i_e, Idx_t stride_i_t, Idx_t stride_i_h, Idx_t stride_yq_e,
291
Idx_t stride_yq_t, Idx_t stride_yq_h, Idx_t stride_ys_e, Idx_t stride_ys_t,
292
Idx_t stride_ys_g, Idx_t stride_counts_e) {
241
- static constexpr __nv_bfloat16 fp8_min = get_fp8_min<fp8_type>();
242
- static constexpr __nv_bfloat16 fp8_max = get_fp8_max<fp8_type>();
243
- // We assign EPS with its 16-bit unsigned counterpart to allow constexpr.
244
- static constexpr __nv_bfloat16 EPS = (__nv_bfloat16_raw{.x = 11996});
293
+#ifndef USE_ROCM
294
+ static constexpr int NUM_WARPS = THREADS / WARP_SIZE;
295
+
296
+ static constexpr int LOAD_STAGE_SIZE = 2 * GROUP_SIZE / 8;
297
+ static constexpr int LOAD_STAGE_MOD = NUM_STAGES * LOAD_STAGE_SIZE;
298
246
- // We pack 8 16-bit bfloat16 values into a 128-bit __int128_t.
247
- static constexpr int32_t BFLOAT16_PER_GROUP = 8;
299
+ static constexpr int COMPUTE_STAGE_SIZE = 2 * GROUP_SIZE / 4;
300
+ static constexpr int COMPUTE_STAGE_MOD = COMPUTE_STAGE_SIZE * NUM_STAGES;
301
249
- // We split the shared memory in half, corresponding to gate and up matrices:
250
- // [...gate_i, ...up_i] where 0 <= i < stages.
251
- static constexpr int32_t S_NUM_128 =
252
- 2u * (GROUP_SIZE / BFLOAT16_PER_GROUP) * NUM_WARPS * NUM_STAGES;
253
- static constexpr auto THREAD_COUNT = NUM_WARPS * WARP_SIZE;
254
- static constexpr int HALF_THREAD_COUNT = THREAD_COUNT / 2;
255
- static constexpr int32_t S_NUM_64 = S_NUM_128 * 2;
256
- __shared__ __int128_t __align__(16) s_buff_128[S_NUM_128];
302
+ extern __shared__ __align__(16) __int128_t smem_128[];
303
258
- const int32_t tid = threadIdx.x;
259
- const int32_t warp_id = tid / WARP_SIZE;
260
- const int32_t lane_id = tid % WARP_SIZE;
304
+ int* s_expert_offsets =
305
+ reinterpret_cast<int*>(smem_128 + (SMEM_SIZE_BYTES_Y / 16));
306
262
- auto s_buff_compute_32 = reinterpret_cast<__nv_bfloat162*>(s_buff_128);
307
+ static constexpr __nv_bfloat16 fp8_min = get_fp8_min<fp8_type>();
308
+ static constexpr __nv_bfloat16 fp8_max = get_fp8_max<fp8_type>();
309
+ // We assign EPS with it's 16-bit unsigned counterpart to allow constexpr.
310
+ static constexpr __nv_bfloat16 EPS = (__nv_bfloat16_raw{.x = 11996});
311
+ int tid = threadIdx.x;
312
+ int warp_id = tid >> 5;
313
+ int lane_id = tid & 0x1f;
314
+
315
+ int running_sum{};
316
+ if (!warp_id) {
317
+ for (int i = 0; i < E; i += WARP_SIZE) {
318
+ bool valid = (i + threadIdx.x) < E;
319
+ int value =
320
+ (valid ? tokens_per_expert[i + threadIdx.x * stride_counts_e] : 0) +
321
+ (!lane_id ? running_sum : 0);
322
+
323
+ for (int offset = 1; offset < 32; offset *= 2) {
324
+ int n = __shfl_up_sync(0xFFFFFFFFu, value, offset);
325
+ if (lane_id >= offset) value += n;
326
+ }
327
264
- // block handles one (expert e, group g)
265
- int32_t pid = blockIdx.x;
266
- int32_t e = pid / G;
267
- int32_t g = pid % G;
328
+ if (valid) {
329
+ s_expert_offsets[i + threadIdx.x + 1] = value;
330
+ }
331
269
- const int32_t n_tokens = counts[e * stride_counts_e];
332
+ running_sum = __shfl_sync(0xFFFFFFFFu, value, WARP_SIZE - 1);
333
+ }
334
271
- if (!n_tokens) {
272
- return; // Exit ASAP.
335
+ if (!lane_id) {
336
+ s_expert_offsets[0] = 0;
337
+ }
338
}
339
275
- const Idx_t stride_i_t_128 = stride_i_t / 8u;
340
+ __syncthreads();
341
+
342
+ int32_t total_tokens = s_expert_offsets[E];
343
277
- int32_t n_tokens_lower, n_tokens_upper;
344
+ const int warp_position_yq = warp_id * (H / NUM_WARPS);
345
+ const int warp_position_scales = warp_id * (H / (GROUP_SIZE * NUM_WARPS));
346
347
+ // A single block will handle tokens_per_block tokens.
348
// Each block i iterates over tokens of a slice of n_tokens =
349
// expert_counts[i], with the size of chunk being
350
// (n_tokens / NUM_PARALLEL_TOKENS) + residual, instead of
351
// updiv(n_tokens, NUM_PARALLEL_TOKENS) for better scheduling.
283
- if (n_tokens < NUM_PARALLEL_TOKENS && blockIdx.y < n_tokens) {
284
- // Specialize this, but can be likely fused.
285
- if (blockIdx.y >= NUM_PARALLEL_TOKENS) {
286
- return;
287
- }
288
- n_tokens_lower = blockIdx.y;
289
- n_tokens_upper = blockIdx.y + 1;
290
- } else {
291
- auto chunk_size = n_tokens / NUM_PARALLEL_TOKENS;
292
- auto residual = n_tokens - chunk_size * NUM_PARALLEL_TOKENS;
293
- auto calc_id = [&](int32_t id) {
294
- if (id < residual) {
295
- return min(n_tokens, id * (chunk_size + 1));
296
- } else {
297
- return min(n_tokens, id * chunk_size + residual);
298
- }
299
- };
300
- n_tokens_lower = calc_id(blockIdx.y);
301
- n_tokens_upper = calc_id(blockIdx.y + 1);
302
- }
352
304
- if (n_tokens_lower >= n_tokens_upper) {
353
+ // Each warp will get space to store its hidden dim for gate and up.
354
+ __int128_t* s_hidden_load = smem_128 + warp_id * ((2 * 128 / 8) * NUM_STAGES);
355
+ __int128_t* smem_load_ptr = s_hidden_load + lane_id;
356
+
357
+ const __nv_bfloat16 fp8_inv = __hdiv(__float2bfloat16(1.f), fp8_max);
358
+
359
+ int32_t compute_pipeline_offset_64 = 0;
360
+ int32_t load_stage_offset{};
361
+ const __nv_bfloat16 one_bf16 = __float2bfloat16_rn(1.f);
362
+
363
+ __int64_t* smem_compute_ptr = reinterpret_cast<__int64_t*>(smem_128) +
364
+ warp_id * (2 * (GROUP_SIZE / 4) * NUM_STAGES) +
365
+ lane_id;
366
+ __int64_t* s_gate64_ptr = smem_compute_ptr;
367
+ __int64_t* s_up64_ptr = smem_compute_ptr + GROUP_SIZE / 4;
368
+
369
+ int tokens_lower, tokens_upper;
370
+
371
+ token_bounds<BLOCK_COUNT>(total_tokens, blockIdx.x, tokens_lower,
372
+ tokens_upper);
373
+
374
+ Idx_t expert_id{}, expert_offset{}, next_expert_offset{};
375
+ int token_id = tokens_lower;
376
+ int32_t t_load{};
377
+
378
+ if (token_id < tokens_upper) {
379
+ expert_id = warp_expert_search<int>(lane_id, E, s_expert_offsets, token_id);
380
+ expert_offset = s_expert_offsets[expert_id];
381
+ next_expert_offset = s_expert_offsets[expert_id + 1];
382
+ } else {
383
+ // This thread block has no work to do.
384
return;
385
}
386
308
- // We do calculations here, using constexpr wherever possible.
309
- const Idx_t base_i = e * stride_i_e + NUM_WARPS * g * GROUP_SIZE * stride_i_h;
310
- const Idx_t base_ys = e * stride_ys_e + NUM_WARPS * g * stride_ys_g;
311
- const Idx_t base_yq =
312
- e * stride_yq_e + NUM_WARPS * g * GROUP_SIZE * stride_yq_h;
313
- Idx_t gate_off_128 = (base_i / static_cast<Idx_t>(8u));
314
- auto input_128_ptr = reinterpret_cast<const __int128_t*>(_input);
315
- auto gate_128_ptr = input_128_ptr + gate_off_128 + (tid % HALF_THREAD_COUNT) +
316
- stride_i_t_128 * n_tokens_lower;
317
- auto up_128_ptr = gate_128_ptr + (H * stride_i_h) / 8u;
318
- auto y_s_ptr =
319
- _y_s + base_ys + warp_id * stride_ys_g + n_tokens_lower * stride_ys_t;
320
- auto y_q_ptr = _y_q + base_yq + warp_id * GROUP_SIZE +
321
- stride_yq_t * n_tokens_lower + 4 * lane_id;
322
- int32_t t_load = n_tokens_lower, load_stage_id = 0;
323
- auto s_buff_gate_load_128 = s_buff_128 + (tid % HALF_THREAD_COUNT);
324
- auto s_buff_up_load_128 = s_buff_gate_load_128 + S_NUM_128 / 2u;
325
- int32_t stage_offset{};
326
-
327
- static constexpr int32_t LOAD_STAGE_SIZE = (NUM_WARPS * WARP_SIZE / 2);
328
- static constexpr int32_t LOAD_STAGE_MOD =
329
- NUM_STAGES * (NUM_WARPS * WARP_SIZE / 2);
330
-
331
- // Two halves of all threads in a block conduct global loads for gate and up,
332
- // repsectively.
387
+ int t_load_bound = H / (GROUP_SIZE * NUM_WARPS);
388
+
389
+ Idx_t base_i = ((expert_id * stride_i_e) / 8) +
390
+ (token_id - expert_offset) * stride_i_t / 8;
391
+ const Idx_t gate_warp_offset =
392
+ warp_id * ((stride_i_h * H) / (8 * NUM_WARPS)) + (lane_id & 0b1111);
393
+
394
+ const __int128_t* input_128_ptr =
395
+ reinterpret_cast<const __int128_t*>(_input) + gate_warp_offset +
396
+ ((lane_id < 16) ? 0 : ((H * stride_i_h) / 8));
397
+ __int128_t* load_ptr = const_cast<__int128_t*>(input_128_ptr + base_i);
398
+
399
+ auto token_offset = token_id - expert_offset;
400
+
401
auto load_and_advance_y_pred = [&] {
334
- if (t_load < n_tokens_upper) {
335
- auto s_gate_stage_128_staged_ptr = s_buff_gate_load_128 + stage_offset;
336
- auto s_up_stage_128_staged_ptr = s_buff_up_load_128 + stage_offset;
402
+ if (t_load < t_load_bound) {
403
+ // Here we are simply continuing to load data
404
+ // from the current token.
405
+ auto smem_load_ptr_staged = smem_load_ptr + load_stage_offset;
406
407
// It is very important that LOAD_STAGE_SIZE is constexpr to avoid
408
// unnecessary ALU ops.
340
- stage_offset += LOAD_STAGE_SIZE;
341
- stage_offset %= LOAD_STAGE_MOD;
409
+ load_stage_offset += LOAD_STAGE_SIZE;
410
+ load_stage_offset %= LOAD_STAGE_MOD;
411
343
- if (tid < HALF_THREAD_COUNT) {
344
- cp_async4(s_gate_stage_128_staged_ptr, gate_128_ptr);
345
- gate_128_ptr += stride_i_t_128;
412
+ cp_async4(smem_load_ptr_staged, load_ptr);
413
+ load_ptr += GROUP_SIZE / 8;
414
+ ++t_load;
415
+ } else if (token_id + 1 < tokens_upper) {
416
+ // We loaded everything from the current token, let's move on
417
+ // to the next one, and we checked that we have more tokens to load.
418
+ ++token_id;
419
+ t_load = 0;
420
+ if (token_id >= next_expert_offset) {
421
+ // We need to find the next expert.
422
+ do {
423
+ // This is a loop because it's possible
424
+ // that some experts are assigned 0 tokens.
425
+ // NOTE: We are guaranteed that there's at least
426
+ // one more token left so we don't have to check for
427
+ // expert_id bounds.
428
+ ++expert_id;
429
+ // This skips 1 memory read.
430
+ expert_offset = next_expert_offset;
431
+ next_expert_offset = s_expert_offsets[expert_id + 1];
432
+ } while (next_expert_offset == expert_offset);
433
+
434
+ base_i = expert_id * (stride_i_e / 8);
435
+ token_offset = 0;
436
+ load_ptr = const_cast<__int128_t*>(input_128_ptr + base_i);
437
} else {
347
- cp_async4(s_up_stage_128_staged_ptr, up_128_ptr);
348
- up_128_ptr += stride_i_t_128;
438
+ // We remain within the same expert, so just
439
+ // move by H/4 __int128_t (2 * H/8).
440
+ base_i += stride_yq_t / 4;
441
+ token_offset++;
442
}
443
+
444
+ load_ptr = const_cast<__int128_t*>(input_128_ptr + base_i);
445
+
446
+ auto smem_load_ptr_staged = smem_load_ptr + load_stage_offset;
447
+
448
+ // It is very important that LOAD_STAGE_SIZE is constexpr to avoid
449
+ // unnecessary ALU ops.
450
+ load_stage_offset += LOAD_STAGE_SIZE;
451
+ load_stage_offset %= LOAD_STAGE_MOD;
452
+
453
+ cp_async4(smem_load_ptr_staged, load_ptr);
454
+ load_ptr += GROUP_SIZE / 8;
455
++t_load;
351
- ++load_stage_id;
456
}
457
// We fence even if there is nothing to load to simplify pipelining.
458
cp_async_fence();
459
};
460
461
+ // We need to warm-up the pipeline.
462
#pragma unroll
463
for (int i = 0; i < NUM_STAGES - 1; i++) {
464
load_and_advance_y_pred();
465
}
466
362
- __int64_t* s_gate_ptr = reinterpret_cast<__int64_t*>(
363
- s_buff_compute_32 + warp_id * (GROUP_SIZE / 2)) +
364
- lane_id;
365
- __int64_t* s_up_ptr = s_gate_ptr + S_NUM_64 / 2;
366
-
367
- static constexpr int32_t STAGE_SIZE = (GROUP_SIZE * NUM_WARPS) / 4u;
368
- static constexpr int32_t STAGE_MOD = STAGE_SIZE * NUM_STAGES;
369
-
370
- int32_t compute_pipeline_offset_64 = 0;
371
-
372
- for (int32_t t = n_tokens_lower; t < n_tokens_upper; ++t) {
373
- __nv_bfloat162 results_bf162[2];
467
+ __nv_fp8x4_e4m3* y_q_base_ptr =
468
+ reinterpret_cast<__nv_fp8x4_e4m3*>(_y_q) + lane_id;
469
+ auto y_scale_base_ptr = _y_s + warp_position_scales * stride_ys_g;
470
375
- cp_async_wait<NUM_STAGES - 2>();
376
- __syncthreads();
471
+ for (auto j = tokens_lower; j < tokens_upper; j++) {
472
+ const Idx_t base_ys = expert_id * stride_ys_e;
473
+ auto y_s_ptr = y_scale_base_ptr + base_ys + token_offset * stride_ys_t;
474
+ __nv_fp8x4_e4m3* y_q_ptr =
475
+ y_q_base_ptr + (expert_id * stride_yq_e + token_offset * stride_yq_t +
476
+ warp_position_yq * stride_yq_h) /
477
+ 4;
478
+ const int COMPUTE_LIMIT = H / (GROUP_SIZE * NUM_WARPS);
479
378
- // We double-buffer pipelined loads so that the next load will
379
- // concurrently run with compute without overwrites.
380
- load_and_advance_y_pred();
480
+ for (int i = 0; i < COMPUTE_LIMIT; i++) {
481
+ cp_async_wait<NUM_STAGES - 2>();
482
+ __syncthreads();
483
+ load_and_advance_y_pred();
484
382
- auto s_gate_compute_64 = s_gate_ptr + compute_pipeline_offset_64;
383
- auto s_up_compute_64 = s_up_ptr + compute_pipeline_offset_64;
485
+ __int64_t* gate64_ptr = s_gate64_ptr + compute_pipeline_offset_64;
486
+ __int64_t* up64_ptr = s_up64_ptr + compute_pipeline_offset_64;
487
385
- // STAGE_SIZE must also be constexpr!
386
- compute_pipeline_offset_64 += STAGE_SIZE;
387
- compute_pipeline_offset_64 %= STAGE_MOD;
488
+ // COMPUTE_STAGE_SIZE/MOD must also be constexpr!
489
+ compute_pipeline_offset_64 += COMPUTE_STAGE_SIZE;
490
+ compute_pipeline_offset_64 %= COMPUTE_STAGE_MOD;
491
389
- // Each thread loads (gate/up) 2X 4X bfloat16 values into registers.
390
- __int64_t gate64 = *s_gate_compute_64;
391
- __nv_bfloat162* s_gate_compute_32 =
392
- reinterpret_cast<__nv_bfloat162*>(&gate64);
492
+ __int64_t gate64 = *gate64_ptr;
493
+ __int64_t up64 = *up64_ptr;
494
394
- __int64_t up64 = *s_up_compute_64;
395
- __nv_bfloat162* s_up_compute_32 = reinterpret_cast<__nv_bfloat162*>(&up64);
495
+ // Compute
496
+ __nv_bfloat162 res[2];
497
+ __nv_bfloat162* s_up_comp = reinterpret_cast<__nv_bfloat162*>(&up64);
498
+ __nv_bfloat162* s_gate_comp = reinterpret_cast<__nv_bfloat162*>(&gate64);
499
500
#pragma unroll
398
- for (int i = 0; i < 2; i++) {
399
- // For silu, we make sure that div is emitted.
400
- float2 gate = silu2(__bfloat1622float2(s_gate_compute_32[i]));
401
- results_bf162[i] = __float22bfloat162_rn(gate);
402
- }
403
-
404
- #pragma unroll
405
- for (int i = 0; i < 2; i++) {
406
- results_bf162[i] = __hmul2(results_bf162[i], s_up_compute_32[i]);
407
- }
501
+ for (int32_t k = 0; k < 2; ++k) {
502
+ __nv_bfloat162 gate = silu2_v2(__bfloat1622float2(s_gate_comp[k]));
503
+ res[k] = __hmul2(gate, s_up_comp[k]);
504
+ }
505
409
- auto _y_max2 =
410
- __hmax2(__habs2(results_bf162[0]), __habs2(results_bf162[1]));
506
+ auto _y_max2 = __hmax2(__habs2(res[0]), __habs2(res[1]));
507
412
- __nv_bfloat16 y_max_bf16 = __hmax(EPS, __hmax(_y_max2.x, _y_max2.y));
508
+ _y_max2.x = __hmax(__hmax(_y_max2.x, _y_max2.y), EPS);
509
414
- // An entire group is assigned to a single warp, so a simple warp reduce
415
- // is used.
416
- __nv_bfloat16 y_s = warp_max(y_max_bf16) / fp8_max;
510
+ __nv_bfloat16 y_s = __hmul(warp_max(_y_max2.x), fp8_inv);
511
418
- if constexpr (USE_UE8M0) {
419
- y_s = hexp2(hceil(hlog2(y_s)));
420
- }
512
+ if constexpr (USE_UE8M0) {
513
+ y_s = hexp2(hceil(hlog2(y_s)));
514
+ }
515
422
- auto inv_y = __float2bfloat16_rn(1.f) / y_s;
516
+ __nv_bfloat16 inv_y = __hdiv(one_bf16, y_s);
517
424
- auto y_s2 = make_bfloat162(inv_y, inv_y);
518
+ auto y_s2 = make_bfloat162(inv_y, inv_y);
519
520
#pragma unroll
427
- for (int32_t i = 0; i < 2; ++i) {
428
- results_bf162[i] =
429
- clip(__hmul2(results_bf162[i], y_s2), __bfloat162bfloat162(fp8_min),
430
- __bfloat162bfloat162(fp8_max));
431
- }
521
+ for (int32_t k = 0; k < 2; ++k) {
522
+ res[k] = clip(__hmul2(res[k], y_s2), __bfloat162bfloat162(fp8_min),
523
+ __bfloat162bfloat162(fp8_max));
524
+ }
525
433
- auto fp8x4 = __nv_fp8x4_e4m3(results_bf162[0], results_bf162[1]);
434
- *reinterpret_cast<__nv_fp8x4_e4m3*>(y_q_ptr) = fp8x4;
435
- y_q_ptr += stride_yq_t;
526
+ *y_q_ptr = __nv_fp8x4_e4m3(res[0], res[1]);
527
+ y_q_ptr += WARP_SIZE * stride_yq_h;
528
437
- if (lane_id == 0) {
438
- *y_s_ptr = y_s;
439
- y_s_ptr += stride_ys_t;
529
+ if (!lane_id) {
530
+ *y_s_ptr = y_s;
531
+ y_s_ptr += stride_ys_g;
532
+ }
533
}
534
}
442
-}
535
#endif
536
+}
537
538
} // namespace vllm
539
@@ -475,14 +568,14 @@ void silu_and_mul_quant(torch::Tensor& out, // [..., d]
568
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
569
}
570
478
-void silu_mul_fp8_quant_deep_gemm_cuda(
479
- const at::Tensor& input, // (E, T, 2*H)
480
- const at::Tensor& counts, // (E)
481
- at::Tensor& y_q, // (E, T, H) [OUT]
482
- at::Tensor& y_s, // (E, T, H//group_size) [OUT]
483
- int64_t group_size, bool use_ue8m0, int64_t num_parallel_tokens) {
571
+void persistent_masked_m_silu_mul_quant(
572
+ const at::Tensor& input, // (E, T, 2*H)
573
+ const at::Tensor& tokens_per_expert, // (E)
574
+ at::Tensor& y_q, // (E, T, H) [OUT]
575
+ at::Tensor& y_s, // (E, T, H//group_size) [OUT]
576
+ bool use_ue8m0) {
577
#ifndef USE_ROCM
485
- // This kernel relies heavily on cp.async and fp8 support.
578
+
579
// This kernel currently only supports H % 128 == 0 and assumes a
580
// fixed GROUP_SIZE of 128.
581
TORCH_CHECK(input.dtype() == torch::kBFloat16);
@@ -491,10 +584,6 @@ void silu_mul_fp8_quant_deep_gemm_cuda(
584
TORCH_CHECK(y_s.dtype() == torch::kFloat32);
585
TORCH_CHECK(input.size(-1) % 256 == 0);
586
494
- // Check that num_parallel_tokens is of power of 2 and between 1 and 64.
495
- TORCH_CHECK(1 <= num_parallel_tokens && num_parallel_tokens <= 64);
496
- TORCH_CHECK(!(num_parallel_tokens & (num_parallel_tokens - 1)));
497
-
587
using Idx_t = int64_t;
588
589
Idx_t E = input.size(0);
@@ -510,81 +599,54 @@ void silu_mul_fp8_quant_deep_gemm_cuda(
599
Idx_t stride_ys_t = y_s.stride(1);
600
Idx_t stride_ys_g = y_s.stride(2);
601
513
- Idx_t stride_counts_e = counts.stride(0);
602
+ Idx_t stride_counts_e = tokens_per_expert.stride(0);
603
604
static constexpr int GROUP_SIZE = 128;
605
517
- #define KERNEL_FN \
518
- if (use_ue8m0) { \
519
- vllm::silu_mul_fp8_quant_deep_gemm_kernel<fp8_t, NUM_WARPS, Idx_t, \
520
- NUM_PARALLEL_TOKENS, true> \
521
- <<<grid, block, 0, stream>>>( \
522
- reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
523
- (fp8_t*)y_q.data_ptr(), y_s.data_ptr<float>(), \
524
- reinterpret_cast<int32_t*>(counts.data_ptr<int>()), H, G, \
525
- stride_i_e, stride_i_t, stride_i_h, stride_yq_e, stride_yq_t, \
526
- stride_yq_h, stride_ys_e, stride_ys_t, stride_ys_g, \
527
- stride_counts_e); \
528
- } else { \
529
- vllm::silu_mul_fp8_quant_deep_gemm_kernel<fp8_t, NUM_WARPS, Idx_t, \
530
- NUM_PARALLEL_TOKENS, false> \
531
- <<<grid, block, 0, stream>>>( \
532
- reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
533
- (fp8_t*)y_q.data_ptr(), y_s.data_ptr<float>(), \
534
- reinterpret_cast<int32_t*>(counts.data_ptr<int>()), H, G, \
535
- stride_i_e, stride_i_t, stride_i_h, stride_yq_e, stride_yq_t, \
536
- stride_yq_h, stride_ys_e, stride_ys_t, stride_ys_g, \
537
- stride_counts_e); \
538
- }
606
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
607
540
- #define KERNEL_CALL_H \
541
- if (H % (4 * GROUP_SIZE) == 0) { \
542
- static constexpr int NUM_WARPS = 4; \
543
- populate_launch_params(NUM_WARPS, NUM_PARALLEL_TOKENS); \
544
- KERNEL_FN \
545
- } else { \
546
- static constexpr int NUM_WARPS = 1; \
547
- populate_launch_params(NUM_WARPS, NUM_PARALLEL_TOKENS); \
548
- KERNEL_FN \
608
+ #define KERNEL(BLOCK_COUNT, USE_UE8M0, THREAD_COUNT, STAGES) \
609
+ static constexpr int NUM_WARPS = THREAD_COUNT / WARP_SIZE; \
610
+ int sms = SILU_V2_BLOCK_COUNT; \
611
+ static constexpr int max_shared_mem_bytes = \
612
+ GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2; \
613
+ dim3 grid(sms), block(THREAD_COUNT); \
614
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
615
+ VLLM_DISPATCH_FP8_TYPES( \
616
+ y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] { \
617
+ vllm::silu_mul_fp8_quant_deep_gemm_kernel< \
618
+ BLOCK_COUNT, max_shared_mem_bytes, fp8_t, THREAD_COUNT, Idx_t, \
619
+ USE_UE8M0, GROUP_SIZE, STAGES> \
620
+ <<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>( \
621
+ reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
622
+ (fp8_t*)y_q.data_ptr(), y_s.data_ptr<float>(), \
623
+ reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
624
+ T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
625
+ stride_yq_t, stride_yq_h, stride_ys_e, stride_ys_t, \
626
+ stride_ys_g, stride_counts_e); \
627
+ });
628
+
629
+ static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
630
+
631
+ if (!use_ue8m0) {
632
+ if (H >= 4096) {
633
+ static constexpr int NUM_STAGES = 4;
634
+ static constexpr int THREAD_COUNT = 256;
635
+ KERNEL(SILU_V2_BLOCK_COUNT, false, THREAD_COUNT, NUM_STAGES);
636
+ } else {
637
+ static constexpr int THREAD_COUNT = 32;
638
+ KERNEL(SILU_V2_BLOCK_COUNT, false, THREAD_COUNT, 2);
639
}
550
-
551
- #define KERNEL_CALL_TOP_LEVEL \
552
- if (num_parallel_tokens == 1) { \
553
- static constexpr int NUM_PARALLEL_TOKENS = 1; \
554
- KERNEL_CALL_H \
555
- } else if (num_parallel_tokens == 2) { \
556
- static constexpr int NUM_PARALLEL_TOKENS = 2; \
557
- KERNEL_CALL_H \
558
- } else if (num_parallel_tokens == 4) { \
559
- static constexpr int NUM_PARALLEL_TOKENS = 4; \
560
- KERNEL_CALL_H \
561
- } else if (num_parallel_tokens == 8) { \
562
- static constexpr int NUM_PARALLEL_TOKENS = 8; \
563
- KERNEL_CALL_H \
564
- } else if (num_parallel_tokens == 16) { \
565
- static constexpr int NUM_PARALLEL_TOKENS = 16; \
566
- KERNEL_CALL_H \
567
- } else if (num_parallel_tokens == 32) { \
568
- static constexpr int NUM_PARALLEL_TOKENS = 32; \
569
- KERNEL_CALL_H \
570
- } else if (num_parallel_tokens == 64) { \
571
- static constexpr int NUM_PARALLEL_TOKENS = 64; \
572
- KERNEL_CALL_H \
640
+ } else {
641
+ if (H >= 4096) {
642
+ static constexpr int NUM_STAGES = 4;
643
+ static constexpr int THREAD_COUNT = 256;
644
+ KERNEL(SILU_V2_BLOCK_COUNT, true, THREAD_COUNT, NUM_STAGES);
645
+ } else {
646
+ static constexpr int THREAD_COUNT = 32;
647
+ KERNEL(SILU_V2_BLOCK_COUNT, true, THREAD_COUNT, 2);
648
}
574
-
575
- Idx_t G;
576
- dim3 block, grid;
577
- auto populate_launch_params = [&](int num_warps, int _num_parallel_tokens) {
578
- G = H / Idx_t(group_size * num_warps);
579
- grid = dim3(E * G, _num_parallel_tokens);
580
- block = dim3(num_warps * WARP_SIZE);
581
- };
582
-
583
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
584
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
585
- VLLM_DISPATCH_FP8_TYPES(y_q.scalar_type(),
586
- "silu_mul_fp8_quant_deep_gemm_kernel",
587
- [&] { KERNEL_CALL_TOP_LEVEL });
649
+ }
650
651
#endif
652
}
csrc/torch_bindings.cpp
+5
-5
@@ -33,11 +33,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
33
#endif
34
35
ops.def(
36
- "silu_mul_fp8_quant_deep_gemm_cuda(Tensor input, Tensor counts, Tensor! "
37
- "y_q, Tensor! y_s, int group_size, "
38
- "bool use_ue8m0, int num_parallel_tokens) -> ()");
39
- ops.impl("silu_mul_fp8_quant_deep_gemm_cuda", torch::kCUDA,
40
- &silu_mul_fp8_quant_deep_gemm_cuda);
36
+ "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
37
+ "y_q, Tensor! y_s,"
38
+ "bool use_ue8m0) -> ()");
39
+ ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA,
40
+ &persistent_masked_m_silu_mul_quant);
41
42
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
43
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
+6
-5
@@ -5,7 +5,7 @@ import pytest
5
import torch
6
7
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
8
- silu_mul_fp8_quant_deep_gemm_cuda,
8
+ persistent_masked_m_silu_mul_quant,
9
)
10
from vllm.platforms import current_platform
11
from vllm.utils import cdiv
@@ -50,15 +50,15 @@ def test_silu_mul_fp8_quant_deep_gemm(E, T, H, fp8_type):
50
# Input tensor of shape (E, T, 2*H)
51
y = torch.randn((E, T, 2 * H), dtype=torch.bfloat16, device="cuda")
52
tokens_per_expert = torch.randint(
53
- low=T // 2,
53
+ low=0,
54
high=T,
55
size=(E,),
56
dtype=torch.int32,
57
device="cuda",
58
)
59
60
- # Run the Triton kernel
61
- y_q, y_s = silu_mul_fp8_quant_deep_gemm_cuda(
60
+ # Run the SiLU V2 kernel
61
+ y_q, y_s = persistent_masked_m_silu_mul_quant(
62
y, tokens_per_expert, group_size=group_size
63
)
64
@@ -115,10 +115,11 @@ def test_silu_mul_fp8_quant_deep_gemm(E, T, H, fp8_type):
115
y_se = y_s[e].float()
116
y_qe = y_q[e].float()
117
118
- torch.testing.assert_close(y_se[:nt], ref_s[:nt], atol=1e-4, rtol=1e-2)
118
torch.testing.assert_close(
119
y_qe[:nt].to(torch.float32),
120
ref_q[:nt].to(torch.float32),
121
atol=2,
122
rtol=2e-1,
123
)
124
+
125
+ torch.testing.assert_close(y_se[:nt], ref_s[:nt], atol=1e-4, rtol=1e-2)
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+36
-30
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from math import log2
3
from typing import Optional
4
5
import torch
@@ -94,7 +93,7 @@ def _silu_mul_fp8_quant_deep_gemm(
93
tl.store(y_s_ptr + base_ys_offset + t * stride_ys_t, y_s)
94
95
97
-def silu_mul_fp8_quant_deep_gemm_cuda(
96
+def persistent_masked_m_silu_mul_quant(
97
y: torch.Tensor, # (E, T, 2*H)
98
tokens_per_expert: torch.Tensor, # (E,) number of valid tokens per expert
99
num_parallel_tokens=16,
@@ -103,9 +102,41 @@ def silu_mul_fp8_quant_deep_gemm_cuda(
102
"""Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales
103
y has shape (E, T, 2*H). The first half of the last dimension is
104
silu-activated, multiplied by the second half, then quantized into FP8.
105
+ We launch a fixed grid of threads to accommodate CUDA graphs. Let `P2`
106
+ be a parallelization factor for persistent_masked_m_silu_mul_quant over the
107
+ hidden dimension.
108
+
109
+ Let `expert_offsets = [0] + [num_tokens.cumsum()]` and
110
+ `total_tokens = expert_offsets[-1]`.
111
+ persistent_masked_m_silu_mul_quant launches `total_tokens x P2` number of
112
+ thread blocks. Each thread block contains `NUM_WARPS` warps.
113
+
114
+ Every thread block needs to find it's corresponding expert by warp-parallel scanning
115
+ over the `expert_offsets` array.
116
+
117
+ The i-th warp in the first thread block processes
118
+ `[i * warp_chunk_size, (i + 1) * warp_chunk_size]` groups
119
+ sequentially, where `warp_chunk_size = ((H / GROUP_SIZE) / P2) / NUM_WARPS`,
120
+ pipelining loads and computes.
121
+
122
+ The shared memory layout for 4 warps with a 2-stage pipeline for SiLU V2
123
+ can is visualized like so:
124
+
125
+ stage0 stage1
126
+ ┌─────┬───┬─────┬───┬─────┬───┬─────┬───┬─────┬───┬─────┬───┬─────┬───┬─────┬───┐
127
+ │gate0│up0│gate1│up1│gate2│up2│gate3│up3│gate0│up0│gate1│up1│gate2│up2│gate3│up3│
128
+ └─────┴───┴─────┴───┴─────┴───┴─────┴───┴─────┴───┴─────┴───┴─────┴───┴─────┴───┘
129
+
130
+ with the main difference between V1 and V2 being the global load
131
+ stride between warps, and between half-warps. Regarding the latter stride,
132
+ we assign the first half warp of every warp for `gate` loads and the second
133
+ half-warp to `up` loads.
134
+
135
Returns `(y_q, y_s)` where
136
* `y_q`: FP8 tensor, shape (E, T, H), same layout as y[..., :H]
137
* `y_s`: FP32 tensor, shape (E, T, H // group_size), strides (T*G, 1, T)
138
+ Let NUM_WARPS be the number of warps in a single thread block and
139
+ `GROUP_SIZE = 128` be the size of the quantization group.
140
"""
141
assert y.ndim == 3, "y must be (E, T, 2*H)"
142
E, T, H2 = y.shape
@@ -133,30 +164,15 @@ def silu_mul_fp8_quant_deep_gemm_cuda(
164
165
use_ue8m0 = is_deep_gemm_e8m0_used()
166
136
- if E <= 16:
137
- max_empirical_parallelism = 64
138
- elif E <= 32:
139
- max_empirical_parallelism = 16
140
- else:
141
- max_empirical_parallelism = 4
142
-
143
- # We never want to launch more than Tx number of threads
144
- # This computes the clip.
145
- num_parallel_tokens = max(
146
- 1, min(max_empirical_parallelism, 2 ** int(log2(min(num_parallel_tokens, T))))
147
- )
167
cuda_arch = current_platform.get_device_capability(
168
device_id=y.device.index
169
).to_int()
170
171
if cuda_arch >= 80:
153
- torch.ops._C.silu_mul_fp8_quant_deep_gemm_cuda(
154
- y, tokens_per_expert, y_q, y_s, group_size, use_ue8m0, num_parallel_tokens
172
+ torch.ops._C.persistent_masked_m_silu_mul_quant(
173
+ y, tokens_per_expert, y_q, y_s, use_ue8m0
174
)
175
else:
157
- # Default to triton if not on cuda or if arch is too old
158
- y_q = torch.empty((E, T, H), dtype=fp8_dtype, device=y.device)
159
-
176
stride_cnt_e = tokens_per_expert.stride()[0]
177
178
# Static grid over experts and H-groups.
@@ -166,16 +182,6 @@ def silu_mul_fp8_quant_deep_gemm_cuda(
182
stride_i_e, stride_i_t, stride_i_h = y.stride()
183
stride_yq_e, stride_yq_t, stride_yq_h = y_q.stride()
184
169
- # desired scale strides (elements): (T*G, 1, T)
170
- stride_ys_e = T * G
171
- stride_ys_t = 1
172
- stride_ys_g = T
173
- y_s = torch.empty_strided(
174
- (E, T, G),
175
- (stride_ys_e, stride_ys_t, stride_ys_g),
176
- dtype=torch.float32,
177
- device=y.device,
178
- )
185
f_info = torch.finfo(fp8_dtype)
186
fp8_max = f_info.max
187
fp8_min = f_info.min
@@ -313,7 +319,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
319
expected_m,
320
)
321
316
- a2q, a2q_scale = silu_mul_fp8_quant_deep_gemm_cuda(
322
+ a2q, a2q_scale = persistent_masked_m_silu_mul_quant(
323
workspace1, expert_num_tokens
324
)
325