From 9c483cccfe97e3ef9db9b7ff047aa54858d51476 Mon Sep 17 00:00:00 2001 From: Pranjal Shankhdhar Date: Fri, 24 Jul 2026 16:19:08 -0700 Subject: [PATCH] Support a same-size mixed q dtype in the fused RoPE kernels (#31834) Co-authored-by: pranjalssh --- .../kernels/jit/csrc/elementwise/rope.cuh | 191 +++++++++--------- python/sglang/kernels/ops/attention/rope.py | 12 +- python/sglang/srt/layers/radix_attention.py | 16 +- .../kernels/ops/attention/test_rope.py | 84 ++++++++ 4 files changed, 195 insertions(+), 108 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/elementwise/rope.cuh b/python/sglang/kernels/jit/csrc/elementwise/rope.cuh index 36f583e49..785acd890 100644 --- a/python/sglang/kernels/jit/csrc/elementwise/rope.cuh +++ b/python/sglang/kernels/jit/csrc/elementwise/rope.cuh @@ -46,17 +46,83 @@ constexpr auto next_pow2(uint32_t target, uint32_t factor = 1) { return power; } -template +// Rotate one head row in place, reading and writing as T. Factored out so the +// q and k branches can run different (same-size) dtypes -- e.g. fp16 q emitted +// by a fused QK-norm while k stays bf16. When cache_out is non-null the rotated +// row is additionally stored there (the fused KV-store k path); predicated so +// the q and k branches keep one instruction stream per instantiation. +template +__device__ __forceinline__ void +rope_rotate_head(void* input, const void* cos_ptr, const void* sin_ptr, uint32_t lane_id, void* cache_out = nullptr) { + using namespace device; + using T2 = packed_t; + using Storage = AlignedVector; + if constexpr (kIsNeox) { + using CacheStorage = AlignedVector; + const auto input_x = input; + const auto input_y = pointer::offset(input, (kRopeDim / 2) * sizeof(T)); + auto input_vec_x = load_as(input_x, lane_id); + auto input_vec_y = load_as(input_y, lane_id); + const auto cos_pair = load_as(cos_ptr, lane_id); + const auto sin_pair = load_as(sin_ptr, lane_id); +#pragma unroll + for (int64_t j = 0; j < kVecSize; ++j) { + const auto [x0, x1] = cast(input_vec_x[j]); + const auto [y0, y1] = cast(input_vec_y[j]); + const auto [cos_0, cos_1] = cos_pair[j]; + const auto [sin_0, sin_1] = sin_pair[j]; + const auto out_x0 = x0 * cos_0 - y0 * sin_0; + const auto out_y0 = x0 * sin_0 + y0 * cos_0; + const auto out_x1 = x1 * cos_1 - y1 * sin_1; + const auto out_y1 = x1 * sin_1 + y1 * cos_1; + input_vec_x[j] = cast({out_x0, out_x1}); + input_vec_y[j] = cast({out_y0, out_y1}); + } + store_as(input_x, input_vec_x, lane_id); + store_as(input_y, input_vec_y, lane_id); + if (cache_out != nullptr) { + store_as(cache_out, input_vec_x, lane_id); + const auto cache_out_y = pointer::offset(cache_out, (kRopeDim / 2) * sizeof(T)); + store_as(cache_out_y, input_vec_y, lane_id); + } + } else { + using CacheStorage = AlignedVector; + auto input_vec = load_as(input, lane_id); + const auto cos_vec = load_as(cos_ptr, lane_id); + const auto sin_vec = load_as(sin_ptr, lane_id); +#pragma unroll + for (int64_t j = 0; j < kVecSize; ++j) { + const auto [x, y] = cast(input_vec[j]); + const auto cos = cos_vec[j]; + const auto sin = sin_vec[j]; + const auto out_x = x * cos - y * sin; + const auto out_y = x * sin + y * cos; + input_vec[j] = cast({out_x, out_y}); + } + store_as(input, input_vec, lane_id); + if (cache_out != nullptr) { + store_as(cache_out, input_vec, lane_id); + } + } +} + +template < + bool kIsNeox, + int64_t kRopeDim, + bool kUsePDL, + typename DType, + typename QDType, + typename IdType, + uint32_t kWorkThreads> __global__ void fused_rope_kernel(const __grid_constant__ FusedRopeParams params) { using namespace device; constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float); constexpr int64_t kVecSize = next_pow2(kRopeDim, (2 * kWorkThreads * (1 + kIsNeox))); - using DType2 = packed_t; - using InputStorage = AlignedVector; constexpr int64_t kDimPerThread = kVecSize * 2 * (1 + kIsNeox); constexpr uint32_t kLaneCount = kRopeDim / kDimPerThread; static_assert(kRopeDim % kDimPerThread == 0 && kLaneCount <= kWorkThreads); + static_assert(sizeof(QDType) == sizeof(DType), "q/k must be same-size dtypes"); const auto &[ q, k, cos_sin_cache_ptr, positions, // pointers @@ -90,60 +156,32 @@ __global__ void fused_rope_kernel(const __grid_constant__ FusedRopeParams params const auto input = pointer::offset(input_, head_id * head_stride_bytes); const auto cos_ptr = pointer::offset(cos_cache_ptr, pos * kCosSinStrideBytes); const auto sin_ptr = pointer::offset(sin_cache_ptr, pos * kCosSinStrideBytes); - if constexpr (kIsNeox) { - using CacheStorage = AlignedVector; - const auto input_x = input; - const auto input_y = pointer::offset(input, (kRopeDim / 2) * sizeof(DType)); - auto input_vec_x = load_as(input_x, lane_id); - auto input_vec_y = load_as(input_y, lane_id); - const auto cos_pair = load_as(cos_ptr, lane_id); - const auto sin_pair = load_as(sin_ptr, lane_id); -#pragma unroll - for (int64_t j = 0; j < kVecSize; ++j) { - const auto [x0, x1] = cast(input_vec_x[j]); - const auto [y0, y1] = cast(input_vec_y[j]); - const auto [cos_0, cos_1] = cos_pair[j]; - const auto [sin_0, sin_1] = sin_pair[j]; - const auto out_x0 = x0 * cos_0 - y0 * sin_0; - const auto out_y0 = x0 * sin_0 + y0 * cos_0; - const auto out_x1 = x1 * cos_1 - y1 * sin_1; - const auto out_y1 = x1 * sin_1 + y1 * cos_1; - input_vec_x[j] = cast({out_x0, out_x1}); - input_vec_y[j] = cast({out_y0, out_y1}); - } - store_as(input_x, input_vec_x, lane_id); - store_as(input_y, input_vec_y, lane_id); + if (load_q) { + rope_rotate_head(input, cos_ptr, sin_ptr, lane_id); } else { - using CacheStorage = AlignedVector; - auto input_vec = load_as(input, lane_id); - const auto cos_vec = load_as(cos_ptr, lane_id); - const auto sin_vec = load_as(sin_ptr, lane_id); -#pragma unroll - for (int64_t j = 0; j < kVecSize; ++j) { - const auto [x, y] = cast(input_vec[j]); - const auto cos = cos_vec[j]; - const auto sin = sin_vec[j]; - const auto out_x = x * cos - y * sin; - const auto out_y = x * sin + y * cos; - input_vec[j] = cast({out_x, out_y}); - } - store_as(input, input_vec, lane_id); + rope_rotate_head(input, cos_ptr, sin_ptr, lane_id); } } PDLTriggerSecondary(); } -template +template < + bool kIsNeox, + int64_t kRopeDim, + bool kUsePDL, + typename DType, + typename QDType, + typename IdType, + uint32_t kWorkThreads> __global__ void fused_rope_store_kernel(const __grid_constant__ FusedRopeStoreParams params) { using namespace device; constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float); constexpr int64_t kVecSize = kRopeDim / (2 * kWorkThreads * (1 + kIsNeox)); - using DType2 = packed_t; - using InputStorage = AlignedVector; constexpr int64_t kDimPerThread = kVecSize * 2 * (1 + kIsNeox); static_assert(kRopeDim == kDimPerThread * kWorkThreads); + static_assert(sizeof(QDType) == sizeof(DType), "q/k must be same-size dtypes"); const auto& [base_params, v_ptr, k_cache, v_cache, out_loc, v_stride_bytes, cache_stride_bytes] = params; const auto &[ @@ -180,55 +218,12 @@ __global__ void fused_rope_store_kernel(const __grid_constant__ FusedRopeStorePa const auto input = pointer::offset(input_, head_id * head_stride_bytes); const auto cos_ptr = pointer::offset(cos_cache_ptr, pos * kCosSinStrideBytes); const auto sin_ptr = pointer::offset(sin_cache_ptr, pos * kCosSinStrideBytes); - if constexpr (kIsNeox) { - using CacheStorage = AlignedVector; - const auto input_x = input; - const auto input_y = pointer::offset(input, (kRopeDim / 2) * sizeof(DType)); - auto input_vec_x = load_as(input_x, lane_id); - auto input_vec_y = load_as(input_y, lane_id); - const auto cos_pair = load_as(cos_ptr, lane_id); - const auto sin_pair = load_as(sin_ptr, lane_id); -#pragma unroll - for (int64_t j = 0; j < kVecSize; ++j) { - const auto [x0, x1] = cast(input_vec_x[j]); - const auto [y0, y1] = cast(input_vec_y[j]); - const auto [cos_0, cos_1] = cos_pair[j]; - const auto [sin_0, sin_1] = sin_pair[j]; - const auto out_x0 = x0 * cos_0 - y0 * sin_0; - const auto out_y0 = x0 * sin_0 + y0 * cos_0; - const auto out_x1 = x1 * cos_1 - y1 * sin_1; - const auto out_y1 = x1 * sin_1 + y1 * cos_1; - input_vec_x[j] = cast({out_x0, out_x1}); - input_vec_y[j] = cast({out_y0, out_y1}); - } - store_as(input, input_vec_x, lane_id); - const auto input_y_out = pointer::offset(input, (kRopeDim / 2) * sizeof(DType)); - store_as(input_y_out, input_vec_y, lane_id); - if (!load_q) { - const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes); - store_as(k_out, input_vec_x, lane_id); - const auto k_out_y = pointer::offset(k_out, (kRopeDim / 2) * sizeof(DType)); - store_as(k_out_y, input_vec_y, lane_id); - } + if (load_q) { + // q rotates in place only (no cache store); may be a different dtype. + rope_rotate_head(input, cos_ptr, sin_ptr, lane_id); } else { - using CacheStorage = AlignedVector; - auto input_vec = load_as(input, lane_id); - const auto cos_vec = load_as(cos_ptr, lane_id); - const auto sin_vec = load_as(sin_ptr, lane_id); -#pragma unroll - for (int64_t j = 0; j < kVecSize; ++j) { - const auto [x, y] = cast(input_vec[j]); - const auto cos = cos_vec[j]; - const auto sin = sin_vec[j]; - const auto out_x = x * cos - y * sin; - const auto out_y = x * sin + y * cos; - input_vec[j] = cast({out_x, out_y}); - } - store_as(input, input_vec, lane_id); - if (!load_q) { - const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes); - store_as(k_out, input_vec, lane_id); - } + const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes); + rope_rotate_head(input, cos_ptr, sin_ptr, lane_id, k_out); } } @@ -247,18 +242,22 @@ __global__ void fused_rope_store_kernel(const __grid_constant__ FusedRopeStorePa PDLTriggerSecondary(); } -template +// QDType (default = DType) is q's dtype; k/v/caches use DType. Same-size +// dtypes only (fp16 q with bf16 k for the fp8-KV fp16-Q feed). +template struct FusedRopeKernel { static constexpr uint32_t kDimPerThread = std::gcd(16 / sizeof(DType), kRopeDim); static constexpr uint32_t kWorkThreads = next_pow2(kRopeDim, kDimPerThread); static constexpr bool kSupportFused = kWorkThreads * kDimPerThread == kRopeDim; static_assert(kRopeDim % kDimPerThread == 0); static_assert(kBlockSize % kWorkThreads == 0); + static_assert(sizeof(QDType) == sizeof(DType), "q/k must be same-size dtypes"); template - static constexpr auto _kernel_0 = fused_rope_kernel; + static constexpr auto _kernel_0 = fused_rope_kernel; template - static constexpr auto _kernel_1 = fused_rope_store_kernel; + static constexpr auto _kernel_1 = + fused_rope_store_kernel; static auto get_num_sm(DLDevice device) { static const auto kNumSM = host::runtime::get_sm_count(device.device_id); @@ -284,7 +283,7 @@ struct FusedRopeKernel { device.set_options(); TensorMatcher({N, Q, D}) // q input .with_strides({Dq, Dd, 1}) - .with_dtype() + .with_dtype() .with_device(device) .verify(q); TensorMatcher({N, K, D}) // k input @@ -382,7 +381,7 @@ struct FusedRopeKernel { TensorMatcher({N, Q, D}) // q input .with_strides({Dq, Dd, 1}) - .with_dtype() + .with_dtype() .with_device(device) .verify(q); TensorMatcher({N, K, D}) // k input diff --git a/python/sglang/kernels/ops/attention/rope.py b/python/sglang/kernels/ops/attention/rope.py index 26b2f78e9..a5f6be632 100644 --- a/python/sglang/kernels/ops/attention/rope.py +++ b/python/sglang/kernels/ops/attention/rope.py @@ -27,8 +27,12 @@ def _jit_rotary_embedding_module() -> Module: @cache_once -def _jit_fused_rope_module(is_neox: bool, rope_dim: int, dtype: torch.dtype) -> Module: - args = make_cpp_args(is_neox, rope_dim, is_arch_support_pdl(), dtype) +def _jit_fused_rope_module( + is_neox: bool, rope_dim: int, dtype: torch.dtype, q_dtype: torch.dtype +) -> Module: + # q_dtype supports a same-size mixed-dtype q, e.g. fp16 q with bf16 k; + # k/v/caches use dtype. + args = make_cpp_args(is_neox, rope_dim, is_arch_support_pdl(), dtype, q_dtype) return load_jit( "fused_rope", *args, @@ -133,7 +137,7 @@ def apply_rope_inplace( rope_dim: Rotary embedding dimension. Defaults to cos_sin_cache.size(-1). """ rope_dim = rope_dim or cos_sin_cache.size(-1) - module = _jit_fused_rope_module(is_neox, rope_dim, q.dtype) + module = _jit_fused_rope_module(is_neox, rope_dim, k.dtype, q.dtype) module.run_rope(q, k, cos_sin_cache, positions) @@ -171,7 +175,7 @@ def apply_rope_inplace_with_kvcache( """ rope_dim = rope_dim or cos_sin_cache.size(-1) v = v.view_as(k) - module = _jit_fused_rope_module(is_neox, rope_dim, q.dtype) + module = _jit_fused_rope_module(is_neox, rope_dim, k.dtype, q.dtype) module.run_rope_store(q, k, v, k_cache, v_cache, cos_sin_cache, positions, out_loc) diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index 4bcfa541b..876b0f00d 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/python/sglang/srt/layers/radix_attention.py @@ -195,14 +195,14 @@ class RadixAttention(nn.Module): idx_v=idx_v, ) return idx_out, attn_out - # FP8 q (e.g. mxfp8 KV-cache attention) still produces a bf16 - # attention output; sizing the buffer off q's dtype would silently - # cast-copy the result to fp8. - out_dtype = ( - torch.bfloat16 - if q.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) - else q.dtype - ) + # Output dtype follows v (the model dtype) when available: qk-norm + # may emit q in a different dtype without changing the dtype the + # backend writes. FP8 q/v (e.g. mxfp8 KV-cache attention) still + # produce a bf16 attention output; sizing the buffer off an fp8 + # dtype would silently cast-copy the result to fp8. + out_dtype = v.dtype if v is not None else q.dtype + if out_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + out_dtype = torch.bfloat16 if self.qk_head_dim != self.v_head_dim: output = q.new_empty( (q.shape[0], self.tp_q_head_num * self.v_head_dim), diff --git a/test/registered/kernels/ops/attention/test_rope.py b/test/registered/kernels/ops/attention/test_rope.py index 11bb5295f..b34608857 100644 --- a/test/registered/kernels/ops/attention/test_rope.py +++ b/test/registered/kernels/ops/attention/test_rope.py @@ -211,6 +211,90 @@ def test_rope_position_dtypes(dtype: torch.dtype) -> None: triton.testing.assert_close(k_fi, k_jit, atol=atol, rtol=rtol) +@pytest.mark.parametrize("is_neox", IS_NEOX_LIST) +@pytest.mark.parametrize("rope_dim", get_ci_test_range([64, 128], [64])) +def test_rope_mixed_q_dtype(is_neox: bool, rope_dim: int) -> None: + """fp16 q + bf16 k (a fused QK-norm may emit q in a different dtype): + q must match the all-fp16 kernel bitwise, k the all-bf16 kernel bitwise.""" + batch_size, num_kv_heads, gqa_ratio = 129, 4, 8 + num_qo_heads = num_kv_heads * gqa_ratio + q16 = torch.randn( + batch_size, num_qo_heads, rope_dim, device=DEVICE, dtype=torch.float16 + ) + kbf = torch.randn( + batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=torch.bfloat16 + ) + positions = torch.randint( + 0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=torch.int64 + ) + cos_sin_cache = create_cos_sin_cache(rope_dim) + + q_mixed, k_mixed = q16.clone(), kbf.clone() + sglang_jit_rope(q_mixed, k_mixed, cos_sin_cache, positions, is_neox) + + q_ref, k_f16 = q16.clone(), kbf.to(torch.float16) + sglang_jit_rope(q_ref, k_f16, cos_sin_cache, positions, is_neox) + q_bf16, k_ref = q16.to(torch.bfloat16), kbf.clone() + sglang_jit_rope(q_bf16, k_ref, cos_sin_cache, positions, is_neox) + + assert torch.equal(q_mixed, q_ref) + assert torch.equal(k_mixed, k_ref) + + +@pytest.mark.parametrize("is_neox", IS_NEOX_LIST) +def test_rope_store_mixed_q_dtype(is_neox: bool) -> None: + """Fused RoPE + KV store with fp16 q + bf16 k: q bitwise vs the all-fp16 + kernel; k, k_cache, v_cache bitwise vs the all-bf16 kernel.""" + from sglang.kernels.ops.attention.rope import apply_rope_inplace_with_kvcache + + batch_size, num_kv_heads, gqa_ratio, rope_dim = 129, 4, 8, 64 + num_qo_heads = num_kv_heads * gqa_ratio + row_size = num_kv_heads * rope_dim + q16 = torch.randn( + batch_size, num_qo_heads, rope_dim, device=DEVICE, dtype=torch.float16 + ) + kbf = torch.randn( + batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=torch.bfloat16 + ) + vbf = torch.randn( + batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=torch.bfloat16 + ) + positions = torch.randint( + 0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=torch.int64 + ) + out_loc = torch.randperm(CACHE_SIZE, device=DEVICE, dtype=torch.int64)[:batch_size] + cos_sin_cache = create_cos_sin_cache(rope_dim) + + def run(q, k, v): + k_cache = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=k.dtype) + v_cache = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=k.dtype) + apply_rope_inplace_with_kvcache( + q, + k, + v, + k_cache, + v_cache, + cos_sin_cache, + positions, + out_loc, + is_neox=is_neox, + ) + return k_cache, v_cache + + q_mixed = q16.clone() + k_mixed, v_mixed = kbf.clone(), vbf.clone() + kc_mixed, vc_mixed = run(q_mixed, k_mixed, v_mixed) + + q_ref = q16.clone() + kc16, _ = run(q_ref, kbf.to(torch.float16), vbf.to(torch.float16)) + q_bf16 = q16.to(torch.bfloat16) + kc_ref, vc_ref = run(q_bf16, kbf.clone(), vbf.clone()) + + assert torch.equal(q_mixed, q_ref) + assert torch.equal(kc_mixed, kc_ref) + assert torch.equal(vc_mixed, vc_ref) + + @pytest.mark.parametrize("batch_size", BS_LIST) @pytest.mark.parametrize("is_neox", IS_NEOX_LIST) @pytest.mark.parametrize("rope_dim", PARTIAL_ROPE_DIM_LIST)