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

[VLM] Tiny: Unify VLM environment variables (#15572)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2025-12-22 17:32:37 +08:00
committed by GitHub
parent 575a49dc0e
commit 828dec1c7c
4 changed files with 11 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/environ.py
View File

@@ -342,9 +342,10 @@ class Envs:
SGLANG_RESIZE_RESAMPLE = EnvStr("")
SGLANG_MM_BUFFER_SIZE_MB = EnvInt(0)
SGLANG_MM_PRECOMPUTE_HASH = EnvBool(False)
SGLANG_VIT_ENABLE_CUDA_GRAPH = EnvBool(False)
# VLM Item CUDA IPC Transport
SGLANG_USE_CUDA_IPC_TRANSPORT=EnvBool(False)
SGLANG_USE_CUDA_IPC_TRANSPORT = EnvBool(False)
SGLANG_MM_FEATURE_CACHE_MB = EnvInt(4 * 1024)
SGLANG_MM_ITEM_MEM_POOL_RECYCLE_INTERVAL_SEC = EnvFloat(0.05)

5
python/sglang/srt/layers/attention/vision.py
View File

@@ -11,6 +11,7 @@ import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size
from sglang.srt.utils import (
get_bool_env_var,
@@ -294,7 +295,7 @@ class VisionTritonAttention(nn.Module):
Returns:
[b * s, h, head_size]
"""
if get_bool_env_var("SGLANG_VIT_ENABLE_CUDA_GRAPH"):
if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
if "output_ws" not in kwargs:
raise RuntimeError("output_ws should be prepared for cuda-graph mode")
@@ -363,7 +364,7 @@ class VisionFlash3Attention(nn.Module):
Returns:
[b * s, h, head_size]
"""
if get_bool_env_var("SGLANG_VIT_ENABLE_CUDA_GRAPH"):
if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
max_seqlen = cu_seqlens[1]
output = flash_attn_varlen_func(
q,

5
python/sglang/srt/models/qwen2_5_vl.py
View File

@@ -46,6 +46,7 @@ from sglang.srt.distributed import (
get_tensor_model_parallel_world_size,
)
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.environ import envs
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import (
@@ -74,7 +75,7 @@ from sglang.srt.models.utils import RotaryPosMixin, WeightsMapper, permute_inv
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix, get_bool_env_var
from sglang.srt.utils import add_prefix
logger = logging.getLogger(__name__)
@@ -389,7 +390,7 @@ class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin):
x: torch.Tensor,
grid_thw: torch.Tensor,
) -> torch.Tensor:
if get_bool_env_var("SGLANG_VIT_ENABLE_CUDA_GRAPH"):
if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
return self.forward_with_cuda_graph(x, grid_thw)
# patchify

5
python/sglang/srt/models/qwen3_vl.py
View File

@@ -31,6 +31,7 @@ from sglang.srt.distributed import (
get_tensor_model_parallel_world_size,
)
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.environ import envs
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
from sglang.srt.layers.logits_processor import LogitsProcessor
@@ -59,7 +60,7 @@ from sglang.srt.models.utils import (
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix, get_bool_env_var, get_int_env_var
from sglang.srt.utils import add_prefix, get_int_env_var
from sglang.srt.utils.hf_transformers_utils import get_processor
logger = logging.getLogger(__name__)
@@ -466,7 +467,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
x: torch.Tensor,
grid_thw: torch.Tensor,
) -> torch.Tensor:
if get_bool_env_var("SGLANG_VIT_ENABLE_CUDA_GRAPH"):
if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
return self.forward_with_cuda_graph(x, grid_thw)
x = x.to(device=self.device, dtype=self.dtype)