diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 72081aa1b..99fa1f44b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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 diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index b769cce03..67bdb5ea8 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 86859e526..f94ce61ae 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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 diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index cd30caaea..2bd30e03d 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py @@ -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.) diff --git a/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py b/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py index 3b9821835..99ecf1a0b 100644 --- a/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py +++ b/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py @@ -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()