diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 589447e07..3a724e7d6 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -683,6 +683,7 @@ class Envs: # RoPE cache configuration SGLANG_SPEC_EXPANSION_SAFETY_FACTOR = EnvInt(2) + SGLANG_ROPE_CACHE_FP32 = EnvBool(False) SGLANG_ROPE_CACHE_SAFETY_MARGIN = EnvInt(256) SGLANG_ROPE_CACHE_ALIGN = EnvInt(128) diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index 39cb8f7e2..e04cd876f 100644 --- a/python/sglang/srt/layers/rotary_embedding/base.py +++ b/python/sglang/srt/layers/rotary_embedding/base.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union import torch +from sglang.srt.environ import envs from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb from sglang.srt.layers.utils import MultiPlatformOp from sglang.srt.platforms import current_platform @@ -93,8 +94,8 @@ class RotaryEmbedding(MultiPlatformOp): self.dtype = dtype cache = self._compute_cos_sin_cache() - # NOTE(ByronHsu): cache needs to be in FP32 for numerical stability - if not _is_cuda: + # NOTE(ByronHsu): cache needs to be in FP32 for numerical stability. + if not (_is_cuda or envs.SGLANG_ROPE_CACHE_FP32.get()): cache = cache.to(dtype) if ( @@ -178,8 +179,6 @@ class RotaryEmbedding(MultiPlatformOp): def _ensure_cos_sin_cache_length(self, needed_max_pos: int): """Ensure cos_sin_cache length > needed_max_pos.""" - from sglang.srt.environ import envs - cur_len = int(self.cos_sin_cache.shape[0]) if needed_max_pos < cur_len: return