diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 89854f01c..d8419c48e 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -82,6 +82,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_MORI_FP8_COMB` | Use FP8 for combine | `"false"` | | `SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | Maximum number of dispatch tokens per rank for MORI-EP buffer allocation | `4096` | | `SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD` | Threshold for switching between `InterNodeV1` and `InterNodeV1LL` kernel types. `InterNodeV1LL` is used if `SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK` is less than or equal to this threshold; otherwise, `InterNodeV1` is used. | `256` | +| `SGLANG_MORI_PREALLOC_MAX_RECV_TOKENS` | This argument devives `SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK` which indicates customized amount of tokens preallocated for a rank, valid range from 1 to world_size*SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK, by default `0` means maximum. Setting a smaller value will reduce memory footprint but too small value could cause buffer overflow. | `0` | | `SGLANG_MORI_QP_PER_TRANSFER` | Number of RDMA Queue Pairs (QPs) used per transfer operation | `1` | | `SGLANG_MORI_POST_BATCH_SIZE` | Number of RDMA work requests posted in a single batch to each QP | `-1` | | `SGLANG_MORI_NUM_WORKERS` | Number of worker threads in the RDMA executor thread pool | `1` | diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index 3113c3cbe..18d1b6c8a 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py @@ -19,7 +19,11 @@ from sglang.srt.layers.moe.utils import ( DeepEPMode, is_tbo_enabled, ) -from sglang.srt.utils import get_bool_env_var, get_int_env_var, is_hip +from sglang.srt.utils import ( + get_bool_env_var, + get_int_env_var, + is_hip, +) if TYPE_CHECKING: from sglang.srt.single_batch_overlap import CombineOverlapArgs @@ -261,10 +265,23 @@ def init_mori_op( f"{combine_quant_type=}" ) - mori_config = mori.ops.EpDispatchCombineConfig( + def check_mori_compatibility(kwargs: dict) -> None: + """Remove kwargs not accepted by the installed mori's EpDispatchCombineConfig.""" + import dataclasses + + config_cls = mori.ops.EpDispatchCombineConfig + valid_kwargs = {f.name for f in dataclasses.fields(config_cls)} + + invalid_kwargs = set(kwargs.keys()) - valid_kwargs + for arg in invalid_kwargs: + logger.warning(f"[MORI compat] Removing incompatible argument {arg} ") + del kwargs[arg] + + # Definition refer to https://github.com/ROCm/mori/blob/f9be5ee2e5ac87256b9523399ae9d4d0e8a54f53/python/mori/ops/dispatch_combine.py#L66-L121 + common_kwargs = dict( + data_type=data_type, rank=rank, world_size=world_size, - data_type=data_type, hidden_dim=hidden_dim, scale_dim=scale_dim, scale_type_size=scale_type_size, @@ -274,12 +291,19 @@ def init_mori_op( num_experts_per_token=router_topk, warp_num_per_block=warp_num_per_block, block_num=block_num, + max_total_recv_tokens=get_int_env_var( + "SGLANG_MORI_PREALLOC_MAX_RECV_TOKENS", 0 + ), kernel_type=kernel_type, gpu_per_node=gpu_per_node, rdma_block_num=rdma_block_num, - num_qp_per_pe=2, + num_qp_per_pe=2, # Number of queue pairs per processing element quant_type=combine_quant_type, ) + + check_mori_compatibility(common_kwargs) + + mori_config = mori.ops.EpDispatchCombineConfig(**common_kwargs) mori_op = mori.ops.EpDispatchCombineOp(mori_config) return mori_op diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index ae4feb3cf..34835518a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -5188,7 +5188,7 @@ class ServerArgs: type=str, choices=["normal", "low_latency", "auto"], default="auto", - help="Select the mode when enable DeepEP MoE, could be `normal`, `low_latency` or `auto`. Default is `auto`, which means `low_latency` for decode batch and `normal` for prefill batch.", + help="Select the mode when enable DeepEP or MoriEP MoE, could be `normal`, `low_latency` or `auto`. Default is `auto`, which means `low_latency` for decode batch and `normal` for prefill batch.", ) parser.add_argument( "--ep-num-redundant-experts",