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, 00:30 UTC
DSv4 fused Q-norm kernel grid refactor (#42353)
gnovack committed
May 22, 2026 at 15:21 UTC
f743254143f25e7f3519b2cd5f3f766692b6e397
2 files changed
+330
-216
csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+303
-192
@@ -87,6 +87,12 @@ constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad)
87
constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576
88
constexpr float kFp8Max = 448.0f;
89
90
+#ifndef USE_ROCM
91
+// When num_tokens is less than this threshold,
92
+// run the reduced grid variant on cuda
93
+constexpr float NUM_TOKEN_CUTOFF = 1024;
94
+#endif
95
+
96
// Per-warp layout: 32 lanes × 16 elems/lane = 512 elems = HEAD_DIM.
97
constexpr int kNumLanes = 32;
98
constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 16
@@ -112,6 +118,197 @@ __device__ __forceinline__ float warpSum(float val) {
118
return val;
119
}
120
121
+// ────────────────────────────────────────────────────────────────────────────
122
+// Per-slot inner pipeline
123
+// ────────────────────────────────────────────────────────────────────────────
124
+// Shared by both kernel variants: 1 CTA per (token, head) pair vs. 1 CTA per
125
+// token
126
+template <typename scalar_t_in>
127
+__device__ __forceinline__ void processDeepseekV4Slot(
128
+ uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx,
129
+ int const dim_base, int const laneId, int const num_heads_q,
130
+ float const eps, scalar_t_in* __restrict__ q_inout,
131
+ uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping,
132
+ int64_t const* __restrict__ position_ids,
133
+ float const* __restrict__ cos_sin_cache, int const cache_block_size,
134
+ int const kv_block_stride) {
135
+ using Converter = vllm::_typeConvert<scalar_t_in>;
136
+ bool const isKV = (slotIdx == num_heads_q);
137
+
138
+ // ── Decode the bf16 → 16 fp32 registers ─────────────────────────────
139
+ float elements[kElemsPerLane];
140
+ {
141
+ typename Converter::packed_hip_type const* p0 =
142
+ reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
143
+ typename Converter::packed_hip_type const* p1 =
144
+ reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
145
+#pragma unroll
146
+ for (int i = 0; i < 4; i++) {
147
+ float2 f2 = Converter::convert(p0[i]);
148
+ elements[2 * i] = f2.x;
149
+ elements[2 * i + 1] = f2.y;
150
+ }
151
+#pragma unroll
152
+ for (int i = 0; i < 4; i++) {
153
+ float2 f2 = Converter::convert(p1[i]);
154
+ elements[8 + 2 * i] = f2.x;
155
+ elements[8 + 2 * i + 1] = f2.y;
156
+ }
157
+ }
158
+
159
+ // ── Q branch: RMSNorm (no weight) ───────────────────────────────────
160
+ if (!isKV) {
161
+ float sumOfSquares = 0.0f;
162
+#pragma unroll
163
+ for (int i = 0; i < kElemsPerLane; i++) {
164
+ sumOfSquares += elements[i] * elements[i];
165
+ }
166
+ sumOfSquares = warpSum<float>(sumOfSquares);
167
+ float const rms_rcp =
168
+ rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
169
+#pragma unroll
170
+ for (int i = 0; i < kElemsPerLane; i++) {
171
+ elements[i] = elements[i] * rms_rcp;
172
+ }
173
+ }
174
+
175
+ // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
176
+ // All math in fp32. cos_sin_cache is loaded as fp32 (its native storage).
177
+ bool const is_rope_lane = dim_base >= kNopeDim;
178
+ if (is_rope_lane) {
179
+ int64_t const pos = position_ids[tokenIdx];
180
+ constexpr int kHalfRope = kRopeDim / 2;
181
+ float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
182
+ float const* sin_ptr = cos_ptr + kHalfRope;
183
+
184
+ int const rope_local_base = dim_base - kNopeDim;
185
+ int const half_base = rope_local_base >> 1;
186
+
187
+ // Load phase: 4 vectorized LDGs issue back-to-back.
188
+ float4 const c0 = *reinterpret_cast<float4 const*>(cos_ptr + half_base);
189
+ float4 const c1 = *reinterpret_cast<float4 const*>(cos_ptr + half_base + 4);
190
+ float4 const s0 = *reinterpret_cast<float4 const*>(sin_ptr + half_base);
191
+ float4 const s1 = *reinterpret_cast<float4 const*>(sin_ptr + half_base + 4);
192
+ float const cos_arr[8] = {c0.x, c0.y, c0.z, c0.w, c1.x, c1.y, c1.z, c1.w};
193
+ float const sin_arr[8] = {s0.x, s0.y, s0.z, s0.w, s1.x, s1.y, s1.z, s1.w};
194
+
195
+#pragma unroll
196
+ for (int p = 0; p < kElemsPerLane / 2; p++) {
197
+ float const x_even = elements[2 * p];
198
+ float const x_odd = elements[2 * p + 1];
199
+ elements[2 * p] = x_even * cos_arr[p] - x_odd * sin_arr[p];
200
+ elements[2 * p + 1] = x_even * sin_arr[p] + x_odd * cos_arr[p];
201
+ }
202
+ }
203
+
204
+ // ═══════════════════════════════════════════════════════════════════
205
+ // Q / KV branch dispatch. Restructured as if/else (no early `return`)
206
+ // so every code path lands at the same exit point — callers own PDL
207
+ // triggering and per-iteration buffer rotation.
208
+ // ═══════════════════════════════════════════════════════════════════
209
+ if (!isKV) {
210
+ // ── Q: cast back to bf16 and store. ────────────────────────────
211
+ uint4 out0, out1;
212
+ typename Converter::packed_hip_type* po0 =
213
+ reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
214
+ typename Converter::packed_hip_type* po1 =
215
+ reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
216
+#pragma unroll
217
+ for (int i = 0; i < 4; i++) {
218
+ po0[i] =
219
+ Converter::convert(make_float2(elements[2 * i], elements[2 * i + 1]));
220
+ }
221
+#pragma unroll
222
+ for (int i = 0; i < 4; i++) {
223
+ po1[i] = Converter::convert(
224
+ make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
225
+ }
226
+ scalar_t_in* dst =
227
+ q_inout +
228
+ (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
229
+ dim_base;
230
+ *reinterpret_cast<uint4*>(dst) = out0;
231
+ *reinterpret_cast<uint4*>(dst + 8) = out1;
232
+ } else {
233
+ // ── KV: FP8 quant on NoPE + bf16 store on RoPE + cache insert.
234
+ int64_t const slot_id = slot_mapping[tokenIdx];
235
+ if (slot_id >= 0) {
236
+ int64_t const block_idx = slot_id / cache_block_size;
237
+ int64_t const pos_in_block = slot_id % cache_block_size;
238
+ uint8_t* block_base =
239
+ k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
240
+ uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
241
+ uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
242
+ uint8_t* token_scale_ptr =
243
+ block_base +
244
+ static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
245
+ pos_in_block * kScaleBytesPerToken;
246
+
247
+#pragma unroll
248
+ for (int i = 0; i < kElemsPerLane; i++) {
249
+ elements[i] = Converter::convert(Converter::convert(elements[i]));
250
+ }
251
+
252
+ float local_absmax = 0.0f;
253
+#pragma unroll
254
+ for (int i = 0; i < kElemsPerLane; i++) {
255
+ local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
256
+ }
257
+ float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
258
+ float const exponent = ceilf(log2f(absmax / kFp8Max));
259
+ float const inv_scale = exp2f(-exponent);
260
+
261
+ if (!is_rope_lane) {
262
+ uint8_t out_bytes[kElemsPerLane];
263
+#pragma unroll
264
+ for (int i = 0; i < kElemsPerLane; i++) {
265
+ float scaled = elements[i] * inv_scale;
266
+ scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
267
+#ifndef USE_ROCM
268
+ __nv_fp8_storage_t s =
269
+ __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
270
+ out_bytes[i] = static_cast<uint8_t>(s);
271
+#else
272
+ out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
273
+#endif
274
+ }
275
+ *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
276
+ *reinterpret_cast<uint4 const*>(out_bytes);
277
+
278
+ if ((laneId & 3) == 0) {
279
+ int const q_block_idx = laneId >> 2;
280
+ float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
281
+ token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
282
+ }
283
+ if (laneId == 0) {
284
+ token_scale_ptr[kNumQuantBlocks] = 0;
285
+ }
286
+ } else {
287
+ uint4 out0, out1;
288
+ typename Converter::packed_hip_type* po0 =
289
+ reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
290
+ typename Converter::packed_hip_type* po1 =
291
+ reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
292
+#pragma unroll
293
+ for (int i = 0; i < 4; i++) {
294
+ po0[i] = Converter::convert(
295
+ make_float2(elements[2 * i], elements[2 * i + 1]));
296
+ }
297
+#pragma unroll
298
+ for (int i = 0; i < 4; i++) {
299
+ po1[i] = Converter::convert(
300
+ make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
301
+ }
302
+ int const rope_local_base = dim_base - kNopeDim;
303
+ scalar_t_in* bf16_dst =
304
+ reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
305
+ *reinterpret_cast<uint4*>(bf16_dst) = out0;
306
+ *reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
307
+ }
308
+ }
309
+ }
310
+}
311
+
312
// ────────────────────────────────────────────────────────────────────────────
313
// Kernel
314
// ────────────────────────────────────────────────────────────────────────────
@@ -149,8 +346,6 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
346
return;
347
} else {
348
#endif
152
- using Converter = vllm::_typeConvert<scalar_t_in>;
153
-
349
int const warpsPerBlock = blockDim.x / 32;
350
int const warpId = threadIdx.x / 32;
351
int const laneId = threadIdx.x % 32;
@@ -176,10 +371,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
371
// Dim range this lane owns within the 512-wide head.
372
int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16
373
179
- // ── Load 16 bf16 → 16 fp32 registers (one 16-byte + one 16-byte LDG) ────
180
- float elements[kElemsPerLane];
181
- float sumOfSquares = 0.0f;
182
-
374
+ // Two 16-byte loads per thread (8 bf16 each). Use uint4 as the vector
375
+ // type; the shared per-slot helper bitcasts to scalar_t_in packed pairs.
376
scalar_t_in const* src_ptr;
377
if (isKV) {
378
src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
@@ -189,196 +382,103 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
382
dim_base;
383
src_ptr = q_inout + q_row_offset;
384
}
385
+ uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
386
+ uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
387
193
- // Two 16-byte loads per thread (8 bf16 each). Use uint4 as the vector
194
- // type and bitcast to scalar_t_in packed pairs for conversion.
195
- uint4 v0 = *reinterpret_cast<uint4 const*>(src_ptr);
196
- uint4 v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
197
-
198
- {
199
- typename Converter::packed_hip_type const* p0 =
200
- reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
201
- typename Converter::packed_hip_type const* p1 =
202
- reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
203
-// Each packed_hip_type holds 2 bf16 → 4 packed = 8 elems per uint4.
204
-#pragma unroll
205
- for (int i = 0; i < 4; i++) {
206
- float2 f2 = Converter::convert(p0[i]);
207
- elements[2 * i] = f2.x;
208
- elements[2 * i + 1] = f2.y;
209
- }
210
-#pragma unroll
211
- for (int i = 0; i < 4; i++) {
212
- float2 f2 = Converter::convert(p1[i]);
213
- elements[8 + 2 * i] = f2.x;
214
- elements[8 + 2 * i + 1] = f2.y;
215
- }
216
- }
388
+ processDeepseekV4Slot<scalar_t_in>(
389
+ v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_inout,
390
+ k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
391
+ kv_block_stride);
392
218
- // ── Q branch: RMSNorm with no weight (has_weight=False) ─────────────────
219
- // Variance + rsqrt + multiply all in fp32, no intermediate bf16 round.
220
- // The downstream bf16 round only happens at the final store.
221
- if (!isKV) {
222
-#pragma unroll
223
- for (int i = 0; i < kElemsPerLane; i++) {
224
- sumOfSquares += elements[i] * elements[i];
225
- }
226
- sumOfSquares = warpSum<float>(sumOfSquares);
227
- float const rms_rcp =
228
- rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
229
-#pragma unroll
230
- for (int i = 0; i < kElemsPerLane; i++) {
231
- elements[i] = elements[i] * rms_rcp;
232
- }
233
- }
234
-
235
- // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
236
- // All math in fp32. cos_sin_cache is loaded as fp32 (its native storage).
237
- bool const is_rope_lane = dim_base >= kNopeDim;
238
- if (is_rope_lane) {
239
- int64_t const pos = position_ids[tokenIdx];
240
- constexpr int kHalfRope = kRopeDim / 2; // 32
241
- float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
242
- float const* sin_ptr = cos_ptr + kHalfRope;
243
-
244
- int const rope_local_base = dim_base - kNopeDim; // in [0, 64) step 16
245
-#pragma unroll
246
- for (int p = 0; p < kElemsPerLane / 2; p++) {
247
- int const pair_dim = rope_local_base + 2 * p;
248
- int const half_idx = pair_dim / 2;
249
- float const cos_v = VLLM_LDG(cos_ptr + half_idx);
250
- float const sin_v = VLLM_LDG(sin_ptr + half_idx);
251
- float const x_even = elements[2 * p];
252
- float const x_odd = elements[2 * p + 1];
253
- elements[2 * p] = x_even * cos_v - x_odd * sin_v;
254
- elements[2 * p + 1] = x_even * sin_v + x_odd * cos_v;
255
- }
256
- }
257
-
258
- // ═══════════════════════════════════════════════════════════════════════
259
- // Q branch: cast to bf16 and store back in place.
260
- // ═══════════════════════════════════════════════════════════════════════
261
- if (!isKV) {
262
- uint4 out0, out1;
263
- typename Converter::packed_hip_type* po0 =
264
- reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
265
- typename Converter::packed_hip_type* po1 =
266
- reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
267
-#pragma unroll
268
- for (int i = 0; i < 4; i++) {
269
- po0[i] = Converter::convert(
270
- make_float2(elements[2 * i], elements[2 * i + 1]));
271
- }
272
-#pragma unroll
273
- for (int i = 0; i < 4; i++) {
274
- po1[i] = Converter::convert(
275
- make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
276
- }
277
- scalar_t_in* dst =
278
- q_inout +
279
- (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
280
- dim_base;
281
- *reinterpret_cast<uint4*>(dst) = out0;
282
- *reinterpret_cast<uint4*>(dst + 8) = out1;
393
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
284
- cudaTriggerProgrammaticLaunchCompletion();
394
+ cudaTriggerProgrammaticLaunchCompletion();
395
#endif
286
- return;
287
- }
288
-
289
- // ═══════════════════════════════════════════════════════════════════════
290
- // KV branch.
291
- // ═══════════════════════════════════════════════════════════════════════
292
- int64_t const slot_id = slot_mapping[tokenIdx];
293
- if (slot_id < 0) {
294
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
295
- cudaTriggerProgrammaticLaunchCompletion();
396
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
397
+ }
398
#endif
297
- return;
298
- }
399
+}
400
300
- int64_t const block_idx = slot_id / cache_block_size;
301
- int64_t const pos_in_block = slot_id % cache_block_size;
302
- uint8_t* block_base =
303
- k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
304
- uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
305
- uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
306
- uint8_t* token_scale_ptr =
307
- block_base + static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
308
- pos_in_block * kScaleBytesPerToken;
309
-
310
- // Round K to bf16 first, matching the unfused reference path where K is
311
- // materialized as bf16 before K quantization. absmax, clamp, and FP8
312
- // quant below all run on these bf16-rounded values.
313
-#pragma unroll
314
- for (int i = 0; i < kElemsPerLane; i++) {
315
- elements[i] = Converter::convert(Converter::convert(elements[i]));
316
- }
401
+// ────────────────────────────────────────────────────────────────────────────
402
+// Kernel
403
+// ────────────────────────────────────────────────────────────────────────────
404
+//
405
+// Grid: 1D, gridDim.x = num_tokens_full
406
+// Block: blockDim.x = 256 threads (8 warps per block) Each
407
+// warp handles one token, iterating over each head.
408
+// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q
409
+// KV branch (RoPE + UE8M0 quant + insert)
410
+//
411
+template <typename scalar_t_in>
412
+__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
413
+ scalar_t_in* __restrict__ q_inout, // [N, H, 512] bf16, in place
414
+ scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache,
415
+ int64_t const* __restrict__ slot_mapping,
416
+ int64_t const* __restrict__ position_ids,
417
+ float const* __restrict__ cos_sin_cache, float const eps,
418
+ int const num_tokens_full, int const num_tokens_insert,
419
+ int const num_heads_q, int const cache_block_size,
420
+ int const kv_block_stride) {
421
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
422
+ if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
423
+ return;
424
+ } else {
425
+#endif
426
+ int const warpsPerBlock = blockDim.x / 32;
427
+ int const warpId = threadIdx.x / 32;
428
+ int const laneId = threadIdx.x % 32;
429
318
- // Per-quant-block absmax must be computed by ALL 32 lanes (warp-collective
319
- // shuffle requires full participation). RoPE lanes contribute garbage,
320
- // but their values are gated out below via `!is_rope_lane`.
321
- float local_absmax = 0.0f;
322
-#pragma unroll
323
- for (int i = 0; i < kElemsPerLane; i++) {
324
- local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
325
- }
326
- float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
327
- float const exponent = ceilf(log2f(absmax / kFp8Max));
328
- float const inv_scale = exp2f(-exponent);
430
+ int const tokenIdx = blockIdx.x;
431
+ if (tokenIdx >= num_tokens_full) return;
432
330
- if (!is_rope_lane) {
331
- // ── NoPE lane: UE8M0 FP8 quant ───────────────────────────────────────
332
- uint8_t out_bytes[kElemsPerLane];
333
-#pragma unroll
334
- for (int i = 0; i < kElemsPerLane; i++) {
335
- float scaled = elements[i] * inv_scale;
336
- scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
337
-#ifndef USE_ROCM
338
- __nv_fp8_storage_t s =
339
- __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
340
- out_bytes[i] = static_cast<uint8_t>(s);
341
-#else
342
- out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
433
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
434
+ cudaGridDependencySynchronize();
435
#endif
436
+
437
+ int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16
438
+ int const slot_end =
439
+ (tokenIdx >= num_tokens_insert) ? num_heads_q : (num_heads_q + 1);
440
+
441
+ auto src_for_slot = [&](int s) -> scalar_t_in const* {
442
+ if (s == num_heads_q) {
443
+ return kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
444
}
345
- // One 16-byte STG per lane.
346
- *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
347
- *reinterpret_cast<uint4 const*>(out_bytes);
348
-
349
- // Lane (4k) of each 4-lane group writes the scale byte for block k<7.
350
- if ((laneId & 3) == 0) {
351
- int const q_block_idx = laneId >> 2; // 0..6 for NoPE lanes
352
- float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
353
- token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
354
- }
355
- // Lane 0 also writes the padding byte at index 7.
356
- if (laneId == 0) {
357
- token_scale_ptr[kNumQuantBlocks] = 0; // pad
358
- }
359
- } else {
360
- // ── RoPE lane: cast back to bf16 and store to cache bf16 tail ────────
361
- uint4 out0, out1;
362
- typename Converter::packed_hip_type* po0 =
363
- reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
364
- typename Converter::packed_hip_type* po1 =
365
- reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
366
-#pragma unroll
367
- for (int i = 0; i < 4; i++) {
368
- po0[i] = Converter::convert(
369
- make_float2(elements[2 * i], elements[2 * i + 1]));
370
- }
371
-#pragma unroll
372
- for (int i = 0; i < 4; i++) {
373
- po1[i] = Converter::convert(
374
- make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
375
- }
376
- int const rope_local_base = dim_base - kNopeDim; // in [0, 64)
377
- scalar_t_in* bf16_dst =
378
- reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
379
- *reinterpret_cast<uint4*>(bf16_dst) = out0;
380
- *reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
381
- }
445
+ return q_inout +
446
+ (static_cast<int64_t>(tokenIdx) * num_heads_q +
447
+ static_cast<int64_t>(s)) *
448
+ kHeadDim +
449
+ dim_base;
450
+ };
451
+
452
+ if (warpId < slot_end) {
453
+ int curr_slot = warpId;
454
+ scalar_t_in const* src_curr = src_for_slot(curr_slot);
455
+ uint4 v0_curr = *reinterpret_cast<uint4 const*>(src_curr);
456
+ uint4 v1_curr = *reinterpret_cast<uint4 const*>(src_curr + 8);
457
+
458
+ while (curr_slot < slot_end) {
459
+ int const next_slot = curr_slot + warpsPerBlock;
460
+ bool const has_next = (next_slot < slot_end);
461
+
462
+ // Prefetch src for the next slot
463
+ uint4 v0_next, v1_next;
464
+ if (has_next) {
465
+ scalar_t_in const* src_next = src_for_slot(next_slot);
466
+ v0_next = *reinterpret_cast<uint4 const*>(src_next);
467
+ v1_next = *reinterpret_cast<uint4 const*>(src_next + 8);
468
+ }
469
+
470
+ processDeepseekV4Slot<scalar_t_in>(
471
+ v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
472
+ num_heads_q, eps, q_inout, k_cache, slot_mapping, position_ids,
473
+ cos_sin_cache, cache_block_size, kv_block_stride);
474
+
475
+ // ── Buffer rotation: hand the prefetched LDGs to the next iter.
476
+ v0_curr = v0_next;
477
+ v1_curr = v1_next;
478
+ curr_slot = next_slot;
479
+ } // while
480
+ } // if (warpId < slot_end)
481
+
482
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
483
cudaTriggerProgrammaticLaunchCompletion();
484
#endif
@@ -430,11 +530,22 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
530
config.attrs = attrs;
531
config.numAttrs = (sm_version >= 90) ? 1 : 0;
532
433
- cudaLaunchKernelEx(
434
- &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
435
- q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
436
- num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
437
- kv_block_stride);
533
+ if (num_tokens_full < NUM_TOKEN_CUTOFF) {
534
+ cudaLaunchKernelEx(
535
+ &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
536
+ q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
537
+ num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
538
+ kv_block_stride);
539
+ } else {
540
+ config.gridDim = dim3(num_tokens_full);
541
+ cudaLaunchKernelEx(
542
+ &config,
543
+ fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<scalar_t_in>,
544
+ q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
545
+ num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
546
+ kv_block_stride);
547
+ }
548
+
549
#else
550
// ROCm: use standard kernel launch syntax (no PDL/stream serialization)
551
// clang-format off
@@ -508,4 +619,4 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
619
num_tokens_full, num_tokens_insert, num_heads_q,
620
cache_block_size_i, kv_block_stride, stream);
621
});
511
-}
622
+}
\ No newline at end of file
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+27
-24
@@ -67,29 +67,26 @@ def apply_rope_gptj_last_k(
67
head_dim = x.shape[-1]
68
nope_dim = head_dim - rope_dim
69
70
- # Gather cos/sin for each token position: [num_tokens, rope_dim]
71
- cs = cos_sin_cache[positions].to(torch.float32) # [N, rope_dim]
72
- cos = cs[..., :half] # [N, half]
73
- sin = cs[..., half:] # [N, half]
74
-
75
- # Reshape leading dims so we can broadcast: x shape [..., head_dim].
76
- # Bring token dim to front; assume x is [num_tokens, ..., head_dim].
77
- # We rely on positions being per-token and all other dims sharing the same pos.
78
- rope = x[..., nope_dim:].float() # [..., rope_dim]
79
- # Make rope pairs: reshape last dim to [half, 2]
70
+ cs = cos_sin_cache[positions].to(torch.float32)
71
+ cos = cs[..., :half]
72
+ sin = cs[..., half:]
73
+
74
+ rope = x[..., nope_dim:].float()
75
shape = rope.shape
76
rope = rope.reshape(*shape[:-1], half, 2)
82
- even = rope[..., 0] # [..., half]
77
+ even = rope[..., 0]
78
odd = rope[..., 1]
79
85
- # Broadcast cos/sin over any heads dim in between. cos/sin are [N, half].
86
- # Add singleton dims for intermediate axes.
80
for _ in range(rope.ndim - 3):
81
cos = cos.unsqueeze(1)
82
sin = sin.unsqueeze(1)
83
91
- new_even = even * cos - odd * sin
92
- new_odd = even * sin + odd * cos
84
+ # Use addcmul (compiles to FMA on CUDA) for the 2x2 rotation. nvcc lowers
85
+ # the kernel's `e*c - o*s` to fma(e, c, -o*s); matching that here keeps
86
+ # near-cancellation pairs on the same bf16 grid as the kernel output and
87
+ # avoids spurious 1-ULP boundary flips at high num_tokens.
88
+ new_even = torch.addcmul(-odd * sin, even, cos)
89
+ new_odd = torch.addcmul(odd * cos, even, sin)
90
rope_rotated = torch.stack((new_even, new_odd), dim=-1).reshape(shape)
91
92
out = x.clone().float()
@@ -99,11 +96,15 @@ def apply_rope_gptj_last_k(
96
97
def rmsnorm_no_weight(x: torch.Tensor, eps: float) -> torch.Tensor:
98
"""RMSNorm with no learnable weight, matching
102
- `RMSNorm(head_dim, has_weight=False)`."""
103
- orig_dtype = x.dtype
99
+ `RMSNorm(head_dim, has_weight=False)`.
100
+
101
+ Returns fp32 so callers can chain RoPE without an intermediate bf16 round
102
+ (the kernel keeps the whole RMSNorm→RoPE pipeline in fp32 and rounds once
103
+ at the final store).
104
+ """
105
xf = x.float()
106
variance = xf.pow(2).mean(dim=-1, keepdim=True)
106
- return (xf * torch.rsqrt(variance + eps)).to(orig_dtype)
107
+ return xf * torch.rsqrt(variance + eps)
108
109
110
# ── Dispatch to the CUDA op (skip test cleanly if it isn't built in) ─────────
@@ -128,7 +129,7 @@ def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)
129
# ── Test 1: Q path numerical parity ──────────────────────────────────────────
130
131
131
-@pytest.mark.parametrize("num_tokens", [1, 4, 17, 64])
132
+@pytest.mark.parametrize("num_tokens", [1, 4, 17, 64, 2048])
133
@pytest.mark.parametrize("n_heads", [8, 64])
134
def test_q_path_matches_reference(num_tokens: int, n_heads: int):
135
torch.manual_seed(0)
@@ -142,8 +143,10 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int):
143
cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
144
145
# Reference: RMSNorm (no weight) per head, then GPT-J RoPE on last 64.
146
+ # Keep the chain in fp32 (rmsnorm_no_weight returns fp32) and round to
147
+ # bf16 once at the end, matching the kernel.
148
q_ref = rmsnorm_no_weight(q, eps)
146
- q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache)
149
+ q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache).to(dtype)
150
151
# Fused call with dummy KV tensors (KV branch will write slot_mapping=-1 → noop).
152
num_blocks = 2
@@ -173,7 +176,7 @@ def _ue8m0_per_block_scales(kv_roped_nope_f32: torch.Tensor, qblock: int):
176
return torch.pow(2.0, exponent) # [n_tok, n_blocks]
177
178
176
-@pytest.mark.parametrize("num_tokens", [1, 4, 17, 64])
179
+@pytest.mark.parametrize("num_tokens", [1, 4, 17, 64, 2048])
180
@pytest.mark.parametrize("block_size", [16, 64])
181
def test_kv_path_matches_reference(num_tokens: int, block_size: int):
182
torch.manual_seed(1)
@@ -261,7 +264,7 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int):
264
# ── Test 2b: DP padding (slot_mapping shorter than q/kv) ─────────────────────
265
266
264
-@pytest.mark.parametrize("num_tokens", [4, 17])
267
+@pytest.mark.parametrize("num_tokens", [4, 17, 2048])
268
@pytest.mark.parametrize("pad", [1, 5])
269
@pytest.mark.parametrize("block_size", [16, 64])
270
def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
@@ -312,7 +315,7 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
315
# ── Test 3: combined single-call Q + KV parity ───────────────────────────────
316
317
315
-@pytest.mark.parametrize("num_tokens", [1, 4, 17])
318
+@pytest.mark.parametrize("num_tokens", [1, 4, 17, 2048])
319
@pytest.mark.parametrize("n_heads", [8, 64])
320
@pytest.mark.parametrize("block_size", [16, 64])
321
def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
@@ -332,7 +335,7 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
335
336
# Reference.
337
q_ref = rmsnorm_no_weight(q, eps)
335
- q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache)
338
+ q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache).to(dtype)
339
kv_ref = apply_rope_gptj_last_k(kv, positions, cos_sin_cache)
340
k_cache_ref = torch.zeros(
341
num_blocks, block_size * HEAD_BYTES, dtype=torch.uint8, device=device