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

Add opt-in SGLANG_ROPE_CACHE_FP32 to keep RoPE cache in fp32 on non-CUDA (#29729)

Browse Source
This commit is contained in:
Chengze Fan
2026-07-09 02:10:39 -07:00
committed by GitHub
parent e703f9e566
commit 122b3266a2
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+1
View File
@@ -683,6 +683,7 @@ class Envs:
# RoPE cache configuration # RoPE cache configuration
SGLANG_SPEC_EXPANSION_SAFETY_FACTOR = EnvInt(2) SGLANG_SPEC_EXPANSION_SAFETY_FACTOR = EnvInt(2)
SGLANG_ROPE_CACHE_FP32 = EnvBool(False)
SGLANG_ROPE_CACHE_SAFETY_MARGIN = EnvInt(256) SGLANG_ROPE_CACHE_SAFETY_MARGIN = EnvInt(256)
SGLANG_ROPE_CACHE_ALIGN = EnvInt(128) SGLANG_ROPE_CACHE_ALIGN = EnvInt(128)
python/sglang/srt/layers/rotary_embedding/base.py
+3 -4
View File
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
import torch import torch
from sglang.srt.environ import envs
from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb
from sglang.srt.layers.utils import MultiPlatformOp from sglang.srt.layers.utils import MultiPlatformOp
from sglang.srt.platforms import current_platform from sglang.srt.platforms import current_platform
@@ -93,8 +94,8 @@ class RotaryEmbedding(MultiPlatformOp):
self.dtype = dtype self.dtype = dtype
cache = self._compute_cos_sin_cache() cache = self._compute_cos_sin_cache()
# NOTE(ByronHsu): cache needs to be in FP32 for numerical stability # NOTE(ByronHsu): cache needs to be in FP32 for numerical stability.
if not _is_cuda: if not (_is_cuda or envs.SGLANG_ROPE_CACHE_FP32.get()):
cache = cache.to(dtype) cache = cache.to(dtype)
if ( if (
@@ -178,8 +179,6 @@ class RotaryEmbedding(MultiPlatformOp):
def _ensure_cos_sin_cache_length(self, needed_max_pos: int): def _ensure_cos_sin_cache_length(self, needed_max_pos: int):
"""Ensure cos_sin_cache length > needed_max_pos.""" """Ensure cos_sin_cache length > needed_max_pos."""
from sglang.srt.environ import envs
cur_len = int(self.cos_sin_cache.shape[0]) cur_len = int(self.cos_sin_cache.shape[0])
if needed_max_pos < cur_len: if needed_max_pos < cur_len:
return return