diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 41ff73fc2..f3f8e644d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -266,16 +266,6 @@ class Envs: # Release & Resume Memory SGLANG_MEMORY_SAVER_CUDA_GRAPH = EnvBool(False) - # Ktransformers - SGLANG_KT_MOE_NUM_GPU_EXPERTS = EnvInt(None) - SGLANG_KT_MOE_CPUINFER = EnvInt(None) - SGLANG_KT_THREADPOOL_COUNT = EnvInt(None) - 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) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index bc0e1061e..543bd877b 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -25,6 +25,10 @@ from sglang.srt.layers.moe import ( get_moe_a2a_backend, get_moe_runner_backend, ) +from sglang.srt.layers.moe.kt_ep_wrapper import ( + KTEPWrapperMethod, + create_kt_config_from_server_args, +) from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher from sglang.srt.layers.moe.token_dispatcher.standard import ( @@ -36,15 +40,11 @@ from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, QuantizationConfig, ) -from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors_moe import ( - CompressedTensorsWNA16AMXEPMoEMethod, - CompressedTensorsWNA16AMXMoEMethod, - CompressedTensorsWNA16MoEMethod, -) from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod from sglang.srt.layers.quantization.modelopt_quant import ModelOptNvFp4FusedMoEMethod from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod from sglang.srt.model_loader.weight_utils import narrow_padded_param_and_loaded_weight +from sglang.srt.server_args import get_global_server_args from sglang.srt.two_batch_overlap import MaybeTboDeepEPDispatcher from sglang.srt.utils import ( cpu_has_amx_support, @@ -206,10 +206,19 @@ class FusedMoE(torch.nn.Module): ) self.quant_method: Optional[FusedMoEMethodBase] = None - if quant_config is not None: - self.quant_method = quant_config.get_quant_method(self, prefix) - if self.quant_method is None: - self.quant_method = UnquantizedFusedMoEMethod(self.use_triton_kernels) + server_args = get_global_server_args() + kt_config = create_kt_config_from_server_args(server_args, layer_id) + if kt_config is not None: + if quant_config is not None: + gpu_method = quant_config.get_quant_method(self, prefix) + else: + gpu_method = UnquantizedFusedMoEMethod(self.use_triton_kernels) + self.quant_method = KTEPWrapperMethod(gpu_method, kt_config) + else: + if quant_config is not None: + self.quant_method = quant_config.get_quant_method(self, prefix) + if self.quant_method is None: + self.quant_method = UnquantizedFusedMoEMethod(self.use_triton_kernels) self.quant_method.create_weights( layer=self, @@ -222,8 +231,6 @@ class FusedMoE(torch.nn.Module): if not use_weight_loader_fused else self.weight_loader_fused ), - intermediate_size_full=intermediate_size, - top_k=top_k, with_bias=with_bias, ) @@ -541,11 +548,7 @@ class FusedMoE(torch.nn.Module): if isinstance( self.quant_method, - ( - CompressedTensorsWNA16MoEMethod, - CompressedTensorsWNA16AMXMoEMethod, - CompressedTensorsWNA16AMXEPMoEMethod, - ), + KTEPWrapperMethod, ): if self.quant_method.num_gpu_experts != -1: if expert_id >= self.quant_method.num_gpu_experts: @@ -573,15 +576,17 @@ class FusedMoE(torch.nn.Module): # compressed-tensors checkpoints with packed weights are stored flipped # TODO (mgoin): check self.quant_method.quant_config.quant_format # against known CompressionFormat enum values that have this quality + method = self.quant_method + if method.__class__.__name__ == "KTEPWrapperMethod": + method = method.gpu_method + loaded_weight = ( loaded_weight.t().contiguous() if ( - self.quant_method.__class__.__name__ + method.__class__.__name__ in [ "CompressedTensorsWNA16MarlinMoEMethod", "CompressedTensorsWNA16MoEMethod", - "CompressedTensorsWNA16AMXMoEMethod", - "CompressedTensorsWNA16AMXEPMoEMethod", ] ) else loaded_weight diff --git a/python/sglang/srt/layers/moe/kt_ep_wrapper.py b/python/sglang/srt/layers/moe/kt_ep_wrapper.py new file mode 100644 index 000000000..3d8901e8f --- /dev/null +++ b/python/sglang/srt/layers/moe/kt_ep_wrapper.py @@ -0,0 +1,393 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +KT Expert Parallelism Wrapper for MoE layers. + +This module provides a generic wrapper that enables CPU-GPU expert parallelism +for any MoE quantization method. It coordinates parallel execution of GPU experts +(using any quantization method) and CPU experts (using AMX/AVX instructions). +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.distributed import get_tensor_model_parallel_rank +from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase +from sglang.srt.utils import get_compiler_backend + +if TYPE_CHECKING: + from sglang.srt.layers.moe import MoeRunnerConfig + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + from sglang.srt.server_args import ServerArgs + +try: + from kt_kernel import KTMoEWrapper + + KTRANSFORMERS_AVAILABLE = True +except ImportError: + KTRANSFORMERS_AVAILABLE = False + + +@dataclass +class KTConfig: + """Configuration for KTransformers heterogeneous computing CPU part. + + Args: + layer_idx: Layer index in the model + num_gpu_experts: Number of experts to run on GPU + cpuinfer_threads: Number of CPU inference threads + threadpool_count: Number of thread pools for CPU computation + weight_path: Path to CPU quantized weights + chunked_prefill_size: Chunk size for prefill computation + method: CPU computation method (e.g., "int4") + num_layers: Total number of layers in the model (optional) + """ + + layer_idx: int + num_gpu_experts: int + cpuinfer_threads: int + threadpool_count: int + weight_path: str + chunked_prefill_size: int + max_deferred_experts_per_token: int + method: str + num_layers: Optional[int] = None + + +def create_kt_config_from_server_args( + server_args: "ServerArgs", layer_idx: int +) -> Optional[KTConfig]: + """Create KTConfig from ServerArgs if KT is configured. + + Args: + server_args: Global server arguments + layer_idx: Layer index in the model + + Returns: + KTConfig if KT is configured, None otherwise + """ + if server_args.kt_weight_path is None: + return None + + # Try to get num_layers from model config + num_layers = None + try: + hf_config = server_args.get_hf_config() + num_layers = getattr(hf_config, "num_hidden_layers", None) + except Exception: + # If we can't get the config, num_layers will be None + pass + + return KTConfig( + layer_idx=layer_idx, + num_gpu_experts=server_args.kt_num_gpu_experts, + cpuinfer_threads=server_args.kt_cpuinfer, + threadpool_count=server_args.kt_threadpool_count, + weight_path=server_args.kt_weight_path, + chunked_prefill_size=server_args.chunked_prefill_size, + method=server_args.kt_method, + max_deferred_experts_per_token=server_args.kt_max_deferred_experts_per_token, + num_layers=num_layers, + ) + + +@torch.compile(dynamic=True, backend=get_compiler_backend()) +def mask_cpu_expert_ids(topk_ids: torch.Tensor, num_gpu_experts: int) -> torch.Tensor: + """Mask CPU expert IDs by setting them to -1. + + This function masks expert IDs that should be computed on CPU (IDs >= num_gpu_experts) + so they won't be computed on GPU. The masked IDs are set to -1, which causes the + GPU MoE kernel to skip those experts. + + Args: + topk_ids: Tensor of shape [num_tokens, top_k] containing expert IDs + num_gpu_experts: Number of experts that should run on GPU (experts 0 to num_gpu_experts-1) + + Returns: + Modified topk_ids tensor with CPU expert IDs masked as -1 + """ + topk_ids[topk_ids >= num_gpu_experts] = -1 + return topk_ids + + +class KTEPWrapperMethod(FusedMoEMethodBase): + """Wrapper for any MoE quantization method to enable CPU-GPU expert parallelism. + + This wrapper coordinates parallel execution of: + - GPU experts (0 to num_gpu_experts-1) using any quantization method + - CPU experts (num_gpu_experts to total_experts-1) using AMX/AVX instructions + + The wrapper implements the submit-compute-sync pattern: + 1. Submit CPU expert computation (non-blocking) + 2. Execute GPU expert computation in parallel + 3. Synchronize and merge CPU+GPU results + + Example: + # Wrap any GPU method with AMX/AVX CPU expert support + gpu_method = CompressedTensorsWNA16MoEMethod(quant_config, prefix) + kt_config = KTConfig(layer_idx=0, num_gpu_experts=4, ...) + method = KTEPWrapperMethod(gpu_method, kt_config) + """ + + def __init__( + self, + gpu_method: FusedMoEMethodBase, + kt_config: KTConfig, + ): + """Initialize the KT EP wrapper. + + Args: + gpu_method: The quantization method to use for GPU experts + kt_config: Configuration for KT CPU expert computation + """ + if not KTRANSFORMERS_AVAILABLE: + raise ImportError( + "kt_kernel is not installed. To use KTransformers EP wrapper, please install kt_kernel." + ) + + self.gpu_method = gpu_method + self.kt_config = kt_config + self.num_gpu_experts = kt_config.num_gpu_experts + self.override_num_local_experts = True + self.gpu_method.num_gpu_experts = self.num_gpu_experts + self.tp_rank = get_tensor_model_parallel_rank() + + # KT wrapper will be initialized in create_weights + self.wrapper: Optional[KTMoEWrapper] = None + + # Store parameters needed for KT initialization + self._layer_params = None + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + """Create weights for both GPU and CPU experts. + + Args: + layer: The MoE layer module + num_experts: Total number of experts (GPU + CPU) + hidden_size: Hidden dimension size + intermediate_size_per_partition: Intermediate size per TP partition + params_dtype: Data type for parameters + **extra_weight_attrs: Additional weight attributes + """ + self.global_num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size_per_partition = intermediate_size_per_partition + + # Get required parameters from layer object + # top_k: number of experts selected per token + num_experts_per_tok = layer.top_k + + # intermediate_size_full: full intermediate size before TP partitioning + intermediate_size_full = ( + layer.intermediate_size_per_partition * layer.moe_tp_size + ) + + layer_max_deferred = self.kt_config.max_deferred_experts_per_token or 0 + if ( + self.kt_config.max_deferred_experts_per_token is not None + and self.kt_config.num_layers is not None + and self.kt_config.layer_idx == self.kt_config.num_layers - 1 + ): + layer_max_deferred = 0 + + # 1. Create weights for GPU experts using the wrapped method + # GPU experts: 0 to num_gpu_experts-1 + self.gpu_method.create_weights( + layer=layer, + num_experts=self.num_gpu_experts, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size_per_partition, + params_dtype=params_dtype, + **extra_weight_attrs, + ) + + # 2. Initialize KT wrapper for CPU experts + # CPU experts: num_gpu_experts to num_experts-1 + if self.tp_rank == 0: + self.wrapper = KTMoEWrapper( + layer_idx=self.kt_config.layer_idx, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + hidden_size=hidden_size, + moe_intermediate_size=intermediate_size_full, + num_gpu_experts=self.num_gpu_experts, + cpuinfer_threads=self.kt_config.cpuinfer_threads, + threadpool_count=self.kt_config.threadpool_count, + weight_path=self.kt_config.weight_path, + chunked_prefill_size=self.kt_config.chunked_prefill_size, + method=self.kt_config.method, + max_deferred_experts_per_token=layer_max_deferred, + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Process weights after loading from checkpoint. + + Args: + layer: The MoE layer module + """ + # 1. Process GPU weights + if hasattr(self.gpu_method, "process_weights_after_loading"): + self.gpu_method.process_weights_after_loading(layer) + + # 2. Load CPU weights using KT wrapper + if self.tp_rank == 0 and self.wrapper is not None: + torch.cuda.synchronize() + + # Get expert location metadata for CPU expert mapping + from sglang.srt.eplb.expert_location_dispatch import ( + get_global_expert_location_metadata, + ) + + physical_to_logical_map_cpu = ( + get_global_expert_location_metadata() + .physical_to_logical_map_cpu[self.kt_config.layer_idx] + .contiguous() + ) + self.wrapper.load_weights(physical_to_logical_map_cpu) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig" + ): + """Create MoE runner for computation. + + Args: + layer: The MoE layer module + moe_runner_config: Configuration for MoE runner + """ + self.moe_runner_config = moe_runner_config + if self.override_num_local_experts: + moe_runner_config.num_local_experts = self.num_gpu_experts + # Delegate to GPU method to create its runner + self.gpu_method.create_moe_runner(layer, moe_runner_config) + + def submit( + self, + layer: torch.nn.Module, + dispatch_output: "StandardDispatchOutput", + ) -> None: + """Submit CPU expert computation asynchronously (non-blocking). + + This method submits the CPU expert computation to AMX/AVX without waiting + for completion, allowing GPU computation to proceed in parallel. + + Args: + layer: The MoE layer module + dispatch_output: Dispatched tokens and routing information + """ + assert ( + self.moe_runner_config.activation == "silu" + ), "Only SiLU activation is supported." + + if self.tp_rank != 0 or self.wrapper is None: + return + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + topk_weights, topk_ids, _ = topk_output + + # Submit forward task to CPU (non-blocking) + self.wrapper.submit_forward( + x, topk_ids, topk_weights, torch.cuda.current_stream(x.device).cuda_stream + ) + + def sync(self, x: torch.Tensor) -> torch.Tensor: + """Synchronize and retrieve CPU expert computation results. + + This method waits for the CPU computation to complete and returns the results. + + Args: + x: Reference tensor for shape and device information + + Returns: + CPU expert computation results + """ + if self.tp_rank != 0 or self.wrapper is None: + return torch.zeros_like(x) + + # Wait for CPU computation and retrieve results + return self.wrapper.sync_forward( + x, torch.cuda.current_stream(x.device).cuda_stream + ) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: "StandardDispatchOutput", + ) -> "CombineInput": + """Execute hybrid CPU+GPU MoE forward pass with parallelism. + + This is the main computation method that coordinates: + 1. Submit CPU expert computation (non-blocking) + 2. Execute GPU expert computation in parallel + 3. Synchronize CPU results and merge with GPU results + + Args: + layer: The MoE layer module + dispatch_output: Dispatched tokens and routing information + + Returns: + Combined computation results from CPU and GPU experts + """ + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + # Step 1: Submit CPU expert computation (non-blocking) + if self.tp_rank == 0: + self.submit(layer, dispatch_output) + + # Step 2: Prepare GPU computation by masking CPU expert IDs + # CPU expert IDs (>= num_gpu_experts) are set to -1 so GPU kernel skips them + topk_ids = topk_output.topk_ids + masked_topk_ids = mask_cpu_expert_ids(topk_ids, self.num_gpu_experts) + + # Create modified dispatch output for GPU computation + masked_topk_output = topk_output._replace(topk_ids=masked_topk_ids) + masked_dispatch_output = dispatch_output._replace( + topk_output=masked_topk_output + ) + + # Step 3: Execute GPU expert computation (any quantization method) + # This runs in parallel with CPU computation + gpu_combine_input = self.gpu_method.apply(layer, masked_dispatch_output) + + # Step 4: Synchronize CPU results and merge with GPU results + output = gpu_combine_input.hidden_states + if self.tp_rank == 0: + cpu_output = self.sync(x) + output = output + cpu_output + + return StandardCombineInput(hidden_states=output) + + def __getattr__(self, name: str): + """Delegate attribute access to the wrapped GPU method. + + This allows the wrapper to transparently expose attributes and methods + from the wrapped GPU quantization method. + + Args: + name: Attribute name + + Returns: + Attribute value from gpu_method + """ + # Avoid infinite recursion for internal attributes + if name in ("gpu_method", "wrapper", "kt_config"): + raise AttributeError( + f"'{type(self).__name__}' object has no attribute '{name}'" + ) + + return getattr(self.gpu_method, name) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index 3960e568a..9e131f1bb 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -19,7 +19,6 @@ from compressed_tensors.quantization import ( ) from pydantic import BaseModel -from sglang.srt.environ import envs from sglang.srt.layers.quantization.base_config import ( LinearMethodBase, QuantizationConfig, @@ -71,8 +70,6 @@ class DeviceCapability(NamedTuple): class CompressedTensorsConfig(QuantizationConfig): - DeepSeekFP8Config = None - def __init__( self, target_scheme_map: Dict[str, Any], @@ -83,6 +80,7 @@ class CompressedTensorsConfig(QuantizationConfig): kv_cache_scheme: Optional[Dict[str, Any]] = None, config: Optional[Dict[str, Any]] = None, packed_modules_mapping: Optional[Dict[str, List[str]]] = None, + linear_fp8_config: Optional[Any] = None, ): super().__init__() self.ignore = ignore @@ -94,6 +92,8 @@ class CompressedTensorsConfig(QuantizationConfig): self.sparsity_ignore_list = sparsity_ignore_list self.config = config self.packed_modules_mapping = packed_modules_mapping or {} + # FP8 config for linear layers, compressed tensor currently does not support block fp8, this is used for ktransformers + self.linear_fp8_config = linear_fp8_config def get_linear_method(self) -> CompressedTensorsLinearMethod: return CompressedTensorsLinearMethod(self) @@ -128,10 +128,10 @@ class CompressedTensorsConfig(QuantizationConfig): return None if isinstance(layer, LinearBase): - if CompressedTensorsConfig.DeepSeekFP8Config is not None: - return Fp8LinearMethod(CompressedTensorsConfig.DeepSeekFP8Config) - if envs.SGLANG_KT_MOE_AMX_WEIGHT_PATH.is_set(): - return UnquantizedLinearMethod() + # If linear_fp8_config is set, use FP8 for linear layers + # This allows mixed quantization: experts with int4, linear layers with fp8 + if self.linear_fp8_config is not None: + return Fp8LinearMethod(self.linear_fp8_config) scheme = self.get_scheme(layer=layer, layer_name=prefix) if scheme is None: return UnquantizedLinearMethod() @@ -140,7 +140,6 @@ class CompressedTensorsConfig(QuantizationConfig): from sglang.srt.layers.moe.fused_moe_triton import FusedMoE if isinstance(layer, FusedMoE): - # Ktransformers use CompressedTensorsWNA16AMXMOEMethod if AMX weights are provided return CompressedTensorsMoEMethod.get_moe_method(self, layer, prefix) return None @@ -154,6 +153,23 @@ class CompressedTensorsConfig(QuantizationConfig): ) packed_modules_mapping = config.get("packed_modules_mapping", {}) + # Parse linear_fp8_config if present (for mixed quantization scenarios) + # Format: {"activation_scheme": "dynamic", "fmt": "e4m3", + # "quant_method": "fp8", "weight_block_size": [128, 128]} + linear_fp8_config = None + if "linear_fp8_config" in config: + from sglang.srt.layers.quantization.fp8 import Fp8Config + + fp8_cfg = config["linear_fp8_config"] + # Check if it's fp8 format based on quant_method field + is_fp8 = fp8_cfg.get("quant_method") == "fp8" + linear_fp8_config = Fp8Config( + is_checkpoint_fp8_serialized=is_fp8, + activation_scheme=fp8_cfg.get("activation_scheme", "dynamic"), + ignored_layers=fp8_cfg.get("ignored_layers"), + weight_block_size=fp8_cfg.get("weight_block_size"), + ) + return cls( target_scheme_map=target_scheme_map, ignore=ignore, @@ -162,6 +178,7 @@ class CompressedTensorsConfig(QuantizationConfig): sparsity_ignore_list=sparsity_ignore_list, config=config, packed_modules_mapping=packed_modules_mapping, + linear_fp8_config=linear_fp8_config, ) @classmethod diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 039990d09..30883e391 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -4,7 +4,6 @@ from __future__ import annotations import enum import logging -import re from enum import Enum from typing import TYPE_CHECKING @@ -15,19 +14,10 @@ try: except ImportError: FUSED_MARLIN_MOE_AVAILABLE = False -try: - from kt_kernel import AMXMoEWrapper - - KTRANSFORMERS_AVAILABLE = True -except ImportError: - KTRANSFORMERS_AVAILABLE = False - import torch from compressed_tensors import CompressionFormat from compressed_tensors.quantization import QuantizationStrategy -from sglang.srt.distributed import get_tensor_model_parallel_rank -from sglang.srt.environ import envs from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase @@ -43,13 +33,7 @@ from sglang.srt.layers.quantization.utils import ( per_tensor_dequantize, replace_parameter, ) -from sglang.srt.utils import ( - get_bool_env_var, - get_compiler_backend, - is_cuda, - is_hip, - set_weight_attrs, -) +from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, set_weight_attrs if TYPE_CHECKING: from sglang.srt.layers.moe.fused_moe_triton import FusedMoE @@ -83,13 +67,6 @@ def _mask_topk_ids_cpu_experts(topk_ids: torch.Tensor, num_gpu_experts: int): topk_ids[topk_ids >= num_gpu_experts] = -1 -@torch.compile(dynamic=True, backend=get_compiler_backend()) -def mask_cpu_expert_ids(topk_ids: torch.Tensor, num_gpu_experts: int): - """mask CPU expert IDs.""" - _mask_topk_ids_cpu_experts(topk_ids, num_gpu_experts) - return topk_ids - - class GPTQMarlinState(Enum): REPACK = enum.auto() READY = enum.auto() @@ -99,7 +76,6 @@ __all__ = [ "CompressedTensorsMoEMethod", "CompressedTensorsW8A8Fp8MoEMethod", "CompressedTensorsWNA16MoEMethod", - "CompressedTensorsWNA16AMXEPMoEMethod", # for Ktransformers ] @@ -118,16 +94,6 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase): # TODO: @dsikka: refactor this to use schemes as other kernels # are supported + check if the layer is being ignored. - if envs.SGLANG_KT_MOE_AMX_WEIGHT_PATH.is_set(): - match = re.search(r"(\d+)\.mlp", prefix) - if not match: - raise ValueError( - f"Unable to extract layer number from prefix '{prefix}'. " - f"Expected format: '.mlp'" - ) - layer_number = int(match.group(1)) - return CompressedTensorsWNA16AMXEPMoEMethod(quant_config, layer_number) - weight_quant = quant_config.target_scheme_map["Linear"].get("weights") input_quant = quant_config.target_scheme_map["Linear"].get("input_activations") if quant_config._is_wNa16_group_channel(weight_quant, input_quant): @@ -432,9 +398,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): params_dtype: torch.dtype, **extra_weight_attrs, ): - if self.num_gpu_experts != -1: - num_experts = self.num_gpu_experts - # Will transpose the loaded weight along the # intermediate and hidden dim sizes. Will # shard for TP along the transposed dims @@ -689,379 +652,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): sort_indices2=layer.w2_g_idx_sort_indices, num_bits=self.num_bits, is_k_full=self.is_k_full, - ) - return StandardCombineInput(hidden_states=output) - - -class CompressedTensorsWNA16AMXMoEMethod(CompressedTensorsMoEMethod): - """AMX MoE method using AMXMoEWrapper for CPU inference.""" - - def __init__( - self, - quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501 - layer_idx, - num_gpu_experts, - cpuinfer, - 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." - ) - - if not FUSED_MARLIN_MOE_AVAILABLE: - raise ImportError("fused_marlin_moe is not available") - - self.tp_rank = get_tensor_model_parallel_rank() - self.layer_idx = layer_idx - self.num_gpu_experts = num_gpu_experts - self.amx_weight_path = amx_weight_path - self.chunked_prefill_size = chunked_prefill_size - 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, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - 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 - self.moe_intermediate_size = extra_weight_attrs.pop("intermediate_size_full") - - if self.tp_rank != 0: - return - self.amx_wrapper = AMXMoEWrapper( - layer_idx=self.layer_idx, - num_experts=num_experts, - num_experts_per_tok=self.num_experts_per_tok, - hidden_size=hidden_size, - moe_intermediate_size=self.moe_intermediate_size, - num_gpu_experts=self.num_gpu_experts, - cpuinfer_threads=self.cpuinfer, - 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, - ) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - if self.tp_rank != 0: - return - - if self.amx_wrapper is None: - raise RuntimeError( - "AMXMoEWrapper not initialized. Call create_weights first." - ) - - torch.cuda.synchronize() - # Load weights using wrapper - from sglang.srt.eplb.expert_location_dispatch import ( - get_global_expert_location_metadata, - ) - - physical_to_logical_map_cpu = ( - get_global_expert_location_metadata() - .physical_to_logical_map_cpu[self.layer_idx] - .contiguous() - ) - self.amx_wrapper.load_weights(physical_to_logical_map_cpu) - - def submit( - self, - layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, - ) -> None: - """Submit AMX inference task asynchronously.""" - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - topk_weights, topk_ids, _ = topk_output - - if self.tp_rank != 0 or self.amx_wrapper is None: - return None - - # Submit forward task using wrapper - self.amx_wrapper.submit_forward( - x, topk_ids, topk_weights, torch.cuda.current_stream(x.device).cuda_stream - ) - return None - - def sync(self, x): - """Synchronize and retrieve AMX inference results.""" - if self.tp_rank != 0 or self.amx_wrapper is None: - return torch.zeros_like(x) - - # Sync forward task using wrapper - return self.amx_wrapper.sync_forward( - x, torch.cuda.current_stream(x.device).cuda_stream - ) - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig - ): - self.moe_runner_config = moe_runner_config - - def apply( - self, - layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - """Execute AMX MoE forward pass synchronously.""" - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - topk_weights, topk_ids, _ = topk_output - - if self.tp_rank != 0 or self.amx_wrapper is None: - return StandardCombineInput(hidden_states=torch.zeros_like(x)) - - # Execute forward using wrapper (submit + sync) - output = self.amx_wrapper.forward( - x, topk_ids, topk_weights, torch.cuda.current_stream(x.device).cuda_stream - ) - return StandardCombineInput(hidden_states=output) - - -def override_config( - cls, - num_gpu_experts, - cpuinfer, - threadpool_count, - 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 - if num_gpu_experts is not None: - envs.SGLANG_KT_MOE_NUM_GPU_EXPERTS.set(num_gpu_experts) - if cpuinfer is not None: - envs.SGLANG_KT_MOE_CPUINFER.set(cpuinfer) - if threadpool_count is not None: - envs.SGLANG_KT_THREADPOOL_COUNT.set(threadpool_count) - if amx_weight_path is not None: - envs.SGLANG_KT_MOE_AMX_WEIGHT_PATH.set(amx_weight_path) - if amx_method is not None: - 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): - - def __init__( - self, - quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501 - layer_idx, - ): - self.tp_rank = get_tensor_model_parallel_rank() - - if ( - not envs.SGLANG_KT_MOE_NUM_GPU_EXPERTS.is_set() - or not envs.SGLANG_KT_MOE_CPUINFER.is_set() - or not envs.SGLANG_KT_MOE_AMX_WEIGHT_PATH.is_set() - ): - raise RuntimeError( - "the following arguments are required: --kt-amx-weight-path, --kt-cpuinfer, --kt-num-gpu-experts" - ) - self.num_gpu_experts = envs.SGLANG_KT_MOE_NUM_GPU_EXPERTS.value - cpuinfer = envs.SGLANG_KT_MOE_CPUINFER.value - 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, - layer_idx, - self.num_gpu_experts, - cpuinfer, - 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 - ) - self.layer_id = layer_idx - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - self.global_num_experts = num_experts - self.AMX_method.create_weights( - layer, - num_experts, - hidden_size, - intermediate_size_per_partition, - params_dtype, - **extra_weight_attrs, - ) - self.marlin_method.create_weights( - layer, - num_experts, - hidden_size, - intermediate_size_per_partition, - params_dtype, - **extra_weight_attrs, - ) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - self.AMX_method.process_weights_after_loading(layer) - self.marlin_method.process_weights_after_loading(layer) - - def submit( - self, - layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - """Submit hybrid GPU+CPU MoE task (AMX submission + GPU execution).""" - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - - topk_weights, topk_ids, router_logits = topk_output - - # Submit AMX task if on rank 0 - if self.tp_rank == 0: - self.AMX_method.submit(layer, dispatch_output) - - # Mask CPU expert IDs (>= num_gpu_experts) as -1 so they won't be computed on GPU - topk_ids = mask_cpu_expert_ids(topk_ids, self.num_gpu_experts) - - # Execute GPU (Marlin) experts - output = fused_marlin_moe( - x, - layer.w13_weight_packed, - layer.w2_weight_packed, - layer.w13_weight_scale, - layer.w2_weight_scale, - router_logits, - topk_weights, - topk_ids, - g_idx1=layer.w13_weight_g_idx, - g_idx2=layer.w2_weight_g_idx, - sort_indices1=layer.w13_g_idx_sort_indices, - sort_indices2=layer.w2_g_idx_sort_indices, - num_bits=self.marlin_method.num_bits, - is_k_full=self.marlin_method.is_k_full, - global_num_experts=self.global_num_experts, expert_map=torch.empty(1, device=x.device), ) return StandardCombineInput(hidden_states=output) - - def sync(self, x): - """Synchronize and retrieve AMX results.""" - if self.tp_rank != 0: - return torch.zeros_like(x) - return self.AMX_method.sync(x) - - def apply( - self, - layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - """Execute hybrid GPU+CPU MoE forward pass with parallelism.""" - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - topk_weights, topk_ids, router_logits = topk_output - - # Step 1: Submit AMX task (non-blocking) if on rank 0 - # This starts CPU computation in parallel - if self.tp_rank == 0: - self.AMX_method.submit(layer, dispatch_output) - - # Step 2: Execute GPU (Marlin) experts in parallel with CPU - - # Mask CPU expert IDs (>= num_gpu_experts) as -1 so they won't be computed on GPU - topk_ids = mask_cpu_expert_ids(topk_ids, self.num_gpu_experts) - - # While GPU computes, CPU is also computing - output = fused_marlin_moe( - x, - layer.w13_weight_packed, - layer.w2_weight_packed, - layer.w13_weight_scale, - layer.w2_weight_scale, - router_logits, - topk_weights, - topk_ids, - g_idx1=layer.w13_weight_g_idx, - g_idx2=layer.w2_weight_g_idx, - sort_indices1=layer.w13_g_idx_sort_indices, - sort_indices2=layer.w2_g_idx_sort_indices, - num_bits=self.marlin_method.num_bits, - is_k_full=self.marlin_method.is_k_full, - global_num_experts=self.global_num_experts, - expert_map=torch.empty(1, device=x.device), - ) - - # Step 3: Sync AMX results and combine with GPU results - if self.tp_rank == 0: - amx_output = self.AMX_method.sync(x) - output += amx_output - - return StandardCombineInput(hidden_states=output) - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig - ): - self.moe_runner_config = moe_runner_config - self.AMX_method.create_moe_runner(layer, moe_runner_config) diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index 4d22f1887..3110cbbb7 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/python/sglang/srt/layers/radix_attention.py @@ -81,6 +81,7 @@ class RadixAttention(nn.Module): self.k_scale_float = None self.v_scale_float = None self.quant_method = None + if quant_config is not None: self.quant_method = quant_config.get_quant_method(self, prefix=prefix) if self.quant_method is not None: diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 243b9a84b..58e95ba6a 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -69,7 +69,7 @@ from sglang.srt.utils.patch_torch import monkey_patch_torch_compile from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter try: - from kt_kernel import AMXMoEWrapper + from kt_kernel import KTMoEWrapper KTRANSFORMERS_AVAILABLE = True except ImportError: @@ -259,7 +259,7 @@ class CudaGraphRunner: self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner) log_info_on_rank0(logger, f"Capture cuda graph bs {self.capture_bs}") if KTRANSFORMERS_AVAILABLE: - AMXMoEWrapper.set_capture_batch_sizes(self.capture_bs) + KTMoEWrapper.set_capture_batch_sizes(self.capture_bs) self.capture_forward_mode = ForwardMode.DECODE self.capture_hidden_mode = CaptureHiddenMode.NULL self.num_tokens_per_bs = 1 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 105bbf5d7..b190bf5d4 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -41,7 +41,6 @@ from sglang.srt.distributed import ( get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) -from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo @@ -79,11 +78,10 @@ from sglang.srt.layers.moe import ( ) from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat -from sglang.srt.layers.quantization import CompressedTensorsConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors_moe import ( - CompressedTensorsWNA16AMXEPMoEMethod, CompressedTensorsWNA16MoEMethod, ) from sglang.srt.layers.quantization.fp8 import Fp8Config @@ -785,9 +783,7 @@ class DeepseekV2MoE(nn.Module): final_hidden_states = self.experts(hidden_states, topk_output) if ( not _is_cuda - or isinstance( - self.experts.quant_method, CompressedTensorsWNA16AMXEPMoEMethod - ) + or isinstance(self.experts.quant_method, KTEPWrapperMethod) or isinstance( self.experts.quant_method, CompressedTensorsWNA16MoEMethod ) @@ -853,9 +849,7 @@ class DeepseekV2MoE(nn.Module): if ( not _is_cuda and not _use_aiter - or isinstance( - self.experts.quant_method, CompressedTensorsWNA16AMXEPMoEMethod - ) + or isinstance(self.experts.quant_method, KTEPWrapperMethod) or isinstance(self.experts.quant_method, CompressedTensorsWNA16MoEMethod) ): # fused in biased_grouped_topk so we can skip here @@ -3036,10 +3030,6 @@ class DeepseekV2ForCausalLM(nn.Module): self.config = config self.tp_size = get_tensor_model_parallel_world_size() self.quant_config = quant_config - if envs.SGLANG_KT_MOE_AMX_WEIGHT_PATH.is_set(): - CompressedTensorsConfig.DeepSeekFP8Config = Fp8Config( - True, "dynamic", None, [128, 128] - ) self.determine_num_fused_shared_experts() self.model = DeepseekV2Model( config, quant_config, prefix=add_prefix("model", prefix) @@ -3183,8 +3173,9 @@ class DeepseekV2ForCausalLM(nn.Module): torch.float8_e4m3fn, torch.float8_e4m3fnuz, ): + # For mixed quantization (experts int4, linear fp8), use linear_fp8_config selected_quant_config = getattr( - self.quant_config, "DeepSeekFP8Config", self.quant_config + self.quant_config, "linear_fp8_config", self.quant_config ) weight_block_size = getattr( selected_quant_config, "weight_block_size", None diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 3b04422b1..c392e19c7 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -61,6 +61,7 @@ from sglang.srt.layers.moe import ( ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz @@ -454,6 +455,42 @@ class Glm4MoeSparseMoeBlock(nn.Module): else: return self.forward_deepep(hidden_states, forward_batch) + def forward_normal_dual_stream( + self, + hidden_states: torch.Tensor, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + shared_output = self._forward_shared_experts(hidden_states) + + with torch.cuda.stream(self.alt_stream): + # router_logits: (num_tokens, n_experts) + router_logits = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + final_hidden_states = self.experts(hidden_states, topk_output) + if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod): + final_hidden_states *= self.routed_scaling_factor + + current_stream.wait_stream(self.alt_stream) + with use_symmetric_memory( + parallel_state.get_tp_group(), disabled=not is_allocation_symmetric() + ): + final_hidden_states_out = torch.empty_like(final_hidden_states) + + torch.add(final_hidden_states, shared_output, out=final_hidden_states_out) + final_hidden_states = final_hidden_states_out + if ( + self.tp_size > 1 + and not should_allreduce_fusion + and not use_reduce_scatter + and not should_use_flashinfer_cutlass_moe_fp4_allgather() + ): + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + return final_hidden_states + def forward_normal( self, hidden_states: torch.Tensor, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6d4e77c61..463a78514 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -438,9 +438,9 @@ class ServerArgs: # LMCache enable_lmcache: bool = False - # Ktransformers - kt_amx_weight_path: Optional[str] = None - kt_amx_method: Optional[str] = None + # Ktransformers/AMX expert parallelism + kt_weight_path: Optional[str] = None + kt_method: Optional[str] = None kt_cpuinfer: Optional[int] = None kt_threadpool_count: Optional[int] = None kt_num_gpu_experts: Optional[int] = None @@ -604,9 +604,6 @@ class ServerArgs: self._handle_amd_specifics() self._handle_grammar_backend() - # Handle Ktransformers specific configs - self._handle_ktransformers_configs() - # Handle data parallelism. self._handle_data_parallelism() @@ -1347,39 +1344,6 @@ class ServerArgs: if self.grammar_backend is None: self.grammar_backend = "xgrammar" - def _handle_ktransformers_configs(self): - from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors_moe import ( - CompressedTensorsWNA16AMXEPMoEMethod, - 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, - self.kt_cpuinfer, - self.kt_threadpool_count, - 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): if self.dp_size == 1: self.enable_dp_attention = False @@ -3053,12 +3017,12 @@ class ServerArgs: # Ktransformer server args parser.add_argument( - "--kt-amx-weight-path", + "--kt-weight-path", type=str, help="[ktransformers parameter] The path of the quantized expert weights for amx kernel. A local folder.", ) parser.add_argument( - "--kt-amx-method", + "--kt-method", type=str, default="AMXINT4", help="[ktransformers parameter] Quantization formats for CPU execution.", @@ -3083,9 +3047,8 @@ class ServerArgs: "--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.", + help="[ktransformers parameter] 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( "--enable-double-sparsity",