A high-throughput and memory-efficient inference and serving engine for LLMs
Revert "[Kernel] Speed up silu_and_mul_per_block_quant with warp-shuf… (#45066)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Micah Williamson committed
Jun 9, 2026 at 16:12 UTC
e1ed89dbee7190f356c6f19da02e60bd457a365e
1 file changed
+58
-89
csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu
+58
-89
@@ -8,23 +8,11 @@
8
9
namespace vllm {
10
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
11
+// Logic: one thread block per (token, group) pair
12
13
template <typename scalar_t, typename scalar_out_t, bool is_scale_transposed,
14
int32_t group_size>
26
-__global__ void
27
-__launch_bounds__(kWarpsPerBlock * 32) silu_and_mul_per_block_quant_kernel(
15
+__global__ void silu_and_mul_per_block_quant_kernel(
16
scalar_out_t* __restrict__ out, // Output: [num_tokens, hidden_size] in
17
// FP8/INT8
18
float* __restrict__ scales, // Output: [num_tokens, hidden_size /
@@ -36,98 +24,82 @@ __launch_bounds__(kWarpsPerBlock * 32) silu_and_mul_per_block_quant_kernel(
24
) {
25
static_assert((group_size & (group_size - 1)) == 0,
26
"group_size must be a power of 2 for correct reduction");
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;
27
+
28
+ // Grid: (num_tokens, num_groups)
29
int const token_idx = blockIdx.x;
30
+ int const group_idx = blockIdx.y;
31
+ int const tid = threadIdx.x; // tid in [0, group_size)
32
int const num_tokens = gridDim.x;
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.
33
+
34
+ // Input layout: [gate || up] concatenated along last dimension
35
int const input_stride = hidden_size * 2;
36
int const group_start = group_idx * group_size;
62
- int const lane_base = group_start + lane_id * EPT;
37
+
38
+ // Pointers to this token's data
39
scalar_t const* token_input_gate =
64
- input + token_idx * input_stride + lane_base;
40
+ input + token_idx * input_stride + group_start;
41
scalar_t const* token_input_up = token_input_gate + hidden_size;
66
- scalar_out_t* token_output = out + token_idx * hidden_size + lane_base;
42
+ scalar_out_t* token_output = out + token_idx * hidden_size + group_start;
43
44
// Scale pointer for this group
45
+ int const num_groups = gridDim.y;
46
float* group_scale_ptr = is_scale_transposed
47
? scales + group_idx * num_tokens + token_idx
48
: scales + token_idx * num_groups + group_idx;
49
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
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
- }
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
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.
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)
67
#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));
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();
73
}
74
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;
75
+ // Step 3: Compute scale (thread 0), broadcast via shared memory
76
+ if (tid == 0) {
77
+ float group_max = shared_max[0];
78
108
- // Apply scale upper bound if provided
109
- if (scale_ub != nullptr) {
110
- group_scale = fminf(group_scale, *scale_ub);
111
- }
79
+ float const quant_range = quant_type_max_v<scalar_out_t>;
80
+ float group_scale = group_max / quant_range;
81
113
- // Use minimum safe scaling factor
114
- group_scale = fmaxf(group_scale, min_scaling_factor<scalar_out_t>::val());
82
+ // Apply scale upper bound if provided
83
+ if (scale_ub != nullptr) {
84
+ group_scale = fminf(group_scale, *scale_ub);
85
+ }
86
116
- if (lane_id == 0) {
87
+ // Use minimum safe scaling factor
88
+ group_scale = fmaxf(group_scale, min_scaling_factor<scalar_out_t>::val());
89
+
90
+ // Store scale to global memory
91
*group_scale_ptr = group_scale;
118
- }
92
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);
93
+ // Reuse shared_max[0] to broadcast scale
94
+ shared_max[0] = group_scale;
95
}
130
- *reinterpret_cast<OutVec*>(token_output) = out_v;
96
+ __syncthreads();
97
+
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);
103
}
104
105
} // namespace vllm
@@ -170,11 +142,8 @@ void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
142
input.get_device_index());
143
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
144
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);
145
+ dim3 grid(num_tokens, num_groups);
146
+ dim3 block(group_size);
147
148
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
149
input.scalar_type(), "silu_and_mul_per_block_quant", [&] {