From 122b3266a24d849364b3369244c2cab14d5b3aa0 Mon Sep 17 00:00:00 2001 From: Chengze Fan Date: Thu, 9 Jul 2026 02:10:39 -0700 Subject: [PATCH] Add opt-in SGLANG_ROPE_CACHE_FP32 to keep RoPE cache in fp32 on non-CUDA (#29729) --- python/sglang/srt/environ.py | 1 + python/sglang/srt/layers/rotary_embedding/base.py | 7 +++---- 2 files changed, 4 insertions(+), 4 deletions(-) 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