[Misc]Register and refactor some environs for dpsk-fp4 and DeepEp (#14538)

This commit is contained in:
Baizhou Zhang
2025-12-06 12:29:16 -08:00
committed by GitHub
parent e592ee6545
commit 9dfa01a435
7 changed files with 36 additions and 17 deletions

View File

@@ -19,9 +19,10 @@ from typing import Optional
import torch
from sglang.srt.environ import envs
from sglang.srt.layers.moe import get_moe_runner_backend
from sglang.srt.layers.moe.utils import is_sbo_enabled
from sglang.srt.utils import get_int_env_var, is_blackwell
from sglang.srt.utils import is_blackwell
class SboFlags:
@@ -87,9 +88,11 @@ def compute_overlap_args(dispatch_output, alt_stream):
total_num_sms = torch.cuda.get_device_properties(
device="cuda"
).multi_processor_count
communicate_num_sms = get_int_env_var(
"SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS", 32 if is_blackwell() else 3
)
if envs.SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS.is_set():
communicate_num_sms = envs.SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS.get()
else:
communicate_num_sms = 32 if is_blackwell() else 3
compute_num_sms = total_num_sms - communicate_num_sms
assert alt_stream is not None