Refactor KTransformers heterogeneous compute with unified GPU-quantization backend (#12834)

Co-authored-by: Chen Hongtao <56470055+chenht2022@users.noreply.github.com>
Co-authored-by: chenht2022 <cht22@mails.tsinghua.edu.cn>
Co-authored-by: skqliao <skqliao@gmail.com>
Co-authored-by: ovowei <1913953267@qq.com>

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)
 

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

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)

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

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: '<layer_number>.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)

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:

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

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

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,

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",