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

Add FP32 dtype support for RoPE - Part2 (#13328)

Browse Source
This commit is contained in:
iLeGend
2025-11-20 13:19:53 +08:00
committed by GitHub
parent bc42c8c415
commit 10e0b83a4c
3 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/layers/rotary_embedding.py
View File

@@ -113,7 +113,7 @@ class RotaryEmbedding(CustomOp):
if not _is_cuda:
cache = cache.to(dtype)
if dtype == torch.float32 or (
if (
(not (_is_cuda or _is_npu) or self.head_size not in [64, 128, 256, 512])
and not (_is_cpu and _is_cpu_amx_available)
and not (_is_xpu)
@@ -273,11 +273,7 @@ class RotaryEmbedding(CustomOp):
offsets: Optional[torch.Tensor] = None,
fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
if (
_is_cuda
and (self.head_size in [64, 128, 256, 512])
and self.dtype != torch.float32
):
if _is_cuda and (self.head_size in [64, 128, 256, 512]):
apply_rope_with_cos_sin_cache_inplace(
positions=positions,
query=query,