A high-throughput and memory-efficient inference and serving engine for LLMs
[Kernel] Speed up silu_and_mul_per_block_quant with warp-shuffle reduction + vectorized I/O (#44173)
Signed-off-by: SII-yangdian <yangdian@sii.edu.cn> Co-authored-by: SII-yangdian <yangdian@sii.edu.cn>
SII-yangdian committed
Jun 10, 2026 at 01:41 UTC
a4b14b98c6f6a10bee1f48154baf4098b7e117ca
1 file changed
+89
-58
csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu
+89
-58
@@ -8,11 +8,23 @@
8
9
namespace vllm {
10
11
-// Logic: one thread block per (token, group) pair
11
+// Logic: one WARP per (token, group). Each of the 32 lanes owns EPT =
12
+// group_size / 32 contiguous elements, sweeps the group in a single coalesced
13
+// wide vector load (gate and up), and the per-group abs-max is a pure
14
+// warp-shuffle reduction -- no shared memory and no __syncthreads.
15
+// kWarpsPerBlock groups are packed into one 128-thread block.
16
+//
17
+// This replaces the previous one-block-per-group design, whose log2(group_size)
18
+// shared-memory tree reduction (one __syncthreads per level) and one-element
19
+// scalar loads left the kernel memory-latency-bound and ~2x off peak. The op
20
+// sequence (fp32 SiLU, fmaxf abs-max) is unchanged, and fmaxf is
21
+// order-invariant, so the per-group scale stays bit-identical.
22
+constexpr int kWarpsPerBlock = 4; // 4 groups per 128-thread block
23
24
template <typename scalar_t, typename scalar_out_t, bool is_scale_transposed,
25
int32_t group_size>
15
-__global__ void silu_and_mul_per_block_quant_kernel(
26
+__global__ void
27
+__launch_bounds__(kWarpsPerBlock * 32) silu_and_mul_per_block_quant_kernel(
28
scalar_out_t* __restrict__ out, // Output: [num_tokens, hidden_size] in
29
// FP8/INT8
30
float* __restrict__ scales, // Output: [num_tokens, hidden_size /
@@ -24,82 +36,98 @@ __global__ void silu_and_mul_per_block_quant_kernel(
36
) {
37
static_assert((group_size & (group_size - 1)) == 0,
38
"group_size must be a power of 2 for correct reduction");
27
-
28
- // Grid: (num_tokens, num_groups)
39
+ static_assert(group_size % 32 == 0,
40
+ "group_size must be a multiple of the warp size");
41
+ // Elements per thread: 2 for group_size=64, 4 for group_size=128. Each maps
42
+ // to a single vector load/store per lane (4B/8B in, 2B/4B out).
43
+ constexpr int EPT = group_size / 32;
44
+
45
+ int const tid = threadIdx.x;
46
+ int const warp_id = tid >> 5;
47
+ int const lane_id = tid & 31;
48
int const token_idx = blockIdx.x;
30
- int const group_idx = blockIdx.y;
31
- int const tid = threadIdx.x; // tid in [0, group_size)
49
int const num_tokens = gridDim.x;
33
-
34
- // Input layout: [gate || up] concatenated along last dimension
50
+ // num_groups is no longer gridDim.y (we pack kWarpsPerBlock groups per
51
+ // block), so recover it from the compile-time group_size.
52
+ int const num_groups = hidden_size / group_size;
53
+ int const group_idx = blockIdx.y * kWarpsPerBlock + warp_id;
54
+ if (group_idx >= num_groups) return; // whole warp exits together (no sync)
55
+
56
+ // Input layout: [gate || up] concatenated along the last dimension. Each lane
57
+ // owns the EPT contiguous elements at group_start + lane_id * EPT, so the
58
+ // warp reads the whole group as one fully-coalesced wide load for gate and
59
+ // for up.
60
int const input_stride = hidden_size * 2;
61
int const group_start = group_idx * group_size;
37
-
38
- // Pointers to this token's data
62
+ int const lane_base = group_start + lane_id * EPT;
63
scalar_t const* token_input_gate =
40
- input + token_idx * input_stride + group_start;
64
+ input + token_idx * input_stride + lane_base;
65
scalar_t const* token_input_up = token_input_gate + hidden_size;
42
- scalar_out_t* token_output = out + token_idx * hidden_size + group_start;
66
+ scalar_out_t* token_output = out + token_idx * hidden_size + lane_base;
67
68
// Scale pointer for this group
45
- int const num_groups = gridDim.y;
69
float* group_scale_ptr = is_scale_transposed
70
? scales + group_idx * num_tokens + token_idx
71
: scales + token_idx * num_groups + group_idx;
72
50
- // Shared memory for reduction (compile-time sized)
51
- __shared__ float shared_max[group_size];
52
-
53
- // Step 1: Each thread loads one element, computes SiLU, stores in register
54
- float gate = static_cast<float>(token_input_gate[tid]);
55
- float up = static_cast<float>(token_input_up[tid]);
56
-
57
- // Compute SiLU(gate) * up
58
- float sigmoid_gate = 1.0f / (1.0f + expf(-gate));
59
- float silu_gate = gate * sigmoid_gate;
60
- float result = silu_gate * up; // Keep in register
61
-
62
- // Step 2: Reduce to find group max
63
- shared_max[tid] = fabsf(result);
64
- __syncthreads();
65
-
66
-// Power-of-2 reduction (group_size guaranteed to be power of 2)
73
+ // Step 1: one wide vector load per lane for gate and up, then SiLU(gate) * up
74
+ // in fp32. (group_start and hidden_size are both multiples of EPT because
75
+ // group_size = 32 * EPT divides hidden_size, so these loads are aligned.)
76
+ struct alignas(sizeof(scalar_t) * EPT) InVec {
77
+ scalar_t v[EPT];
78
+ };
79
+ InVec const gate_v = *reinterpret_cast<InVec const*>(token_input_gate);
80
+ InVec const up_v = *reinterpret_cast<InVec const*>(token_input_up);
81
+
82
+ float result[EPT]; // SiLU(gate) * up, kept in registers
83
+ float thread_max = 0.0f;
84
#pragma unroll
68
- for (int stride = group_size / 2; stride > 0; stride >>= 1) {
69
- if (tid < stride) {
70
- shared_max[tid] = fmaxf(shared_max[tid], shared_max[tid + stride]);
71
- }
72
- __syncthreads();
85
+ for (int k = 0; k < EPT; ++k) {
86
+ float gate = static_cast<float>(gate_v.v[k]);
87
+ float up = static_cast<float>(up_v.v[k]);
88
+ float sigmoid_gate = 1.0f / (1.0f + expf(-gate));
89
+ float silu_gate = gate * sigmoid_gate;
90
+ result[k] = silu_gate * up;
91
+ thread_max = fmaxf(thread_max, fabsf(result[k]));
92
}
93
75
- // Step 3: Compute scale (thread 0), broadcast via shared memory
76
- if (tid == 0) {
77
- float group_max = shared_max[0];
94
+ // Step 2: per-group abs-max via warp-shuffle. fmaxf is order-invariant, so
95
+ // the group max (and therefore the scale) is bit-identical to the tree
96
+ // reduction.
97
+#pragma unroll
98
+ for (int offset = 16; offset > 0; offset >>= 1) {
99
+ thread_max =
100
+ fmaxf(thread_max, __shfl_xor_sync(0xffffffffu, thread_max, offset));
101
+ }
102
79
- float const quant_range = quant_type_max_v<scalar_out_t>;
80
- float group_scale = group_max / quant_range;
103
+ // Step 3: compute the group scale in registers; lane 0 writes it to global.
104
+ float const group_max = thread_max;
105
+ float const quant_range = quant_type_max_v<scalar_out_t>;
106
+ float group_scale = group_max / quant_range;
107
82
- // Apply scale upper bound if provided
83
- if (scale_ub != nullptr) {
84
- group_scale = fminf(group_scale, *scale_ub);
85
- }
108
+ // Apply scale upper bound if provided
109
+ if (scale_ub != nullptr) {
110
+ group_scale = fminf(group_scale, *scale_ub);
111
+ }
112
87
- // Use minimum safe scaling factor
88
- group_scale = fmaxf(group_scale, min_scaling_factor<scalar_out_t>::val());
113
+ // Use minimum safe scaling factor
114
+ group_scale = fmaxf(group_scale, min_scaling_factor<scalar_out_t>::val());
115
90
- // Store scale to global memory
116
+ if (lane_id == 0) {
117
*group_scale_ptr = group_scale;
92
-
93
- // Reuse shared_max[0] to broadcast scale
94
- shared_max[0] = group_scale;
118
}
96
- __syncthreads();
119
98
- float group_scale = shared_max[0];
99
-
100
- // Step 4: Quantize and write output
101
- token_output[tid] =
102
- vllm::ScaledQuant<scalar_out_t, false>::quant_fn(result, group_scale);
120
+ // Step 4: quantize the EPT owned elements and write them with one wide store.
121
+ struct alignas(sizeof(scalar_out_t) * EPT) OutVec {
122
+ scalar_out_t q[EPT];
123
+ };
124
+ OutVec out_v;
125
+#pragma unroll
126
+ for (int k = 0; k < EPT; ++k) {
127
+ out_v.q[k] = vllm::ScaledQuant<scalar_out_t, false>::quant_fn(result[k],
128
+ group_scale);
129
+ }
130
+ *reinterpret_cast<OutVec*>(token_output) = out_v;
131
}
132
133
} // namespace vllm
@@ -142,8 +170,11 @@ void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
170
input.get_device_index());
171
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
172
145
- dim3 grid(num_tokens, num_groups);
146
- dim3 block(group_size);
173
+ // One warp per group; vllm::kWarpsPerBlock groups packed per 128-thread
174
+ // block.
175
+ dim3 grid(num_tokens,
176
+ (num_groups + vllm::kWarpsPerBlock - 1) / vllm::kWarpsPerBlock);
177
+ dim3 block(vllm::kWarpsPerBlock * 32);
178
179
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
180
input.scalar_type(), "silu_and_mul_per_block_quant", [&] {