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

Support Expert Deferral Mechanism in KTransformers (#12586)

Browse Source 
Co-authored-by: Chen Hongtao <56470055+chenht2022@users.noreply.github.com>
Co-authored-by: chenht2022 <cht22@mails.tsinghua.edu.cn>
This commit is contained in:
Atream
2025-11-06 05:41:52 +08:00
committed by GitHub
parent 3651cfbf62
commit 627bac649c
3 changed files with 50 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -270,6 +270,8 @@ class Envs:
SGLANG_KT_MOE_AMX_WEIGHT_PATH = EnvStr(None)
SGLANG_KT_AMX_METHOD = EnvStr(None)
SGLANG_KT_MOE_CHUNKED_PREFILL_SIZE = EnvInt(None)
SGLANG_KT_MOE_MAX_DEFERRED_EXPERTS_PER_TOKEN = EnvInt(None)
SGLANG_KT_MOE_TOTAL_LAYERS = EnvInt(None)
# Sparse Embeddings
SGLANG_EMBEDDINGS_SPARSE_HEAD = EnvStr(None)

25
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py
View File

@@ -705,8 +705,9 @@ class CompressedTensorsWNA16AMXMoEMethod(CompressedTensorsMoEMethod):
threadpool_count,
amx_weight_path,
chunked_prefill_size,
max_deferred_experts_per_token,
total_num_hidden_layers,
):
if not KTRANSFORMERS_AVAILABLE:
raise ImportError(
"kt_kernel is not installed, to use CompressedTensorsWNA16AMXEPMoEMethod, please install kt_kernel."
@@ -723,6 +724,8 @@ class CompressedTensorsWNA16AMXMoEMethod(CompressedTensorsMoEMethod):
self.cpuinfer = cpuinfer
self.threadpool_count = threadpool_count
self.amx_wrapper = None
self.max_deferred_experts_per_token = max_deferred_experts_per_token
self.total_num_hidden_layers = total_num_hidden_layers
def create_weights(
self,
@@ -733,6 +736,13 @@ class CompressedTensorsWNA16AMXMoEMethod(CompressedTensorsMoEMethod):
params_dtype: torch.dtype,
**extra_weight_attrs,
):
layer_max_deferred = self.max_deferred_experts_per_token or 0
if (
self.max_deferred_experts_per_token is not None
and self.total_num_hidden_layers is not None
and self.layer_idx == self.total_num_hidden_layers - 1
):
layer_max_deferred = 0
self.experts_num = num_experts
self.num_experts_per_tok = extra_weight_attrs.pop("top_k")
self.hidden_size = hidden_size
@@ -751,6 +761,7 @@ class CompressedTensorsWNA16AMXMoEMethod(CompressedTensorsMoEMethod):
threadpool_count=self.threadpool_count,
amx_weight_path=self.amx_weight_path,
chunked_prefill_size=self.chunked_prefill_size,
max_deferred_experts_per_token=layer_max_deferred,
amx_method=envs.SGLANG_KT_AMX_METHOD.value,
)
@@ -848,6 +859,8 @@ def override_config(
amx_weight_path,
amx_method,
chunked_prefill_size,
max_deferred_experts_per_token,
num_hidden_layers,
):
"""Override MOE configuration via environment variables."""
# Set environment variables using envs utility class
@@ -863,6 +876,12 @@ def override_config(
envs.SGLANG_KT_AMX_METHOD.set(amx_method)
if chunked_prefill_size is not None:
envs.SGLANG_KT_MOE_CHUNKED_PREFILL_SIZE.set(chunked_prefill_size)
envs.SGLANG_KT_MOE_MAX_DEFERRED_EXPERTS_PER_TOKEN.set(
max_deferred_experts_per_token
)
envs.SGLANG_KT_MOE_TOTAL_LAYERS.set(num_hidden_layers)
cls.max_deferred_experts_per_token = max_deferred_experts_per_token
cls.total_num_hidden_layers = num_hidden_layers
class CompressedTensorsWNA16AMXEPMoEMethod(CompressedTensorsMoEMethod):
@@ -887,6 +906,8 @@ class CompressedTensorsWNA16AMXEPMoEMethod(CompressedTensorsMoEMethod):
threadpool_count = envs.SGLANG_KT_THREADPOOL_COUNT.value
amx_weight_path = envs.SGLANG_KT_MOE_AMX_WEIGHT_PATH.value
chunked_prefill_size = envs.SGLANG_KT_MOE_CHUNKED_PREFILL_SIZE.value
max_deferred = envs.SGLANG_KT_MOE_MAX_DEFERRED_EXPERTS_PER_TOKEN.value
total_layers = envs.SGLANG_KT_MOE_TOTAL_LAYERS.value
self.AMX_method = CompressedTensorsWNA16AMXMoEMethod(
quant_config,
@@ -896,6 +917,8 @@ class CompressedTensorsWNA16AMXEPMoEMethod(CompressedTensorsMoEMethod):
threadpool_count,
amx_weight_path,
chunked_prefill_size,
max_deferred_experts_per_token=max_deferred,
total_num_hidden_layers=total_layers,
)
self.marlin_method = CompressedTensorsWNA16MoEMethod(
quant_config, self.num_gpu_experts

24
python/sglang/srt/server_args.py
View File

@@ -437,6 +437,7 @@ class ServerArgs:
kt_cpuinfer: Optional[int] = None
kt_threadpool_count: Optional[int] = None
kt_num_gpu_experts: Optional[int] = None
kt_max_deferred_experts_per_token: Optional[int] = None
# Double Sparsity
enable_double_sparsity: bool = False
@@ -1329,6 +1330,21 @@ class ServerArgs:
override_config,
)
num_hidden_layers = None
if self.kt_max_deferred_experts_per_token is not None:
try:
model_config = self.get_model_config()
base_config = (
getattr(model_config, "hf_text_config", None)
or model_config.hf_config
)
num_hidden_layers = getattr(base_config, "num_hidden_layers", None)
except Exception as exc: # noqa: BLE001
logger.warning(
"Failed to load model config for kt_max_deferred_experts_per_token: %s",
exc,
)
override_config(
CompressedTensorsWNA16AMXEPMoEMethod,
self.kt_num_gpu_experts,
@@ -1337,6 +1353,8 @@ class ServerArgs:
self.kt_amx_weight_path,
self.kt_amx_method,
self.chunked_prefill_size,
self.kt_max_deferred_experts_per_token,
num_hidden_layers,
)
def _handle_data_parallelism(self):
@@ -3038,6 +3056,12 @@ class ServerArgs:
type=int,
help="[ktransformers parameter] The number of GPU experts.",
)
parser.add_argument(
"--kt-max-deferred-experts-per-token",
type=int,
default=ServerArgs.kt_max_deferred_experts_per_token,
help="Maximum number of experts deferred to CPU per token. All MoE layers except the final one use this value; the final layer always uses 0.",
)
# Double Sparsity
parser.add_argument(