chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add FP32 dtype support for RoPE - Part1 (#13181)

Browse Source
This commit is contained in:
iLeGend
2025-11-16 03:37:18 +08:00
committed by GitHub
parent 0d116b9a0b
commit 20e59f9510
2 changed files with 13 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
sgl-kernel/csrc/elementwise/rope.cu
View File

@@ -14,8 +14,13 @@
* limitations under the License.
*/
#include <ATen/cuda/Exceptions.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/all.h>
#include "pos_enc.cuh"
#include "pytorch_extension_utils.h"
#include "utils.h"
using namespace flashinfer;
@@ -88,7 +93,7 @@ void apply_rope_pos_ids_cos_sin_cache(
size_t k_rope_stride_h = k_rope.stride(1);
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(q.scalar_type(), c_type, [&] {
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(q.scalar_type(), c_type, [&] {
// TODO temporarily only use `BatchQKApplyRotaryPosIdsCosSinCacheEnhanced` when save_kv_cache
// to avoid changing original code path; but this branch is feature-complete and should switch to this later
if (save_kv_cache) {

6
sgl-kernel/tests/test_rotary_embedding.py
View File

@@ -47,6 +47,12 @@ from sgl_kernel.testing.rotary_embedding import (
(128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 32, 8, False),
(128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 16, 4, False),
(512, 128, 311, 10000, False, torch.bfloat16, "cuda", 3, 39, 4, 2, False),
(64, 64, 32, 8000, True, torch.float32, "cuda", 32, 32, 1, 1, False),
(256, 128, 4096, 10000, True, torch.float32, "cuda", 2, 512, 4, 2, False),
(512, 128, 311, 10000, True, torch.float32, "cuda", 3, 39, 4, 2, False),
(128, 128, 2048, 10000, False, torch.float32, "cuda", 2, 512, 32, 8, False),
(128, 128, 2048, 10000, False, torch.float32, "cuda", 2, 512, 16, 4, False),
(512, 128, 311, 10000, False, torch.float32, "cuda", 3, 39, 4, 2, False),
],
)
def test_correctness(