[Kernel] Optional Q-norm in fused DSv4 MLA epilogue; group_size=32 for packed FP8 quant (#56215)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Yongye Zhu committed
Sep 10, 2026 at 00:35 UTC
be1cb9834b3d74114f4d426068d2bce57c9912a0
6 files changed
+267
-148
csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+76
-58
@@ -3,7 +3,7 @@
3
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
*
5
* Horizontally-fused DeepseekV4-MLA kernel:
6
- * - Q side: per-head RMSNorm (no weight) + GPT-J RoPE on last ROPE_DIM
6
+ * - Q side: optional per-head RMSNorm + GPT-J RoPE on last ROPE_DIM
7
* - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged
8
* cache insert
9
*
@@ -167,7 +167,8 @@ __device__ __forceinline__ float warpSum(float val) {
167
// q_out stride fold to compile-time constants.
168
//
169
// Slot layout (per token):
170
-// slot < num_heads_q → live-Q (RMSNorm + RoPE,
170
+// slot < num_heads_q → live-Q (optional RMSNorm +
171
+// RoPE,
172
// read q_in →
173
// write q_out)
174
// num_heads_q <= slot < kNumHeadsQPadded → pad-Q (zero-fill q_out;
@@ -175,7 +176,7 @@ __device__ __forceinline__ float warpSum(float val) {
176
// slot == kNumHeadsQPadded → KV (RoPE + UE8M0 quant
177
// + paged-cache
178
// insert)
178
-template <typename scalar_t_in, int kNumHeadsQPadded>
179
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
180
__device__ __forceinline__ void processDeepseekV4Slot(
181
uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx,
182
int const dim_base, int const laneId, int const num_heads_q,
@@ -224,19 +225,21 @@ __device__ __forceinline__ void processDeepseekV4Slot(
225
}
226
}
227
227
- // ── Q branch: RMSNorm (no weight) ───────────────────────────────────
228
- if (!isKV) {
229
- float sumOfSquares = 0.0f;
228
+ // ── Q branch: optional RMSNorm (no weight) ──────────────────────────
229
+ if constexpr (APPLY_Q_NORM) {
230
+ if (!isKV) {
231
+ float sumOfSquares = 0.0f;
232
#pragma unroll
231
- for (int i = 0; i < kElemsPerLane; i++) {
232
- sumOfSquares += elements[i] * elements[i];
233
- }
234
- sumOfSquares = warpSum<float>(sumOfSquares);
235
- float const rms_rcp =
236
- rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
233
+ for (int i = 0; i < kElemsPerLane; i++) {
234
+ sumOfSquares += elements[i] * elements[i];
235
+ }
236
+ sumOfSquares = warpSum<float>(sumOfSquares);
237
+ float const rms_rcp =
238
+ rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
239
#pragma unroll
238
- for (int i = 0; i < kElemsPerLane; i++) {
239
- elements[i] = elements[i] * rms_rcp;
240
+ for (int i = 0; i < kElemsPerLane; i++) {
241
+ elements[i] = elements[i] * rms_rcp;
242
+ }
243
}
244
}
245
@@ -386,7 +389,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
389
// warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
390
// warp handles one (token, head_slot) pair.
391
// slot < num_heads_q → live-Q branch
389
-// (RMSNorm + RoPE,
392
+// (optional RMSNorm + RoPE,
393
// read q_in → write q_out)
394
// num_heads_q <= slot < kNumHeadsQPadded → pad-Q branch
395
// (zero-fill q_out)
@@ -404,7 +407,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
407
// attention uses them). The KV branch only inserts the first
408
// `num_tokens_insert` tokens (= slot_mapping length) into the paged cache.
409
//
407
-template <typename scalar_t_in, int kNumHeadsQPadded>
410
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
411
__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
412
scalar_t_in const* __restrict__ q_in, // [N, num_heads_q, 512]
413
scalar_t_in* __restrict__ q_out, // [N, kNumHeadsQPadded, 512]
@@ -471,7 +474,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
474
v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
475
}
476
474
- processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
477
+ processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>(
478
v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out,
479
k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
480
kv_block_stride);
@@ -491,10 +494,10 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
494
// Grid: 1D, gridDim.x = num_tokens_full
495
// Block: blockDim.x = 256 threads (8 warps per block) Each
496
// warp handles one token, iterating over each head.
494
-// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q
497
+// Q branch (optional RMSNorm + RoPE, in place) head_slot == num_heads_q
498
// KV branch (RoPE + UE8M0 quant + insert)
499
//
497
-template <typename scalar_t_in, int kNumHeadsQPadded>
500
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
501
__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
502
scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out,
503
scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache,
@@ -558,7 +561,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
561
load_slot(next_slot, v0_next, v1_next);
562
}
563
561
- processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
564
+ processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>(
565
v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
566
num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids,
567
cos_sin_cache, cache_block_size, kv_block_stride);
@@ -581,7 +584,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
584
// ────────────────────────────────────────────────────────────────────────────
585
// Launch wrapper
586
// ────────────────────────────────────────────────────────────────────────────
584
-template <typename scalar_t_in, int kNumHeadsQPadded>
587
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
588
static void launchFusedDeepseekV4Templated(
589
scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
590
uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids,
@@ -622,19 +625,18 @@ static void launchFusedDeepseekV4Templated(
625
config.numAttrs = (sm_version >= 90) ? 1 : 0;
626
627
if (num_tokens_full < NUM_TOKEN_CUTOFF) {
625
- cudaLaunchKernelEx(
626
- &config,
627
- fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in,
628
- kNumHeadsQPadded>,
629
- q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
630
- eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
631
- kv_block_stride);
628
+ cudaLaunchKernelEx(&config,
629
+ fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<
630
+ scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>,
631
+ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
632
+ cos_sin_cache, eps, num_tokens_full, num_tokens_insert,
633
+ num_heads_q, cache_block_size, kv_block_stride);
634
} else {
635
config.gridDim = dim3(num_tokens_full);
636
cudaLaunchKernelEx(
637
&config,
638
fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<
637
- scalar_t_in, kNumHeadsQPadded>,
639
+ scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>,
640
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
641
eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
642
kv_block_stride);
@@ -642,7 +644,8 @@ static void launchFusedDeepseekV4Templated(
644
#else
645
// ROCm: use standard kernel launch syntax (no PDL/stream serialization)
646
// clang-format off
645
- fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in, kNumHeadsQPadded>
647
+ fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<
648
+ scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>
649
<<<grid, kBlockSize, 0, stream>>>(
650
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
651
cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q,
@@ -660,13 +663,20 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
663
int const num_tokens_full, int const num_tokens_insert,
664
int const num_heads_q, int const num_heads_q_padded,
665
int const cache_block_size, int const kv_block_stride,
663
- cudaStream_t stream) {
666
+ bool const apply_q_norm, cudaStream_t stream) {
667
#define DISPATCH(N) \
668
case N: \
666
- launchFusedDeepseekV4Templated<scalar_t_in, N>( \
667
- q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \
668
- cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, \
669
- cache_block_size, kv_block_stride, stream); \
669
+ if (apply_q_norm) { \
670
+ launchFusedDeepseekV4Templated<scalar_t_in, N, true>( \
671
+ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \
672
+ cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \
673
+ num_heads_q, cache_block_size, kv_block_stride, stream); \
674
+ } else { \
675
+ launchFusedDeepseekV4Templated<scalar_t_in, N, false>( \
676
+ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \
677
+ cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \
678
+ num_heads_q, cache_block_size, kv_block_stride, stream); \
679
+ } \
680
return;
681
682
switch (num_heads_q_padded) {
@@ -700,7 +710,8 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
710
// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) / warps).
711
// Each warp handles one (token, slot): slot < num_heads_q → Q, slot ==
712
// num_heads_q → KV.
703
-template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8>
713
+template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8,
714
+ bool APPLY_Q_NORM>
715
__global__ void fusedDeepseekV4FullCacheKernel(
716
scalar_t_in* __restrict__ q_inout, // [N, H, 512], in place (bf16)
717
uint8_t* __restrict__ q_fp8_out, // [N, H, 512] fp8, optional
@@ -777,19 +788,21 @@ __global__ void fusedDeepseekV4FullCacheKernel(
788
}
789
}
790
780
- // ── Q branch: RMSNorm (no weight) ─────────────────────────────────────
781
- if (!isKV) {
782
- float sumOfSquares = 0.0f;
791
+ // ── Q branch: optional RMSNorm (no weight) ────────────────────────────
792
+ if constexpr (APPLY_Q_NORM) {
793
+ if (!isKV) {
794
+ float sumOfSquares = 0.0f;
795
#pragma unroll
784
- for (int i = 0; i < kElemsPerLane; i++) {
785
- sumOfSquares += elements[i] * elements[i];
786
- }
787
- sumOfSquares = warpSum<float>(sumOfSquares);
788
- float const rms_rcp =
789
- rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
796
+ for (int i = 0; i < kElemsPerLane; i++) {
797
+ sumOfSquares += elements[i] * elements[i];
798
+ }
799
+ sumOfSquares = warpSum<float>(sumOfSquares);
800
+ float const rms_rcp =
801
+ rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
802
#pragma unroll
791
- for (int i = 0; i < kElemsPerLane; i++) {
792
- elements[i] = elements[i] * rms_rcp;
803
+ for (int i = 0; i < kElemsPerLane; i++) {
804
+ elements[i] = elements[i] * rms_rcp;
805
+ }
806
}
807
}
808
@@ -899,15 +912,19 @@ static void launchFullCacheKernel(
912
float const* q_fp8_scale_inv, float const eps, int const num_tokens_full,
913
int const num_tokens_insert, int const num_heads_q,
914
int const cache_block_size, int64_t const kv_block_stride,
902
- int64_t const kv_token_stride, char const* op_name, cudaStream_t stream) {
915
+ int64_t const kv_token_stride, bool const apply_q_norm,
916
+ char const* op_name, cudaStream_t stream) {
917
constexpr int kBlockSize = 256;
918
constexpr int kWarpsPerBlock = kBlockSize / 32;
919
int64_t const total_warps =
920
static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
921
int const grid =
922
static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
909
- auto* kernel =
910
- fusedDeepseekV4FullCacheKernel<scalar_t_in, STORE_Q_FP8, STORE_KV_FP8>;
923
+ auto* kernel = apply_q_norm
924
+ ? fusedDeepseekV4FullCacheKernel<
925
+ scalar_t_in, STORE_Q_FP8, STORE_KV_FP8, true>
926
+ : fusedDeepseekV4FullCacheKernel<
927
+ scalar_t_in, STORE_Q_FP8, STORE_KV_FP8, false>;
928
#ifndef USE_ROCM
929
static int const sm_version = getSMVersion();
930
STD_TORCH_CHECK(sm_version >= 80, op_name,
@@ -950,7 +967,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
967
torch::stable::Tensor const& position_ids, // [N] int64
968
torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16
969
int64_t q_head_padded, // padded Q head count for output
953
- double eps, int64_t cache_block_size) {
970
+ double eps, int64_t cache_block_size, bool apply_q_norm) {
971
STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(),
972
"q_in must be contiguous CUDA");
973
STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
@@ -981,7 +998,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
998
"cos_sin_cache must be float32");
999
1000
// With DP padding, slot_mapping can be shorter than q/kv/positions.
984
- // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
1001
+ // Q processing runs on all q.size(0) rows (downstream attention uses them);
1002
// KV quant+insert runs only on the first slot_mapping.size(0) rows.
1003
int const num_tokens_full = static_cast<int>(q_in.size(0));
1004
int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
@@ -1000,7 +1017,8 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
1017
const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index());
1018
1019
// Allocate the padded q output. The kernel writes every element (live
1003
- // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
1020
+ // region gets optional RMSNorm+RoPE; pad region gets zeros), so `empty` is
1021
+ // safe.
1022
auto q_out = torch::stable::new_empty(
1023
q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
1024
@@ -1018,7 +1036,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
1036
cos_sin_cache.const_data_ptr<float>(), static_cast<float>(eps),
1037
num_tokens_full, num_tokens_insert, num_heads_q,
1038
num_heads_q_padded, cache_block_size_i, kv_block_stride,
1021
- stream);
1039
+ apply_q_norm, stream);
1040
});
1041
return q_out;
1042
}
@@ -1033,7 +1051,7 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
1051
torch::stable::Tensor const& slot_mapping, // [num_tokens_insert] int64
1052
torch::stable::Tensor const& position_ids, // [N] int64
1053
torch::stable::Tensor const& cos_sin_cache, // [max_pos, 64] float32
1036
- double eps, int64_t cache_block_size) {
1054
+ double eps, int64_t cache_block_size, bool apply_q_norm) {
1055
using torch::headeronly::ScalarType;
1056
STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
1057
"q must be contiguous CUDA");
@@ -1091,7 +1109,7 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
1109
cos_sin_cache.const_data_ptr<float>(), nullptr, nullptr,
1110
static_cast<float>(eps), num_tokens_full, num_tokens_insert,
1111
num_heads_q, static_cast<int>(cache_block_size), kv_block_stride,
1094
- kv_token_stride,
1112
+ kv_token_stride, apply_q_norm,
1113
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
1114
stream);
1115
});
@@ -1107,7 +1125,7 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
1125
torch::stable::Tensor const& cos_sin_cache, // [max_pos, 64] float32
1126
torch::stable::Tensor const& fp8_scale, // scalar float32 (KV scale)
1127
torch::stable::Tensor const& q_fp8_scale_inv, // scalar float32 (1 / Q scale)
1110
- double eps, int64_t cache_block_size) {
1128
+ double eps, int64_t cache_block_size, bool apply_q_norm) {
1129
using torch::headeronly::ScalarType;
1130
STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
1131
"q must be contiguous CUDA");
@@ -1181,7 +1199,7 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
1199
num_tokens_full, num_tokens_insert, num_heads_q,
1200
static_cast<int>(cache_block_size),
1201
// fp8 cache: 1 byte/element -> stride already in bytes.
1184
- k_cache.stride(0), k_cache.stride(1),
1202
+ k_cache.stride(0), k_cache.stride(1), apply_q_norm,
1203
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
1204
stream);
1205
});
csrc/libtorch_stable/ops.h
+3
-3
@@ -264,14 +264,14 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
264
torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
265
torch::stable::Tensor const& position_ids,
266
torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
267
- double eps, int64_t cache_block_size);
267
+ double eps, int64_t cache_block_size, bool apply_q_norm);
268
269
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
270
torch::stable::Tensor& q, torch::stable::Tensor const& kv,
271
torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
272
torch::stable::Tensor const& position_ids,
273
torch::stable::Tensor const& cos_sin_cache, double eps,
274
- int64_t cache_block_size);
274
+ int64_t cache_block_size, bool apply_q_norm);
275
276
void fused_kimi_k3_mla_key_concat_kv_cache_insert(
277
torch::stable::Tensor& q, torch::stable::Tensor const& k_nope,
@@ -346,7 +346,7 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
346
torch::stable::Tensor const& cos_sin_cache,
347
torch::stable::Tensor const& fp8_scale,
348
torch::stable::Tensor const& q_fp8_scale_inv, double eps,
349
- int64_t cache_block_size);
349
+ int64_t cache_block_size, bool apply_q_norm);
350
351
#ifndef USE_ROCM
352
std::tuple<torch::stable::Tensor, torch::stable::Tensor>
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+80
-71
@@ -295,8 +295,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
295
// Loads two contiguous uint4s (16 B + 16 B = 32 B) per thread; on Blackwell
296
// nvcc fuses these into a single 256-bit LDG.E.256.
297
//
298
-// Constraints: GROUP_SIZE % (THREADS_PER_GROUP * VEC_SIZE) == 0; for
299
-// THREADS_PER_GROUP=8 and bf16/fp16 (VEC_SIZE=16), this means GROUP_SIZE=128.
298
+// Each thread quantizes 16 bf16/fp16 values, using 2 or 8 threads per group.
299
template <typename T, typename DST_DTYPE, int GROUP_SIZE, int kGroupsPerBlockX,
300
int kRowsPerBlock>
301
__global__ void per_token_group_quant_8bit_packed_register_kernel(
@@ -305,9 +304,9 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
304
const int groups_per_row, const int mn, const int output_q_mn_extent,
305
const int tma_aligned_mn, const int64_t num_scale_elems, const float eps,
306
const float min_8bit, const float max_8bit) {
308
- static_assert(GROUP_SIZE == 128, "fast path supports GROUP_SIZE==128");
309
- constexpr int THREADS_PER_GROUP = 8;
307
+ static_assert(GROUP_SIZE == 32 || GROUP_SIZE == 128);
308
constexpr int VEC_SIZE = 32 / sizeof(T); // 16 for bf16/fp16
309
+ constexpr int THREADS_PER_GROUP = GROUP_SIZE / VEC_SIZE;
310
static_assert(GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE,
311
"GROUP_SIZE must equal THREADS_PER_GROUP * VEC_SIZE");
312
static_assert(32 % THREADS_PER_GROUP == 0,
@@ -363,20 +362,20 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
362
}
363
}
364
366
- // 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the
367
- // 8 lanes within the warp that share a group.
365
#ifdef USE_ROCM
366
const int lane_in_wave = threadIdx.x % warpSize;
370
- const unsigned long long mask = 0xFFull << (lane_in_wave & ~7);
371
- local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4, 8));
372
- local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2, 8));
373
- local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1, 8));
367
+ const unsigned long long mask = ((1ull << THREADS_PER_GROUP) - 1)
368
+ << (lane_in_wave & ~(THREADS_PER_GROUP - 1));
369
#else
375
- unsigned mask = 0xffu << (threadIdx.x & 24u);
376
- local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4));
377
- local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2));
378
- local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1));
370
+ const unsigned mask = ((1u << THREADS_PER_GROUP) - 1)
371
+ << ((threadIdx.x % 32) & ~(THREADS_PER_GROUP - 1));
372
#endif
373
+#pragma unroll
374
+ for (int offset = THREADS_PER_GROUP / 2; offset > 0; offset /= 2) {
375
+ local_absmax =
376
+ fmaxf(local_absmax,
377
+ __shfl_xor_sync(mask, local_absmax, offset, THREADS_PER_GROUP));
378
+ }
379
380
float y_s = local_absmax / max_8bit;
381
y_s = fmaxf(y_s, 1e-10f);
@@ -455,15 +454,15 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
454
}
455
456
// Public entry point: register-resident packed quant kernel.
458
-// Constraints: group_size == 128 and bf16/fp16 input.
457
+// Constraints: group_size in {32, 128} and bf16/fp16 input.
458
void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
459
torch::stable::Tensor& output_q,
460
torch::stable::Tensor& output_s_packed,
461
int64_t group_size, double eps,
462
double min_8bit, double max_8bit) {
464
- STD_TORCH_CHECK(group_size == 128,
463
+ STD_TORCH_CHECK(group_size == 32 || group_size == 128,
464
"per_token_group_quant_8bit_packed only supports "
466
- "group_size==128, got ",
465
+ "group_size in {32, 128}, got ",
466
group_size, ".");
467
const auto in_dtype = input.scalar_type();
468
STD_TORCH_CHECK(
@@ -512,7 +511,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
511
input.get_device_index());
512
cudaStream_t stream = get_current_cuda_stream();
513
515
- constexpr int THREADS_PER_GROUP = 8;
514
+ const int threads_per_group = group_size / 16;
515
const int64_t padded_groups_per_row = k_num_packed_sfk * 4;
516
const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn;
517
@@ -523,7 +522,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
522
const int ry = 16 / kx;
523
const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry;
524
const int64_t sf_k_blocks = padded_groups_per_row / kx;
526
- const int num_threads = (kx * ry) * THREADS_PER_GROUP;
525
+ const int num_threads = (kx * ry) * threads_per_group;
526
// CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535).
527
constexpr int64_t kMaxGridDimYZ = 65535;
528
STD_TORCH_CHECK(row_blocks <= static_cast<int64_t>(INT32_MAX) &&
@@ -536,63 +535,72 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
535
// PDL (Programmatic Dependent Launch) is NVIDIA-only; ROCm/HIP has no
536
// equivalent launch attribute, so fall back to a classic launch there.
537
#ifndef USE_ROCM
539
- #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
540
- do { \
541
- cudaLaunchConfig_t config = {}; \
542
- config.gridDim = dim3(static_cast<unsigned int>(row_blocks), \
543
- static_cast<unsigned int>(sf_k_blocks)); \
544
- config.blockDim = dim3(num_threads); \
545
- config.dynamicSmemBytes = 0; \
546
- config.stream = stream; \
547
- cudaLaunchAttribute attrs[1]; \
548
- attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
549
- attrs[0].val.programmaticStreamSerializationAllowed = 1; \
550
- config.numAttrs = 1; \
551
- config.attrs = attrs; \
552
- cudaLaunchKernelEx( \
553
- &config, \
554
- per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, \
555
- KX, RY>, \
556
- static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
557
- reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
558
- static_cast<int>(padded_groups_per_row), \
559
- static_cast<int>(groups_per_row), static_cast<int>(mn), \
560
- static_cast<int>(output_q_mn_extent), \
561
- static_cast<int>(tma_aligned_mn), num_scale_elems, \
562
- static_cast<float>(eps), static_cast<float>(min_8bit), \
563
- static_cast<float>(max_8bit)); \
538
+ #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, GROUP_SIZE, KX, RY) \
539
+ do { \
540
+ cudaLaunchConfig_t config = {}; \
541
+ config.gridDim = dim3(static_cast<unsigned int>(row_blocks), \
542
+ static_cast<unsigned int>(sf_k_blocks)); \
543
+ config.blockDim = dim3(num_threads); \
544
+ config.dynamicSmemBytes = 0; \
545
+ config.stream = stream; \
546
+ cudaLaunchAttribute attrs[1]; \
547
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
548
+ attrs[0].val.programmaticStreamSerializationAllowed = 1; \
549
+ config.numAttrs = 1; \
550
+ config.attrs = attrs; \
551
+ cudaLaunchKernelEx( \
552
+ &config, \
553
+ per_token_group_quant_8bit_packed_register_kernel< \
554
+ T, DST_DTYPE, GROUP_SIZE, KX, RY>, \
555
+ static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
556
+ reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
557
+ static_cast<int>(padded_groups_per_row), \
558
+ static_cast<int>(groups_per_row), static_cast<int>(mn), \
559
+ static_cast<int>(output_q_mn_extent), \
560
+ static_cast<int>(tma_aligned_mn), num_scale_elems, \
561
+ static_cast<float>(eps), static_cast<float>(min_8bit), \
562
+ static_cast<float>(max_8bit)); \
563
} while (0)
564
#else
566
- #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
567
- do { \
568
- dim3 grid(static_cast<unsigned int>(row_blocks), \
569
- static_cast<unsigned int>(sf_k_blocks)); \
570
- dim3 block(num_threads); \
571
- per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
572
- RY> \
573
- <<<grid, block, 0, stream>>>( \
574
- static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
575
- reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
576
- static_cast<int>(padded_groups_per_row), \
577
- static_cast<int>(groups_per_row), static_cast<int>(mn), \
578
- static_cast<int>(output_q_mn_extent), \
579
- static_cast<int>(tma_aligned_mn), num_scale_elems, \
580
- static_cast<float>(eps), static_cast<float>(min_8bit), \
581
- static_cast<float>(max_8bit)); \
565
+ #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, GROUP_SIZE, KX, RY) \
566
+ do { \
567
+ dim3 grid(static_cast<unsigned int>(row_blocks), \
568
+ static_cast<unsigned int>(sf_k_blocks)); \
569
+ dim3 block(num_threads); \
570
+ per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, \
571
+ GROUP_SIZE, KX, RY> \
572
+ <<<grid, block, 0, stream>>>( \
573
+ static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
574
+ reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
575
+ static_cast<int>(padded_groups_per_row), \
576
+ static_cast<int>(groups_per_row), static_cast<int>(mn), \
577
+ static_cast<int>(output_q_mn_extent), \
578
+ static_cast<int>(tma_aligned_mn), num_scale_elems, \
579
+ static_cast<float>(eps), static_cast<float>(min_8bit), \
580
+ static_cast<float>(max_8bit)); \
581
} while (0)
582
#endif
583
585
-#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
586
- do { \
587
- if (kx == 16) { \
588
- LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 16, 1); \
589
- } else if (kx == 8) { \
590
- LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 8, 2); \
591
- } else if (kx == 4) { \
592
- LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 4, 4); \
593
- } else { \
594
- STD_TORCH_CHECK(false, "Unsupported kx value ", kx); \
595
- } \
584
+#define LAUNCH_REG_KERNEL_GROUP_SIZE(T, DST_DTYPE, GROUP_SIZE) \
585
+ do { \
586
+ if (kx == 16) { \
587
+ LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, GROUP_SIZE, 16, 1); \
588
+ } else if (kx == 8) { \
589
+ LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, GROUP_SIZE, 8, 2); \
590
+ } else if (kx == 4) { \
591
+ LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, GROUP_SIZE, 4, 4); \
592
+ } else { \
593
+ STD_TORCH_CHECK(false, "Unsupported kx value ", kx); \
594
+ } \
595
+ } while (0)
596
+
597
+#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
598
+ do { \
599
+ if (group_size == 32) { \
600
+ LAUNCH_REG_KERNEL_GROUP_SIZE(T, DST_DTYPE, 32); \
601
+ } else { \
602
+ LAUNCH_REG_KERNEL_GROUP_SIZE(T, DST_DTYPE, 128); \
603
+ } \
604
} while (0)
605
606
VLLM_STABLE_DISPATCH_HALF_TYPES(
@@ -607,6 +615,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
615
}));
616
617
#undef LAUNCH_REG_KERNEL
618
+#undef LAUNCH_REG_KERNEL_GROUP_SIZE
619
#undef LAUNCH_REG_KERNEL_INST
620
}
621
csrc/libtorch_stable/torch_bindings.cpp
+4
-3
@@ -429,7 +429,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
429
"fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
430
"Tensor q_in, Tensor kv, Tensor! k_cache, "
431
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
432
- "int q_head_padded, float eps, int cache_block_size) -> Tensor");
432
+ "int q_head_padded, float eps, int cache_block_size, "
433
+ "bool apply_q_norm=True) -> Tensor");
434
435
// FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
436
// FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
@@ -437,13 +438,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
438
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert("
439
"Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, "
440
"Tensor position_ids, Tensor cos_sin_cache, float eps, "
440
- "int cache_block_size) -> ()");
441
+ "int cache_block_size, bool apply_q_norm=True) -> ()");
442
ops.def(
443
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert("
444
"Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, "
445
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
446
"Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, "
446
- "int cache_block_size) -> ()");
447
+ "int cache_block_size, bool apply_q_norm=True) -> ()");
448
449
// Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion.
450
ops.def(
tests/kernels/quantization/test_per_token_group_quant.py
+5
@@ -80,6 +80,11 @@ def test_per_token_group_quant_fp8(
80
# Larger shapes with padding
81
(127, 7168, 128),
82
(253, 640, 128),
83
+ (1, 32, 32),
84
+ (3, 64, 32),
85
+ (7, 160, 32),
86
+ (32, 512, 32),
87
+ (127, 4096, 32),
88
],
89
)
90
@pytest.mark.parametrize("poisoned_scales", [False, True])
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+99
-13
@@ -4,11 +4,11 @@
4
Standalone unit test for the horizontally-fused DeepseekV4-MLA kernel:
5
6
fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
7
- - Q side: per-head RMSNorm (no weight) + GPT-J RoPE on last 64 dims
7
+ - Q side: optional per-head RMSNorm + GPT-J RoPE on last 64 dims
8
- KV side: GPT-J RoPE on last 64 + UE8M0 FP8 quant + paged cache insert
9
10
We compare against:
11
- - PyTorch reference for RMSNorm + GPT-J RoPE on Q
11
+ - PyTorch references for RoPE with and without RMSNorm on Q
12
- Existing Triton `quantize_and_insert_k_cache` + round-trip via
13
`dequantize_and_gather_k_cache` for KV
14
@@ -146,9 +146,18 @@ pytestmark = pytest.mark.skipif(
146
147
148
def _call_fused(
149
- q_in, q_head_padded, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
149
+ q_in,
150
+ q_head_padded,
151
+ kv,
152
+ k_cache,
153
+ slot_mapping,
154
+ positions,
155
+ cos_sin_cache,
156
+ eps,
157
+ bs,
158
+ apply_q_norm=None,
159
):
151
- return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
160
+ args = (
161
q_in,
162
kv,
163
k_cache,
@@ -159,6 +168,11 @@ def _call_fused(
168
eps,
169
bs,
170
)
171
+ if apply_q_norm is None:
172
+ return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(*args)
173
+ return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
174
+ *args, apply_q_norm
175
+ )
176
177
178
def _as_stored_fp8(t: torch.Tensor) -> torch.Tensor:
@@ -269,6 +283,42 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: i
283
)
284
285
286
+@pytest.mark.parametrize("num_tokens", [4, 2048])
287
+def test_q_path_without_qnorm_matches_rope_only_reference(num_tokens: int):
288
+ """Disabling Q norm must apply RoPE directly to the projected query."""
289
+ torch.manual_seed(7)
290
+ device = "cuda"
291
+ dtype = torch.bfloat16
292
+ eps = 1e-6
293
+ n_heads = 8
294
+ padded_heads = 16
295
+ block_size = 16
296
+
297
+ q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
298
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
299
+ cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device)
300
+ q_ref = apply_rope_gptj_last_k(q, positions, cos_sin_cache)
301
+ kv = torch.zeros(num_tokens, HEAD_DIM, dtype=dtype, device=device)
302
+ k_cache = torch.zeros(2, block_size * HEAD_BYTES, dtype=torch.uint8, device=device)
303
+ slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
304
+
305
+ q_out = _call_fused(
306
+ q,
307
+ padded_heads,
308
+ kv,
309
+ k_cache,
310
+ slot_mapping,
311
+ positions,
312
+ cos_sin_cache,
313
+ eps,
314
+ block_size,
315
+ apply_q_norm=False,
316
+ )
317
+
318
+ torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=0, atol=0)
319
+ assert q_out[:, n_heads:].count_nonzero().item() == 0
320
+
321
+
322
# ── Test 2: KV path round-trip byte/value parity ─────────────────────────────
323
324
@@ -382,7 +432,7 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int):
432
@pytest.mark.parametrize("block_size", [16, 64])
433
def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
434
"""slot_mapping.size(0) < q.size(0): the kernel must skip padded
385
- tokens in the KV branch while still running Q-norm+RoPE on all rows."""
435
+ tokens in the KV branch while still processing Q on all rows."""
436
torch.manual_seed(3)
437
device = "cuda"
438
dtype = torch.bfloat16
@@ -517,8 +567,9 @@ def _call_full_cache_fp8_fused(
567
q_fp8_scale_inv,
568
eps,
569
bs,
570
+ apply_q_norm=None,
571
):
521
- torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
572
+ args = (
573
q,
574
kv,
575
q_fp8,
@@ -531,6 +582,11 @@ def _call_full_cache_fp8_fused(
582
eps,
583
bs,
584
)
585
+ op = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert
586
+ if apply_q_norm is None:
587
+ op(*args)
588
+ else:
589
+ op(*args, apply_q_norm)
590
591
592
def _call_full_cache_bf16_fused(
@@ -542,8 +598,9 @@ def _call_full_cache_bf16_fused(
598
cos_sin_cache,
599
eps,
600
bs,
601
+ apply_q_norm=None,
602
):
546
- torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
603
+ args = (
604
q,
605
kv,
606
k_cache,
@@ -553,6 +610,11 @@ def _call_full_cache_bf16_fused(
610
eps,
611
bs,
612
)
613
+ op = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert
614
+ if apply_q_norm is None:
615
+ op(*args)
616
+ else:
617
+ op(*args, apply_q_norm)
618
619
620
def _fp8_full_cache_reference(
@@ -567,8 +629,9 @@ def _fp8_full_cache_reference(
629
block_size,
630
fp8_scale,
631
q_fp8_scale_inv,
632
+ apply_q_norm,
633
):
571
- q_ref = rmsnorm_no_weight(q, eps)
634
+ q_ref = rmsnorm_no_weight(q, eps) if apply_q_norm else q.float()
635
q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache)
636
q_fp8.copy_(
637
torch.clamp(q_ref.float() * q_fp8_scale_inv, -FP8_MAX, FP8_MAX).to(
@@ -595,9 +658,10 @@ def _bf16_full_cache_reference(
658
cos_sin_cache,
659
eps,
660
block_size,
661
+ apply_q_norm,
662
):
599
- q_ref = rmsnorm_no_weight(q, eps)
600
- # Kernel keeps RMSNorm+RoPE in fp32 and rounds to bf16 once at the store.
663
+ q_ref = rmsnorm_no_weight(q, eps) if apply_q_norm else q.float()
664
+ # Kernel keeps optional RMSNorm+RoPE in fp32 and rounds once at the store.
665
q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache).to(q.dtype)
666
667
kv_ref = apply_rope_gptj_last_k(kv, positions, cos_sin_cache)
@@ -616,10 +680,19 @@ def _bf16_full_cache_reference(
680
@pytest.mark.parametrize("num_tokens", [4, 17])
681
@pytest.mark.parametrize("n_heads", [8, 17])
682
@pytest.mark.parametrize("positions_dtype", [torch.int32, torch.int64])
683
+@pytest.mark.parametrize(
684
+ "apply_q_norm",
685
+ [
686
+ pytest.param(None, id="qnorm_default"),
687
+ pytest.param(True, id="qnorm_explicit"),
688
+ pytest.param(False, id="rope_only"),
689
+ ],
690
+)
691
def test_full_cache_per_tensor_fp8_matches_reference(
692
num_tokens: int,
693
n_heads: int,
694
positions_dtype: torch.dtype,
695
+ apply_q_norm: bool | None,
696
):
697
torch.manual_seed(4)
698
device = "cuda"
@@ -662,6 +735,7 @@ def test_full_cache_per_tensor_fp8_matches_reference(
735
block_size,
736
fp8_scale,
737
q_fp8_scale_inv,
738
+ apply_q_norm is not False,
739
)
740
_call_full_cache_fp8_fused(
741
q.clone(),
@@ -675,11 +749,12 @@ def test_full_cache_per_tensor_fp8_matches_reference(
749
q_fp8_scale_inv,
750
eps,
751
block_size,
752
+ apply_q_norm,
753
)
754
680
- # Q is RMSNorm(no-weight)+RoPE in fp32 before fp8 quant; the RMSNorm
681
- # reduction and RoPE rotation can land the kernel and the torch reference on
682
- # opposite sides of an fp8 round-to-nearest tie, so allow <=1 fp8 ULP.
755
+ # Q uses optional RMSNorm followed by RoPE in fp32 before fp8 quant. The
756
+ # kernel and torch reference can land on opposite sides of an fp8
757
+ # round-to-nearest tie, so allow <=1 fp8 ULP.
758
q_fused = _as_stored_fp8(q_fp8_fused)
759
q_max_ulp = int(fp8_ulp_distance(q_fused, q_fp8_ref).max().item())
760
assert q_max_ulp <= 1, f"Q fp8 differs by {q_max_ulp} ULP (>1)"
@@ -709,10 +784,19 @@ def test_full_cache_per_tensor_fp8_matches_reference(
784
@pytest.mark.parametrize("num_tokens", [4, 17])
785
@pytest.mark.parametrize("n_heads", [8, 17])
786
@pytest.mark.parametrize("positions_dtype", [torch.int32, torch.int64])
787
+@pytest.mark.parametrize(
788
+ "apply_q_norm",
789
+ [
790
+ pytest.param(None, id="qnorm_default"),
791
+ pytest.param(True, id="qnorm_explicit"),
792
+ pytest.param(False, id="rope_only"),
793
+ ],
794
+)
795
def test_full_cache_bf16_matches_reference(
796
num_tokens: int,
797
n_heads: int,
798
positions_dtype: torch.dtype,
799
+ apply_q_norm: bool | None,
800
):
801
torch.manual_seed(5)
802
device = "cuda"
@@ -743,6 +827,7 @@ def test_full_cache_bf16_matches_reference(
827
cos_sin_cache,
828
eps,
829
block_size,
830
+ apply_q_norm is not False,
831
)
832
_call_full_cache_bf16_fused(
833
q_fused,
@@ -753,6 +838,7 @@ def test_full_cache_bf16_matches_reference(
838
cos_sin_cache,
839
eps,
840
block_size,
841
+ apply_q_norm,
842
)
843
844
torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)