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

Select DeepGEMM standard layouts by memory budget (#33474)

Browse Source 
Co-authored-by: Chunan Zeng <zcnrex@gmail.com>
This commit is contained in:
YAMY
2026-08-05 15:50:43 -07:00
committed by GitHub
co-authored by Chunan Zeng
parent 25035bff8d
commit 7bc90ab394
5 changed files with 200 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+2
View File
@@ -733,6 +733,8 @@ class Envs:
# DeepGemm
SGLANG_ENABLE_JIT_DEEPGEMM = EnvBool(True)
SGLANG_DEEPGEMM_STANDARD_LAYOUT = EnvStr("auto")
SGLANG_DEEPGEMM_MASKED_MEMORY_BUDGET_FRACTION = EnvFloat(0.25)
# Cap the DeepGEMM masked grouped-GEMM per-expert padded capacity at
# round_up(max(masked_m), 256) instead of round_up(rank_tokens, 256):
# shrinks the [num_local_experts, m, *] MoE intermediates ~4x under
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
+86 -5
View File
@@ -73,6 +73,7 @@ else:
_DEEPGEMM_ON_H20 = get_bool_env_var("SGLANG_DEEPGEMM_ON_H20")
_masked_standard_layout_memory_budget_bytes: Optional[int] = None
# TODO(kaixih@nvidia): ideally we should merge this logic into
@@ -100,11 +101,91 @@ def copy_list_to_gpu_no_ce(arr: List[int]):
return tensor_gpu
def _should_use_masked_standard_layout(runner_config: MoeRunnerConfig) -> bool:
"""Use masked GEMM when expert parallelism keeps its buffer small."""
def set_masked_standard_layout_memory_budget(
available_memory_bytes: int,
) -> int:
"""Cache the masked-layout share of free non-static device memory."""
global _masked_standard_layout_memory_budget_bytes
fraction = envs.SGLANG_DEEPGEMM_MASKED_MEMORY_BUDGET_FRACTION.get()
if not 0.0 < fraction <= 1.0:
raise ValueError(
"SGLANG_DEEPGEMM_MASKED_MEMORY_BUDGET_FRACTION must be in (0, 1]"
)
_masked_standard_layout_memory_budget_bytes = int(available_memory_bytes * fraction)
return _masked_standard_layout_memory_budget_bytes
def _estimate_masked_standard_layout_peak_bytes(
runner_config: MoeRunnerConfig,
quant_info: DeepGemmMoeQuantInfo,
hidden_states: torch.Tensor,
) -> int:
padded_m = (hidden_states.shape[0] // 256 + 1) * 256
activation_dtype = (
torch.bfloat16
if quant_info.w13_weight.dtype == torch.bfloat16
else torch.float8_e4m3fn
)
hidden_size = hidden_states.shape[1]
gateup_size = quant_info.w13_weight.shape[1]
gateup_row_bytes = gateup_size * torch.bfloat16.itemsize
down_output_row_bytes = quant_info.w2_weight.shape[1] * torch.bfloat16.itemsize
input_row_bytes = hidden_size * activation_dtype.itemsize
down_input_row_bytes = gateup_size // 2 * activation_dtype.itemsize
if activation_dtype == torch.bfloat16:
input_scale_row_bytes = 0
down_scale_row_bytes = 0
else:
block_k = quant_info.block_shape[1] if quant_info.block_shape else 128
packed_scales = quant_info.use_mxfp8 or deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
scale_item_bytes = (
torch.uint8.itemsize if packed_scales else torch.float32.itemsize
)
input_scale_row_bytes = ceil_div(hidden_size, block_k) * scale_item_bytes
down_scale_row_bytes = ceil_div(gateup_size // 2, block_k) * scale_item_bytes
peak_row_bytes = max(
input_row_bytes + input_scale_row_bytes + gateup_row_bytes,
gateup_row_bytes + down_input_row_bytes + down_scale_row_bytes,
down_input_row_bytes + down_scale_row_bytes + down_output_row_bytes,
)
return runner_config.num_local_experts * padded_m * peak_row_bytes
def _should_use_masked_standard_layout(
runner_config: MoeRunnerConfig,
quant_info: DeepGemmMoeQuantInfo,
hidden_states: torch.Tensor,
) -> bool:
mode = envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.get().lower()
if mode not in ("auto", "masked", "compact"):
raise ValueError(
"SGLANG_DEEPGEMM_STANDARD_LAYOUT must be one of: auto, masked, compact"
)
if mode != "auto":
return mode == "masked"
global _masked_standard_layout_memory_budget_bytes
if _masked_standard_layout_memory_budget_bytes is None:
# Serving sets an all-rank budget before capture. Direct eager callers
# fall back to this rank's free memory without querying inside capture.
# Import lazily to avoid a module-initialization cycle through
# runner_utils -> DeepEP -> MoE -> this module.
from sglang.srt.model_executor.runner_utils.capture_mode import (
get_is_capture_mode,
)
if get_is_capture_mode():
return False
free_memory, _ = torch.cuda.mem_get_info(hidden_states.device)
set_masked_standard_layout_memory_budget(free_memory)
return (
runner_config.num_experts > runner_config.num_local_experts
and runner_config.num_local_experts <= 32
_estimate_masked_standard_layout_peak_bytes(
runner_config, quant_info, hidden_states
)
<= _masked_standard_layout_memory_budget_bytes
)
@@ -775,7 +856,7 @@ def pre_permute_standard_to_deep_gemm(
topk_weights, topk_ids = topk_weights, topk_ids
if _should_use_masked_standard_layout(runner_config):
if _should_use_masked_standard_layout(runner_config, quant_info, hidden_states):
output_dtype = (
torch.bfloat16
if quant_info.w13_weight.dtype == torch.bfloat16
python/sglang/srt/layers/quantization/fp8.py
+16 -8
View File
@@ -1073,20 +1073,28 @@ class Fp8MoEMethod(FusedMoEMethodBase):
), "cutlass_fp8 MoE requires SM90, SM100, or SM120 GPUs"
@staticmethod
def is_deepgemm_moe_runner_backend_enabled() -> bool:
def is_deepgemm_moe_runner_backend_enabled(
moe_runner_backend=None, moe_a2a_backend=None
) -> bool:
"""Check if MoE will actually use DeepGEMM runner for FP8."""
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
moe_runner_backend = get_moe_runner_backend()
if moe_runner_backend is None:
moe_runner_backend = get_moe_runner_backend()
if moe_runner_backend.is_deep_gemm():
return True
if moe_runner_backend.is_auto():
return deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_nixl()
)
if moe_a2a_backend is None:
moe_a2a_backend = get_moe_a2a_backend()
if not (
moe_a2a_backend.is_deepep()
or moe_a2a_backend.is_mooncake()
or moe_a2a_backend.is_nixl()
):
return False
from sglang.srt.layers import deep_gemm_wrapper
return deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
return False
@staticmethod
python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py
+53
View File
@@ -8,9 +8,11 @@ from typing import TYPE_CHECKING, Optional
import msgspec
from sglang.srt.configs.model_config import ModelImpl
from sglang.srt.distributed import get_world_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
prealloc_symmetric_memory_pool,
)
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner
from sglang.srt.hardware_backend.xpu.graph_runner.xpu_graph_runner import XPUGraphRunner
from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner
@@ -125,6 +127,57 @@ def capture_cuda_graphs(
# batch.
eager_runner = EagerRunner(model_runner)
if model_runner.is_draft_worker:
moe_runner_backend = (
model_runner.server_args.speculative_moe_runner_backend
or model_runner.server_args.moe_runner_backend
)
moe_a2a_backend = (
model_runner.server_args.speculative_moe_a2a_backend
or model_runner.server_args.moe_a2a_backend
)
else:
moe_runner_backend = model_runner.server_args.moe_runner_backend
moe_a2a_backend = model_runner.server_args.moe_a2a_backend
uses_deep_gemm_moe_runner = moe_runner_backend == "deep_gemm"
if moe_runner_backend == "auto" and model_runner.model_config.quantization in (
"fp8",
"mxfp8",
):
from sglang.srt.layers.moe.utils import MoeA2ABackend, MoeRunnerBackend
from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod
uses_deep_gemm_moe_runner = Fp8MoEMethod.is_deepgemm_moe_runner_backend_enabled(
MoeRunnerBackend(moe_runner_backend),
MoeA2ABackend(moe_a2a_backend),
)
if (
model_runner.device == "cuda"
and envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.get().lower() == "auto"
and uses_deep_gemm_moe_runner
):
from sglang.srt.layers.moe.moe_runner.deep_gemm import (
set_masked_standard_layout_memory_budget,
)
world_group = get_world_group()
available_memory_gb = get_available_gpu_memory(
model_runner.device,
model_runner.gpu_id,
distributed=world_group.world_size > 1,
cpu_group=world_group.cpu_group,
)
budget_bytes = set_masked_standard_layout_memory_budget(
int(available_memory_gb * (1 << 30))
)
logger.info(
"DeepGEMM masked layout budget: %.2f GiB from %.2f GiB free.",
budget_bytes / (1 << 30),
available_memory_gb,
)
# cuda-graph capture: prefill before decode, so both coalesce onto the
# eager buffer allocated above. (capture_prefill_graph routes prefill
# to the eager runner when the prefill graph is disabled.)
test/registered/kernels/ops/moe/test_minimax_quant_scatter.py
+43 -10
View File
@@ -14,6 +14,7 @@ from sglang.kernels.ops.quantization.minimax_quant_ue8m0 import (
per_token_quant_fp8_ue8m0,
per_token_quant_fp8_ue8m0_scatter,
)
from sglang.srt.environ import envs
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.deep_gemm import (
DeepGemmMoeQuantInfo,
@@ -158,6 +159,37 @@ def test_compact_all_tokens_uses_tight_routing_independent_bound(
)
def test_standard_layout_auto_memory_policy(monkeypatch):
config = MoeRunnerConfig(
num_experts=512,
num_local_experts=512,
hidden_size=4096,
intermediate_size_per_partition=256,
top_k=8,
)
quant_info = DeepGemmMoeQuantInfo(
w13_weight=torch.empty((1, 512, 1), dtype=torch.float8_e4m3fn),
w2_weight=torch.empty((1, 4096, 1), dtype=torch.float8_e4m3fn),
use_fp8=True,
block_shape=[128, 128],
)
monkeypatch.setattr(
deep_gemm_runner,
"_masked_standard_layout_memory_budget_bytes",
int(42.5 * (1 << 30)),
)
with envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.override("auto"):
for num_tokens, expected in ((8192, True), (16384, False)):
hidden_states = torch.empty((num_tokens, 4096), device="meta")
assert (
deep_gemm_runner._should_use_masked_standard_layout(
config, quant_info, hidden_states
)
is expected
)
@pytest.mark.parametrize("weight_dtype", ["fp8", "bf16"])
def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_dtype):
"""Exercise both production grouped GEMMs through the standard path."""
@@ -252,9 +284,9 @@ def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_d
topk_output=(topk_weights, topk_ids, None),
)
def run_with_num_experts(num_experts):
def run_with_layout(layout):
config = MoeRunnerConfig(
num_experts=num_experts,
num_experts=8,
num_local_experts=num_local_experts,
hidden_size=hidden,
intermediate_size_per_partition=intermediate,
@@ -264,12 +296,13 @@ def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_d
inplace=False,
)
running_state = {}
runner_input = pre_permute_standard_to_deep_gemm(
dispatch_output,
quant_info,
config,
running_state,
)
with envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.override(layout):
runner_input = pre_permute_standard_to_deep_gemm(
dispatch_output,
quant_info,
config,
running_state,
)
runner_output = DeepGemmRunnerCore(config).run(
runner_input,
quant_info,
@@ -288,10 +321,10 @@ def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_d
)
compact_is_masked, compact_all_tokens, compact_m_indices, compact_output = (
run_with_num_experts(num_local_experts)
run_with_layout("compact")
)
masked_is_masked, masked_all_tokens, masked_m_indices, masked_output = (
run_with_num_experts(8)
run_with_layout("masked")
)
torch.cuda.synchronize()