diff --git a/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md b/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md index 40c1e3af2..97cc097b7 100644 --- a/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md +++ b/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md @@ -69,7 +69,7 @@ in-flight row as shipped. | Qwen-style shared-expert append into routed top-k output | `_append_shared_to_topk_output`
`fused_append_shared_experts_with_weights`
`num_fused_shared_experts` | `python/sglang/srt/models/qwen2_moe.py`
`python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py` | Qwen-style MoE paths can append shared-expert ids and sigmoid gate weights to routed top-k output in one Triton kernel so the shared experts execute inside the fused MoE path | Treat routed top-k plus shared-expert pad / concat ladders as an existing MoE-prep fusion family first. | | Fused MoE dispatch / permute / combine | token permutation
dispatch / combine
grouped top-k
many small MoE support kernels | `python/sglang/srt/layers/moe/fused_moe_triton/layer.py`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py` | `FusedMoE` plus DeepEP / FlashInfer / FuseEP / standard dispatch backends and `permute_fusion=True` | First ask whether the model is missing an existing `FusedMoE`-style path or backend-specific dispatcher path. | | Fused MoE sum + all-reduce | routed MoE followed by explicit sum-reduce kernels | `python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py` | `fuse_sum_all_reduce=True` path in the second MoE GEMM | Before inventing a new MoE reduction fuse, check whether `enable_fused_moe_sum_all_reduce` is simply off or the quant path is incompatible. | -| Fused MoE activation + quant / re-quant | `silu_and_mul_*quant*`
`npu_dequant_swiglu_quant`
`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`
`python/sglang/jit_kernel/nvfp4.py`
`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`
`python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py` | Quantized MoE backends fuse SwiGLU / SiLU-and-mul with FP8 / FP4 / NPU re-quant before the second expert GEMM | If MoE traces show standalone activation then quant kernels, first check whether the quantized fused path is missing. | +| Fused MoE activation + quant / re-quant | `silu_and_mul_*quant*`
`npu_dequant_swiglu_quant`
`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`
`python/sglang/jit_kernel/nvfp4.py`
`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`
`python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py` | Quantized MoE backends fuse SwiGLU / SiLU-and-mul with FP8 / FP4 / NPU re-quant before the second expert GEMM | If MoE traces show standalone activation then quant kernels, first check whether the quantized fused path is missing. | | DeepSeek comm-prep fused RMSNorm + quant / flatten-quant | `fused_rms_fp8_group_quant`
`fused_rms_mxfp4_quant`
`fused_flatten_fp8_group_quant`
`fused_flatten_mxfp4_quant` | `python/sglang/srt/layers/communicator.py`
`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py` | DeepSeek MLA / MHA ROCm paths fuse RMSNorm or flatten with FP8 / MXFP4 quantization for comm / attention prep | On DeepSeek quant traces, split norm + quant or flatten + quant is an existing family, not a new idea. | | NSA fused top-k transform / page-table build | `fast_topk_transform_fused`
`fast_topk_transform_ragged_fused` | `python/sglang/srt/layers/attention/nsa_backend.py` | NSA can fuse top-k selection with paged / ragged index transform instead of separate top-k plus metadata scatter | If NSA top-k metadata work is split, check `SGLANG_NSA_FUSE_TOPK` and backend support first. | | NSA fused quantize + indexed K-cache store | `fused_store_index_k_cache`
`act_quant`
`index_k_with_scale_buffer` | `python/sglang/jit_kernel/fused_store_index_cache.py`
`python/sglang/srt/layers/attention/nsa/nsa_indexer.py` | Single JIT kernel quantizes bf16 K to fp8 + scale and writes directly into NSA index cache | Treat split `act_quant` + buffer-store on CUDA as missing an existing fused store path. | diff --git a/docs_new/docs/advanced_features/pd_disaggregation.mdx b/docs_new/docs/advanced_features/pd_disaggregation.mdx index fdda50ec8..0c604d7eb 100644 --- a/docs_new/docs/advanced_features/pd_disaggregation.mdx +++ b/docs_new/docs/advanced_features/pd_disaggregation.mdx @@ -557,8 +557,6 @@ export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 export SGLANG_DEEPEP_BF16_DISPATCH=0 export ASCEND_USE_FIA=1 -export SGLANG_NPU_FUSED_MOE_MODE=2 - python3 -m sglang.launch_server \ --model-path /path/to/MiMo-V2-Flash-w8a8-all-0512 \ --attention-backend ascend \ @@ -572,7 +570,6 @@ python3 -m sglang.launch_server \ --base-gpu-id 8 \ --disable-radix-cache \ --cuda-graph-bs 1 2 4 8 10 12 14 16 \ - --quantization modelslim \ --speculative-draft-model-quantization unquant \ --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --enable-multi-layer-eagle \ diff --git a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_optimization.mdx b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_optimization.mdx index 17e7b4ae9..215d1ea7e 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_optimization.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_optimization.mdx @@ -304,7 +304,7 @@ Speculative decoding reduces per-token latency by generating draft tokens that a `--quantization modelslim` - Uses the Ascend ModelSlim quantization tool to load W8A8 pre-quantized model weights. This reduces model weight footprint by approximately 50% compared to BF16, allowing larger models to fit in NPU memory with minimal accuracy degradation. The quantization method is auto-detected from the model's `quant_model_description.json` file. + Uses the Ascend ModelSlim quantization tool to load W4A4/W4A8/W8A8/W4A16 pre-quantized model weights. This reduces model weight footprint by approximately 50% (for w8) or 75% (for w4) compared to BF16, allowing larger models to fit in NPU memory with minimal accuracy degradation. The quantization method is auto-detected from the model's `quant_model_description.json` file. `modelslim` @@ -357,11 +357,6 @@ The following environment variables are used in other best practice configuratio Configures the HCCL communication algorithm scheduling to use AIV (Ascend Intelligent Vision) expansion mode, which can improve communication efficiency for certain collective operations. Used in Qwen MoE and R1 non-DSA configurations - - `SGLANG_NPU_FUSED_MOE_MODE` - Controls the fused MoE optimization mode on Ascend NPU. `1` is default; `2` enables a more aggressive fusion strategy (`DISPATCH_FFN_COMBINE`) that can improve MoE dispatch throughput. Mode `2` requires `--quantization modelslim`. Used primarily with DeepSeek R1 models. - `1` or `2` - `SGLANG_NPU_USE_MLAPO=1` (DeepSeek MLA-specific) Adopts the `MLAPO` fusion operator in the MLA (Multi-Head Latent Attention) preprocessing stage for DeepSeek models with MLA architecture. diff --git a/docs_new/docs/hardware-platforms/ascend-npus/best_practice/deepseek_r1.mdx b/docs_new/docs/hardware-platforms/ascend-npus/best_practice/deepseek_r1.mdx index ca1bde875..1d6942aa4 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/best_practice/deepseek_r1.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/best_practice/deepseek_r1.mdx @@ -611,7 +611,6 @@ do export HCCL_BUFFSIZE=800 export HCCL_SOCKET_IFNAME= export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=131072 - export SGLANG_NPU_FUSED_MOE_MODE=2 export SGLANG_USE_AG_AFTER_QLORA=1 export TASK_QUEUE_ENABLE=2 @@ -625,13 +624,12 @@ do --nnodes 1 \ --tp-size 16 \ --mem-fraction-static 0.778 \ - --quantization modelslim \ --max-running-requests 16 \ --disable-radix-cache \ --chunked-prefill-size -1 \ --max-prefill-tokens 60000 \ --moe-a2a-backend ascend_fuseep \ - --deepep-mode normal \ + --fuseep-mode 2 \ --speculative-algorithm NEXTN \ --speculative-num-steps 1 \ --speculative-eagle-topk 1 \ @@ -664,7 +662,6 @@ do export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=64 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_LM_HEAD_TP=8 - export SGLANG_NPU_FUSED_MOE_MODE=1 export TASK_QUEUE_ENABLE=1 python3 -m sglang.launch_server \ @@ -679,8 +676,8 @@ do --dp-size 32 \ --mem-fraction-static 0.82 \ --max-running-requests 1024 \ - --quantization modelslim \ --moe-a2a-backend ascend_fuseep \ + --fuseep-mode 1 \ --enable-dp-attention \ --deepep-mode low_latency \ --moe-dense-tp 1 \ diff --git a/docs_new/docs/hardware-platforms/ascend-npus/best_practice/minimax_m2_5.mdx b/docs_new/docs/hardware-platforms/ascend-npus/best_practice/minimax_m2_5.mdx index 9ea0f9f28..06d7a37f1 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/best_practice/minimax_m2_5.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/best_practice/minimax_m2_5.mdx @@ -192,7 +192,6 @@ export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=140000 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_EXTERNAL_MODEL_PACKAGE=custom_eagle3 -export SGLANG_NPU_FUSED_MOE_MODE=2 export SGLANG_SET_CPU_AFFINITY=1 export STREAMS_PER_DEVICE=32 export TASK_QUEUE_ENABLE=1 @@ -211,8 +210,7 @@ python3 -m sglang.launch_server \ --max-prefill-tokens 67000 \ --cuda-graph-bs 2 4 8 12 16 18 20 22 24 26 \ --moe-a2a-backend ascend_fuseep \ - --deepep-mode auto \ - --quantization modelslim \ + --fuseep-mode 2 \ --speculative-algorithm EAGLE3 \ --speculative-draft-model-path $DRAFT_MODEL_PATH \ --speculative-num-steps 3 \ @@ -300,7 +298,6 @@ export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=160000 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_EXTERNAL_MODEL_PACKAGE=custom_eagle3 -export SGLANG_NPU_FUSED_MOE_MODE=2 export SGLANG_SET_CPU_AFFINITY=1 export STREAMS_PER_DEVICE=32 export TASK_QUEUE_ENABLE=1 @@ -321,8 +318,7 @@ python3 -m sglang.launch_server \ --max-prefill-tokens 80000 \ --cuda-graph-bs 2 4 6 8 \ --moe-a2a-backend ascend_fuseep \ - --deepep-mode auto \ - --quantization modelslim \ + --fuseep-mode 2 \ --speculative-algorithm EAGLE3 \ --speculative-draft-model-path $DRAFT_MODEL_PATH \ --speculative-num-steps 3 \ diff --git a/docs_new/docs/hardware-platforms/ascend-npus/best_practice/qwen3_235b_a22b.mdx b/docs_new/docs/hardware-platforms/ascend-npus/best_practice/qwen3_235b_a22b.mdx index 9e762b6b6..fce595a31 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/best_practice/qwen3_235b_a22b.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/best_practice/qwen3_235b_a22b.mdx @@ -174,7 +174,7 @@ export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=188416 export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 -export SGLANG_NPU_FUSED_MOE_MODE=2 +export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES=100 export SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE=1 @@ -186,7 +186,6 @@ python3 -m sglang.launch_server \ --node-rank 0 \ --attention-backend ascend \ --device npu \ - --quantization modelslim \ --max-running-requests 432 \ --context-length 8192 \ --dtype bfloat16 \ @@ -196,6 +195,7 @@ python3 -m sglang.launch_server \ --ep-dispatch-algorithm static \ --disable-radix-cache \ --moe-a2a-backend ascend_fuseep \ + --fuseep-mode 2 \ --speculative-algorithm EAGLE3 \ --speculative-draft-model-path $DRAFT_MODEL_PATH \ --speculative-num-steps 3 \ diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index 350d66323..3434b2f73 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -1831,11 +1831,6 @@ SGLang supports various environment variables that can be used to configure its Use the MLAPO path on NPU. false - - SGLANG_NPU_FUSED_MOE_MODE - Fused MoE mode selector for NPU. - 1 - SGLANG_NPU_FORWARD_NATIVE_GELUTANH Use the native gelu-tanh activation forward (for Skywork-Reward-Gemma-2-27B-v0.2). diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 9de2fa48c..b12454003 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -577,6 +577,8 @@ class Envs: # Master switch for the experimental TRT-LLM LoRA fast path; when OFF (default) every # fine-grained opt switch reads False, keeping non-experimental paths byte-identical. SGLANG_EXPERIMENTAL_LORA_OPTI = EnvBool(False) + # Enable int4x2 weights loading + SGLANG_NPU_W4A4_NEW_PACKING = EnvBool(False) # Quantize x to int8 in the dispatch operator DEEP_NORMAL_MODE_USE_INT8_QUANT = EnvBool(False) # This argument is deprecated SGLANG_NPU_FUSED_MOE_MODE = EnvInt(1) diff --git a/python/sglang/srt/hardware_backend/npu/moe/activation.py b/python/sglang/srt/hardware_backend/npu/moe/activation.py new file mode 100644 index 000000000..981191c9f --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/moe/activation.py @@ -0,0 +1,183 @@ +from abc import ABC, abstractmethod +from typing import Any, Optional, Tuple + +import torch +import torch.nn.functional as F + +from sglang.srt.distributed.communication_op import ( + tensor_model_parallel_all_gather, +) +from sglang.srt.layers.activation import GeluAndMul +from sglang.srt.runtime_context import get_parallel + + +# ============================================================================= +# Abstract base for all activation variants +# ============================================================================= +class BaseActivation(ABC): + @abstractmethod + def _apply_activation( + self, *args, **kwargs + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: ... + + +# ============================================================================= +# Concrete activation implementations (unchanged except removed 8.) +# ============================================================================= +class NPUSwiglu(BaseActivation): + def _apply_activation(self, hidden_states: torch.Tensor): + return torch.ops.npu.npu_swiglu(hidden_states), None + + +class NPUSwigluQuant(BaseActivation): + def _apply_activation(self, hidden_states: torch.Tensor): + hidden_states, swiglu_out_scale = torch.ops.npu.npu_dequant_swiglu_quant( + hidden_states, + quant_mode=1, + activate_left=True, + ) + return hidden_states, swiglu_out_scale + + +class NPUSwigluQuantWithScales(BaseActivation): + def _apply_activation( + self, + hidden_states: torch.Tensor, + weight_scale: torch.Tensor, + activation_scale: torch.Tensor, + group_index: torch.Tensor, + bias: Optional[torch.Tensor] = None, + quant_scale: Optional[torch.Tensor] = None, + quant_offset: Optional[torch.Tensor] = None, + ): + hidden_states, swiglu_out_scale = torch.ops.npu.npu_dequant_swiglu_quant( + x=hidden_states, + weight_scale=weight_scale, + activation_scale=activation_scale, + bias=bias, + quant_scale=quant_scale, + quant_offset=quant_offset, + group_index=group_index, + activate_left=True, + quant_mode=1, + ) + return hidden_states, swiglu_out_scale + + +class NPUSwigluDeepEPKernel(BaseActivation): + def __init__(self, need_quant: bool = True): + from sgl_kernel_npu.activation.swiglu_quant import swiglu_quant + + self._kernel = swiglu_quant + self.need_quant = need_quant + + def _apply_activation( + self, + hidden_states: torch.Tensor, + group_list: torch.Tensor, + group_list_type: int, + ): + hidden_states, per_token_scale = self._kernel( + hidden_states, group_list, group_list_type, need_quant=self.need_quant + ) + if self.need_quant: + return hidden_states, per_token_scale + return hidden_states, None + + +class NPUGeluAndMul(BaseActivation): + def __init__(self): + self._gelu = GeluAndMul() + + def _apply_activation(self, hidden_states: torch.Tensor): + return self._gelu(hidden_states), None + + +class NPUSwigluOAI(BaseActivation): + def __init__(self, moe_runner_config=None): + from sgl_kernel_npu.activation.swiglu_oai import swiglu_oai_triton + + self._kernel = swiglu_oai_triton + self._moe_runner_config = moe_runner_config + + def _apply_activation(self, hidden_states: torch.Tensor): + # hidden_states is the output of the grouped matmul with shape + # [num_tokens, 2 * inter]. The old swiglu_oai kernel derived the + # gate_up dimension from layer.w13_weight.shape[2], which now fails + # because w13_weight is stored un-transposed. Instead we pass + # the gate_up dimension explicitly from the tensor itself. + alpha = 1.0 + clamp = None + if self._moe_runner_config is not None: + alpha = getattr(self._moe_runner_config, "gemm1_alpha", 1.0) + clamp = getattr(self._moe_runner_config, "gemm1_clamp_limit", None) + + output = self._kernel( + hidden_states, + hidden_states.shape[-1], # gate_up dim = 2 * inter + alpha, + clamp, + ) + return output, None + + +class NPUSwigluStepAndMul(BaseActivation): + def __init__(self, clamp_limit: Optional[float] = None): + self._clamp_limit = clamp_limit + + def _apply_activation(self, hidden_states: torch.Tensor): + if self._clamp_limit is not None: + return self._swiglustep_and_mul(hidden_states, self._clamp_limit), None + return torch.ops.npu.npu_swiglu(hidden_states), None + + @staticmethod + def _swiglustep_and_mul(x: torch.Tensor, limit: float = 7.0) -> torch.Tensor: + gate, up = x.chunk(2, dim=-1) + gate = F.silu(gate).clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + return gate * up + + +# ============================================================================= +# Generic TP all‑gather wrapper – used by the runner when needed +# ============================================================================= +class AllGatherActivationWrapper(BaseActivation): + """ + Wraps any activation and adds an all‑gather along `dim` if TP > 1. + + This allows the runner to stay TP‑agnostic: the wrapper is applied + transparently at construction time. + """ + + def __init__(self, inner: BaseActivation, dim: int = -1): + self.inner = inner + self.dim = dim + + def _apply_activation(self, *args, **kwargs): + out, scale = self.inner._apply_activation(*args, **kwargs) + if get_parallel().tp_size > 1: + out = tensor_model_parallel_all_gather(out, dim=self.dim) + return out, scale + + +# ============================================================================= +# Factory (unchanged, returns *base* activations) +# ============================================================================= +def get_swiglu_variant(method: str, **kwargs: Any) -> BaseActivation: + variants: dict[str, type[BaseActivation]] = { + "standard": NPUSwiglu, + "dequant_swiglu_quant": NPUSwigluQuant, + "dequant_swiglu_quant_with_scales": NPUSwigluQuantWithScales, + "swiglu_quant_deepep_kernel": NPUSwigluDeepEPKernel, + "gelu_and_mul": NPUGeluAndMul, + } + if method == "swiglu_oai": + # The OAI variant now uses the triton kernel that derives the gate_up + # dimension from the tensor itself. No extra parameters are needed. + return NPUSwigluOAI() + if method == "swiglustep_and_mul": + clamp_limit = kwargs.pop("clamp_limit", None) + return NPUSwigluStepAndMul(clamp_limit=clamp_limit) + if method not in variants: + raise ValueError(f"Unknown SwiGLU variant: {method}") + return variants[method]() diff --git a/python/sglang/srt/hardware_backend/npu/moe/finalize_routing.py b/python/sglang/srt/hardware_backend/npu/moe/finalize_routing.py new file mode 100644 index 000000000..22c129160 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/moe/finalize_routing.py @@ -0,0 +1,99 @@ +""" +NPU MoE finalize routing components. + +These classes reassemble expert outputs into the original token order +after the expert computation. A generic TP‑all‑gather wrapper is provided +to transparently gather the hidden dimension when needed (e.g. GGUF with +full weights). +""" + +from abc import ABC, abstractmethod + +import torch + +from sglang.srt.distributed.communication_op import ( + tensor_model_parallel_all_gather, +) +from sglang.srt.runtime_context import get_parallel + + +class BaseFinalizeRouting(ABC): + @abstractmethod + def _finalize_routing( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + expanded_row_idx: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: ... + + +# --------------------------------------------------------------------------- +# Concrete implementations (unchanged) +# --------------------------------------------------------------------------- +class NPUFinalizeRouting(BaseFinalizeRouting): + def __init__(self, drop_pad_mode: int = 0): + self.drop_pad_mode = drop_pad_mode + + def _finalize_routing( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + expanded_row_idx: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.npu.npu_moe_finalize_routing( + hidden_states, + skip1=None, + skip2=None, + bias=None, + scales=topk_weights, + expanded_src_to_dst_row=expanded_row_idx, + export_for_source_row=topk_ids, + drop_pad_mode=self.drop_pad_mode, + ) + + +class NPUMoETokenUnpermute(BaseFinalizeRouting): + def _finalize_routing( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + expanded_row_idx: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.npu.npu_moe_token_unpermute( + permuted_tokens=hidden_states, + sorted_indices=expanded_row_idx.abs(), + probs=topk_weights, + ) + + +# --------------------------------------------------------------------------- +# Generic TP‑all‑gather wrapper – transparently adds communication +# --------------------------------------------------------------------------- +class AllGatherFinalizeRoutingWrapper(BaseFinalizeRouting): + """ + Wraps any finalize routing and performs an all‑gather along `dim` + after the routing if tensor‑parallelism is active. + + This keeps the runner / permute hooks free of TP logic. + """ + + def __init__(self, inner: BaseFinalizeRouting, dim: int = -1): + self.inner = inner + self.dim = dim + + def _finalize_routing( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + expanded_row_idx: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: + out = self.inner._finalize_routing( + hidden_states, topk_weights, expanded_row_idx, topk_ids + ) + if get_parallel().tp_size > 1: + out = tensor_model_parallel_all_gather(out, dim=self.dim) + return out diff --git a/python/sglang/srt/hardware_backend/npu/moe/fuseep.py b/python/sglang/srt/hardware_backend/npu/moe/fuseep.py index f2ae0f3ca..941c84685 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/fuseep.py +++ b/python/sglang/srt/hardware_backend/npu/moe/fuseep.py @@ -14,9 +14,10 @@ import torch from sglang.srt.distributed import get_tp_group from sglang.srt.environ import envs -from sglang.srt.hardware_backend.npu.utils import FusedMoEMode, npu_format_cast +from sglang.srt.hardware_backend.npu.utils import npu_format_cast from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer from sglang.srt.layers.moe.utils import DeepEPMode +from sglang.srt.runtime_context import get_server_args if TYPE_CHECKING: from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE @@ -56,7 +57,7 @@ def forward_fuseep( envs.SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() ), num_experts=layer.num_experts, - fuse_mode=envs.SGLANG_NPU_FUSED_MOE_MODE.get(), + fuse_mode=get_server_args().fuseep_mode, ) return hidden_states @@ -120,52 +121,58 @@ def _scale_from_float_to_int64(scale: torch.Tensor) -> torch.nn.Parameter: return torch.nn.Parameter(converted, requires_grad=False) -def process_fuseep_weights(layer: torch.nn.Module) -> None: - """Apply the Ascend FuseEP-specific weight layout. +def process_fuseep_weights(layer: torch.nn.Module, weight_prefix: str) -> None: + """Apply the Ascend FuseEP-specific weight layout for a single weight group. - Replaces NPU quant_method weight layouts with the form required by the - fused_deep_moe op. Invoked from NPU ``process_weights_after_loading`` - when ``--moe-a2a-backend ascend_fuseep`` is set. + Invoked by ``maybe_apply_fuseep_weights`` for both ``"w13"`` and ``"w2"``. """ - if envs.SGLANG_NPU_FUSED_MOE_MODE.get() == FusedMoEMode.DISPATCH_FFN_COMBINE.value: - w13_weight = _release_weight_cache(layer.w13_weight) - layer.w13_weight.data = npu_format_cast(w13_weight) - w2_weight = _release_weight_cache(layer.w2_weight) - layer.w2_weight.data = npu_format_cast(w2_weight) + if get_server_args().fuseep_mode == 1: + # -- The fused MoE optimization mode "1": dispatch_gmm_combine_decode -- + if weight_prefix == "w13": + cpu_w13 = layer.w13_weight.data.transpose(1, 2).cpu() + layer.w13_weight.data = _reshape_w13_weight(cpu_w13, -1).npu() + w13_scale = layer.w13_weight_scale.data.squeeze(-1).contiguous() + w13_scale = _permute_w13_weight_scale(w13_scale, 128) + layer.w13_weight_scale = torch.nn.Parameter( + w13_scale.to(torch.float32), requires_grad=False + ) + layer.w13_weight.data = npu_format_cast(layer.w13_weight.data) + else: # weight_prefix == "w2" + layer.w2_weight.data = npu_format_cast(layer.w2_weight.data) + w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous() + layer.w2_weight_scale = torch.nn.Parameter( + w2_scale.to(torch.float32), requires_grad=False + ) + elif get_server_args().fuseep_mode == 2: + # -- The fused MoE optimization mode "2": dispatch_ffn_combine -- + if weight_prefix == "w13": + w13_weight = _release_weight_cache(layer.w13_weight) + layer.w13_weight.data = npu_format_cast(w13_weight) + layer.w13_weight_scale.data = layer.w13_weight_scale.data.view( + layer.w13_weight_scale.data.shape[0], -1 + ) + layer.w13_weight_scale = _scale_from_float_to_int64( + layer.w13_weight_scale.data + ) + else: # weight_prefix == "w2" + w2_weight = _release_weight_cache(layer.w2_weight) + layer.w2_weight.data = npu_format_cast(w2_weight) + w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous() + layer.w2_weight_scale = torch.nn.Parameter( + w2_scale.to(torch.float32), requires_grad=False + ) + layer.w2_weight_scale = _scale_from_float_to_int64( + layer.w2_weight_scale.data + ) - layer.w13_weight_scale.data = layer.w13_weight_scale.data.view( - layer.w13_weight_scale.data.shape[0], -1 - ) - w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous() - layer.w2_weight_scale = torch.nn.Parameter( - w2_scale.to(torch.float32), requires_grad=False - ) - - layer.w13_weight_scale = _scale_from_float_to_int64(layer.w13_weight_scale.data) - layer.w2_weight_scale = _scale_from_float_to_int64(layer.w2_weight_scale.data) - else: - cpu_w13 = layer.w13_weight.data.transpose(1, 2).cpu() - layer.w13_weight.data = _reshape_w13_weight(cpu_w13, -1).npu() - w13_scale = layer.w13_weight_scale.data.squeeze(-1).contiguous() - w13_scale = _permute_w13_weight_scale(w13_scale, 128) - layer.w13_weight_scale = torch.nn.Parameter( - w13_scale.to(torch.float32), requires_grad=False - ) - layer.w13_weight.data = npu_format_cast(layer.w13_weight.data) - layer.w2_weight.data = npu_format_cast(layer.w2_weight.data) - - w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous() - layer.w2_weight_scale = torch.nn.Parameter( - w2_scale.to(torch.float32), requires_grad=False - ) - - if hasattr(layer, "w13_weight_offset"): - layer.w13_weight_offset = torch.nn.Parameter( - layer.w13_weight_offset.data.squeeze(-1).contiguous(), - requires_grad=False, - ) - if hasattr(layer, "w2_weight_offset"): - layer.w2_weight_offset = torch.nn.Parameter( - layer.w2_weight_offset.data.squeeze(-1).contiguous(), - requires_grad=False, + # -- offsets (exist or not, same logic for both prefixes) --------------- + offset_attr = f"{weight_prefix}_weight_offset" + if hasattr(layer, offset_attr): + setattr( + layer, + offset_attr, + torch.nn.Parameter( + getattr(layer, offset_attr).data.squeeze(-1).contiguous(), + requires_grad=False, + ), ) diff --git a/python/sglang/srt/hardware_backend/npu/moe/hidden_states_quant.py b/python/sglang/srt/hardware_backend/npu/moe/hidden_states_quant.py new file mode 100644 index 000000000..99a5c5ed8 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/moe/hidden_states_quant.py @@ -0,0 +1,72 @@ +""" +Hidden state quantization utilities for NPU MoE. + +Each class quantises hidden states and returns a (quantized_tensor, scale) tuple. +For static quantization the scale is ``None``. +""" + +from abc import ABC, abstractmethod +from typing import Optional, Tuple + +import torch + + +class BaseHiddenStatesQuant(ABC): + """Abstract base for NPU hidden state quantisation.""" + + def __init__(self, quant_dtype: torch.dtype) -> None: + self.quant_dtype = quant_dtype + + @abstractmethod + def __call__( + self, hidden_states: torch.Tensor, **kwargs + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: ... + + +class HiddenStatesDynamicQuant(BaseHiddenStatesQuant): + """ + Dynamic per‑token quantisation of hidden states. + + Returns ``(quantized_hidden_states, per‑token_scale)``. + """ + + def __call__( + self, hidden_states: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + quantized, scale = torch.ops.npu.npu_dynamic_quant( + hidden_states, dst_type=self.quant_dtype + ) + return quantized, scale + + +class HiddenStatesStaticQuant(BaseHiddenStatesQuant): + """ + Static quantisation using pre‑computed layer‑specific scales and offsets. + + The ``layer`` argument must expose ``aclnn_input_scale_reciprocal`` and + ``aclnn_input_offset``. Returns ``(quantized_hidden_states, None)``. + """ + + def __call__( + self, + hidden_states: torch.Tensor, + layer: torch.nn.Module, + ) -> Tuple[torch.Tensor, None]: + # Optional defensive check (as suggested in the review) + if not hasattr(layer, "aclnn_input_scale_reciprocal") or not hasattr( + layer, "aclnn_input_offset" + ): + raise AttributeError( + "Static quantisation requires layer attributes " + "'aclnn_input_scale_reciprocal' and 'aclnn_input_offset'." + ) + + quantized = torch.ops.npu.npu_quantize( + hidden_states, + layer.aclnn_input_scale_reciprocal, + layer.aclnn_input_offset, + self.quant_dtype, + -1, + False, + ) + return quantized, None diff --git a/python/sglang/srt/hardware_backend/npu/moe/init_routing.py b/python/sglang/srt/hardware_backend/npu/moe/init_routing.py new file mode 100644 index 000000000..bce2acf19 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/moe/init_routing.py @@ -0,0 +1,129 @@ +""" +NPU MoE init routing components. + +Prepare token routing before expert computation. Two API versions are provided: +- v1: legacy routing using ``npu_moe_init_routing``. +- v2: improved routing using ``npu_moe_init_routing_v2``. +""" + +from abc import ABC, abstractmethod +from typing import Optional, Tuple + +import torch + + +class BaseInitRouting(ABC): + """Abstract base for NPU MoE init routing.""" + + @abstractmethod + def _init_routing( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + top_k: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: ... + + +class NPUMoEInitRouting_v1(BaseInitRouting): + """ + NPU MoE init routing (v1 API). + + Uses ``npu_moe_init_routing`` with a manually constructed ``row_idx`` tensor. + """ + + def _init_routing( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + top_k: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + num_tokens = hidden_states.shape[0] + row_idx_len = num_tokens * top_k + row_idx = ( + torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_ids.device) + .view(topk_ids.shape[1], -1) + .permute(1, 0) + .contiguous() + ) + + hidden_states, expanded_row_idx, expanded_expert_idx = ( + torch.ops.npu.npu_moe_init_routing( + hidden_states, + row_idx=row_idx, + expert_idx=topk_ids, + active_num=num_tokens, + ) + ) + expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens( + expanded_expert_idx, num_experts + ) + expert_tokens = expert_tokens.to(torch.int64) + return hidden_states, expanded_row_idx, expert_tokens, None + + +class NPUMoEInitRouting_v2(BaseInitRouting): + """ + NPU MoE init routing (v2 API). + + Uses ``npu_moe_init_routing_v2``, which integrates expert token counting. + """ + + def __init__(self, quant_mode: int = -1): + self.quant_mode = quant_mode + + def _init_routing( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + top_k: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + num_tokens = hidden_states.shape[0] + hidden_states, expanded_row_idx, expert_tokens, pertoken_scale = ( + torch.ops.npu.npu_moe_init_routing_v2( + hidden_states, + topk_ids, + active_num=num_tokens * top_k, + expert_num=num_experts, + expert_tokens_num_type=1, + expert_tokens_num_flag=True, + active_expert_range=[0, num_experts], + quant_mode=self.quant_mode, + ) + ) + if self.quant_mode == -1: + pertoken_scale = None + expert_tokens = expert_tokens.to(torch.int64) + return hidden_states, expanded_row_idx, expert_tokens, pertoken_scale + + +class NPUMoEInitRouting_Quant(BaseInitRouting): + """ + NPU MoE init routing (Quant API). + + Uses ``npu_moe_init_routing_quant``, which integrates expert token counting. + """ + + def _init_routing( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + num_tokens = hidden_states.shape[0] + + hidden_states, expanded_row_idx, expert_tokens, _, pertoken_scale = ( + torch.ops.npu.npu_moe_init_routing_quant( + hidden_states, + topk_ids, + active_num=num_tokens * topk_ids.shape[1], + expert_num=num_experts, + expert_tokens_num_mode=1, + expert_tokens_before_capacity_flag=False, + quant_mode=1, + ) + ) + expert_tokens = expert_tokens.to(torch.int64) + return hidden_states, expanded_row_idx, expert_tokens, pertoken_scale diff --git a/python/sglang/srt/hardware_backend/npu/moe/matmul.py b/python/sglang/srt/hardware_backend/npu/moe/matmul.py new file mode 100644 index 000000000..1ebcfc833 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/moe/matmul.py @@ -0,0 +1,49 @@ +from abc import ABC, abstractmethod + +import torch + + +class BaseMatmul(ABC): + @abstractmethod + def forward( + self, + layer: torch.nn.Module, + weight_prefix: str, + hidden_states: torch.Tensor, + expert_tokens: torch.Tensor, + output_dtype: torch.dtype, + group_list_type: int, + transposed: bool, + **scale_args, + ) -> torch.Tensor: + pass + + +class GroupedMatmul(BaseMatmul): + def forward( + self, + layer: torch.nn.Module, + weight_prefix: str, + hidden_states: torch.Tensor, + expert_tokens: torch.Tensor, + output_dtype: torch.dtype, + group_list_type: int, + transposed: bool, + **scale_args, + ) -> torch.Tensor: + # Access the weight attribute directly from the layer + weight = getattr(layer, f"{weight_prefix}_weight", None) + if weight is None: + raise AttributeError( + f"Weight attribute '{weight_prefix}_weight' not found in layer" + ) + return torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[weight] if transposed else [weight.transpose(1, 2)], + **scale_args, + split_item=2, + group_list_type=group_list_type, + group_type=0, + group_list=expert_tokens, + output_dtype=output_dtype, + )[0] diff --git a/python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py b/python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py index 197c3de31..9ecfccd4b 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py @@ -3,45 +3,139 @@ from __future__ import annotations from typing import TYPE_CHECKING, Optional import torch +import torch.nn.functional as F +import torch_npu -from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( - NPUW4A16Int4DynamicMoEMethod, +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUWNA16Int4MoEMethod, ) from sglang.srt.layers.quantization.utils import replace_parameter if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput from sglang.srt.layers.quantization.base_config import QuantizationConfig -import torch_npu - class AWQAscendLinearKernel: def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Keep scales as (groups, N) – NPU kernel expects this layout layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) - qweight_tmp = torch.zeros_like(layer.qweight.data) - qzeros_tmp = layer.qzeros.data - qzeros_list = [] + + raw_qweight = layer.qweight.data # (K, N // pack_factor) + raw_qzeros = layer.qzeros.data # (groups, N // pack_factor) + + pack_factor = self.quant_config.pack_factor + # shifts control which 4-bit nibble we extract from each packed byte: + # byte = [nibble_7 | nibble_6 | ... | nibble_0] + # shift = 4*i gives the i-th nibble's bit offset. shifts = [0, 4, 1, 5, 2, 6, 3, 7] - for i in range(0, self.quant_config.pack_factor): - shift_num = shifts[i] * 4 - qzeros_list.append((qzeros_tmp.reshape(-1, 1) >> shift_num) & 0xF) - qweight_tmp.bitwise_or_( - ((layer.qweight.data >> shift_num) & 0xF) << (4 * i) + K = raw_qweight.shape[0] + N = raw_qweight.shape[1] * pack_factor + num_groups = layer.scales.shape[0] + + if K % num_groups != 0: + raise RuntimeError(f"K={K} not divisible by scale groups {num_groups}") + group_size = K // num_groups + + # NPU fast-path constraint: + # The NPU's `npu_weight_quant_batchmatmul` kernel requires group_size + # to be a multiple of 32 and at least 32, but less than K (otherwise + # per-tensor scaling would apply, which is a different code path). + # This aligns with the NPU's SIMD vectorization width (32 elements) + # and ensures efficient memory access patterns. + is_support_npu_quant_mm = (group_size == 0) or ( + group_size % 32 == 0 and 32 <= group_size < K + ) + + if is_support_npu_quant_mm: + # ----- NPU fast path: unsigned weight + raw zero point ----- + # The NPU kernel expects: + # 1. qweight: packed unsigned 4-bit values (no XOR) + # 2. zeros: raw zero-point values (not dequantized) + # + # Step 1: Pack weight as unsigned nibbles (NO XOR). + # We extract each 4-bit nibble from the original packed tensor + # and repack them into a new tensor where each byte contains + # two 4-bit values in the order expected by the NPU kernel. + qweight_tmp = torch.zeros_like(raw_qweight) + qzeros_list = [] + for i in range(pack_factor): + shift_num = shifts[i] * 4 + qzeros_list.append((raw_qzeros.reshape(-1, 1) >> shift_num) & 0xF) + qweight_tmp.bitwise_or_( + ((layer.qweight.data >> shift_num) & 0xF) << (4 * i) + ) + + # Step 2: XOR with 0x88888888 to convert from signed to unsigned + # representation. The original weights are stored as signed int4 + # (values -8..7). XOR with 0x8 flips the sign bit, mapping + # -8 → 0, -7 → 1, ..., 7 → 15. This yields the unsigned + # representation the NPU kernel expects. + # + # Mathematical formula: + # unsigned_val = signed_val ^ 0x8 (for each 4-bit nibble) + # Since we pack two nibbles per byte, we XOR the whole byte + # with 0x88 to flip both sign bits simultaneously. + qweight_tmp.bitwise_xor_( + 0x88888888 + ) # 0x88 per byte = flip sign bit of both nibbles + + # Step 3: Convert zero points from signed to unsigned. + # The zero points are stored as signed int4 (-8..7). + # We convert them to unsigned (0..15) by subtracting 8, + # then negate to get the raw zero-point value expected by the NPU. + # unsigned_zero = signed_zero + 8 + # raw_zero = -unsigned_zero + qzeros_tmp = torch.cat(qzeros_list, dim=-1).reshape(raw_qzeros.shape[0], -1) + qzeros_tmp = -(qzeros_tmp - 8) # convert signed → unsigned → negated + qzeros_tmp = qzeros_tmp.to(layer.scales.data.dtype) + + layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False) + layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False) + + layer.use_npu_matmul = True + layer.npu_group_size = group_size + else: + # ----- Fallback: asymmetric dequantisation on CPU/NPU via standard linear ----- + # When group_size doesn't meet the NPU constraint, we fall back to + # a standard dequantisation + FP16 linear. This is gives memory overhead but correct + # for all group_size values. + weight_u8 = torch.zeros((K, N), dtype=torch.int8, device=raw_qweight.device) + zeros_u8 = torch.zeros( + (num_groups, N), dtype=torch.int8, device=raw_qzeros.device ) - qweight_tmp.bitwise_xor_(0x88888888) + for i in range(pack_factor): + shift = shifts[i] * 4 + nib_w = (raw_qweight >> shift) & 0xF + weight_u8[:, i::pack_factor] = nib_w.to(torch.int8) + nib_z = (raw_qzeros >> shift) & 0xF + zeros_u8[:, i::pack_factor] = nib_z.to(torch.int8) - qzeros_tmp = torch.cat(qzeros_list, dim=-1).reshape(qzeros_tmp.shape[0], -1) - qzeros_tmp = -(qzeros_tmp - 8) - qzeros_tmp = qzeros_tmp.to(layer.scales.data.dtype) + # Dequantize: weight_fp = (weight_u8 - zeros) * scales + if group_size > 0: + zeros_exp = zeros_u8.repeat_interleave(group_size, dim=0) + scales_exp = layer.scales.data.repeat_interleave(group_size, dim=0) + else: + zeros_exp = zeros_u8 + scales_exp = layer.scales.data - layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False) - layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False) + weight_float = (weight_u8.float() - zeros_exp.float()) * scales_exp.float() + weight_float = weight_float.t().contiguous().to(torch.bfloat16) + + layer.register_parameter( + "weight", torch.nn.Parameter(weight_float, requires_grad=False) + ) + delattr(layer, "scales") + layer.use_npu_matmul = False + + # Clean original packed tensors to free memory + for attr in ("qweight", "qzeros"): + if hasattr(layer, attr): + delattr(layer, attr) def apply( self, @@ -49,32 +143,41 @@ class AWQAscendLinearKernel: x: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: - qweight = layer.weight - scales = layer.scales - qzeros = layer.zeros - pack_factor = self.quant_config.pack_factor - out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,) reshaped_x = x.reshape(-1, x.shape[-1]) + pack_factor = self.quant_config.pack_factor - if bias is not None and bias.dtype == torch.bfloat16: - bias = bias.float() + if layer.use_npu_matmul: + qweight = layer.weight # (K, N//pack) int32, unsigned + scales = layer.scales # (groups, N) + offset = layer.zeros # (groups, N) raw zero point - out = torch_npu.npu_weight_quant_batchmatmul( - reshaped_x, - qweight, - antiquant_scale=scales, - antiquant_offset=qzeros, - antiquant_group_size=self.quant_config.group_size, - bias=bias, - ) + out_shape = x.shape[:-1] + (qweight.shape[1] * pack_factor,) + if bias is not None and bias.dtype == torch.bfloat16: + bias = bias.float() - return out.reshape(out_shape) + # NPU-accelerated quantized matmul. + # The kernel internally does: + # out = (x @ qweight_dequantized) + bias + # where qweight_dequantized = (qweight_unsigned - offset) * scales + # with group-wise scaling applied. + out = torch_npu.npu_weight_quant_batchmatmul( + reshaped_x, + qweight, + antiquant_scale=scales, + antiquant_offset=offset, # raw zero point + antiquant_group_size=layer.npu_group_size, + bias=bias, + ) + return out.reshape(out_shape) + else: + return F.linear(x, layer.weight, bias) class AWQAscendMoEKernel: def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config - self.kernel = NPUW4A16Int4DynamicMoEMethod() + self.w13_kernel = NPUWNA16Int4MoEMethod() + self.w2_kernel = NPUWNA16Int4MoEMethod() @staticmethod def _register_or_replace_parameter( @@ -87,88 +190,69 @@ class AWQAscendMoEKernel: name, torch.nn.Parameter(tensor, requires_grad=False) ) - def _convert_awq_weight_to_npu_layout(self, qweight: torch.Tensor) -> torch.Tensor: - num_experts, input_size, _ = qweight.shape - unpacked_weight = ( - self.kernel._unpack_from_int32(qweight.flatten(0, 1), 4) - .view(num_experts, input_size, -1) - .transpose(1, 2) - .contiguous() - .int() - ) - return self.kernel._pack_to_int32(unpacked_weight) - - def _convert_awq_qzeros_to_npu_offset( - self, qzeros: torch.Tensor, dtype: torch.dtype - ) -> torch.Tensor: - num_experts, num_groups, _ = qzeros.shape - offset = ( - -self.kernel._unpack_from_int32(qzeros.flatten(0, 1), 4) - .view(num_experts, num_groups, -1) - .transpose(1, 2) - .contiguous() - ) - return offset.to(dtype) - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - self._register_or_replace_parameter( - layer, - "w13_weight", - self._convert_awq_weight_to_npu_layout(layer.w13_qweight.data), - ) - self._register_or_replace_parameter( - layer, - "w2_weight", - self._convert_awq_weight_to_npu_layout(layer.w2_qweight.data), - ) - self._register_or_replace_parameter( - layer, - "w13_weight_scale", - layer.w13_scales.data.transpose(1, 2).contiguous(), - ) - self._register_or_replace_parameter( - layer, - "w2_weight_scale", - layer.w2_scales.data.transpose(1, 2).contiguous(), - ) - self._register_or_replace_parameter( - layer, - "w13_weight_offset", - self._convert_awq_qzeros_to_npu_offset( - layer.w13_qzeros.data, layer.w13_scales.data.dtype - ), - ) - self._register_or_replace_parameter( - layer, - "w2_weight_offset", - self._convert_awq_qzeros_to_npu_offset( - layer.w2_qzeros.data, layer.w2_scales.data.dtype - ), - ) + w13_qweight_tmp = torch.zeros_like(layer.w13_qweight.data) + w2_qweight_tmp = torch.zeros_like(layer.w2_qweight.data) + w13_qzeros_list = [] + w2_qzeros_list = [] - self.kernel.process_weights_after_loading(layer) + # shifts control which 4-bit nibble we extract from each packed byte. + # For AWQ with pack_factor=8, each byte contains 8 nibbles (4-bit values). + # shifts = [0,4,1,5,2,6,3,7] extracts nibbles in the order: + # nibble_0, nibble_1, nibble_2, ..., nibble_7 + # but interleaved to match the NPU kernel's expected layout. + shifts = [0, 4, 1, 5, 2, 6, 3, 7] - def apply( - self, - layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, - ) -> torch.Tensor: - return self.kernel.apply(layer, dispatch_output) + for i in range(self.quant_config.pack_factor): + shift_num = shifts[i] * 4 + w13_qzeros_list.append( + (layer.w13_qzeros.data.reshape(-1, 1) >> shift_num) & 0xF + ) + w2_qzeros_list.append( + (layer.w2_qzeros.data.reshape(-1, 1) >> shift_num) & 0xF + ) + w13_qweight_tmp.bitwise_or_( + ((layer.w13_qweight.data >> shift_num) * (2 ** (4 * i))) + & (0xF << (4 * i)) + ) + w2_qweight_tmp.bitwise_or_( + ((layer.w2_qweight.data >> shift_num) * (2 ** (4 * i))) + & (0xF << (4 * i)) + ) - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - return self.kernel.apply_without_routing_weights( - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, + # XOR with 0x88888888 converts signed int4 to unsigned int4. + # Each byte contains two 4-bit values, so 0x88 flips the sign bit + # of both nibbles simultaneously. + # + # signed_val: -8 -7 -6 -5 -4 -3 -2 -1 0 1 2 3 4 5 6 7 + # unsigned_val: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 + # signed_val ^ 0x8 maps -8→0, -7→1, ..., 7→15. + w13_qweight_tmp.bitwise_xor_(0x88888888) + w2_qweight_tmp.bitwise_xor_(0x88888888) + + # Convert zero points: signed int4 → unsigned → negated. + # The NPU kernel expects raw zero-point values (not dequantized). + w13_qzeros_tmp = torch.cat(w13_qzeros_list, dim=-1).reshape( + layer.w13_qzeros.shape[0], layer.w13_qzeros.shape[1], -1 + ) + w13_qzeros_tmp = -(w13_qzeros_tmp - 8) # signed → unsigned → negated + w13_qzeros_tmp = w13_qzeros_tmp.to(layer.w13_scales.data.dtype) + + w2_qzeros_tmp = torch.cat(w2_qzeros_list, dim=-1).reshape( + layer.w2_qzeros.shape[0], layer.w2_qzeros.shape[1], -1 + ) + w2_qzeros_tmp = -(w2_qzeros_tmp - 8) + w2_qzeros_tmp = w2_qzeros_tmp.to(layer.w2_scales.data.dtype) + + layer.register_parameter( + "w13_qzeros", torch.nn.Parameter(w13_qzeros_tmp, requires_grad=False) + ) + layer.register_parameter( + "w13_qweight", torch.nn.Parameter(w13_qweight_tmp, requires_grad=False) + ) + layer.register_parameter( + "w2_qzeros", torch.nn.Parameter(w2_qzeros_tmp, requires_grad=False) + ) + layer.register_parameter( + "w2_qweight", torch.nn.Parameter(w2_qweight_tmp, requires_grad=False) ) diff --git a/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py deleted file mode 100644 index 2fed6d73a..000000000 --- a/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py +++ /dev/null @@ -1,1217 +0,0 @@ -from typing import TYPE_CHECKING, Optional - -import numpy as np -import torch - -from sglang.srt.hardware_backend.npu.utils import npu_format_cast -from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase - -if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import ( - CombineInput, - DeepEPLLDispatchOutput, - DeepEPNormalDispatchOutput, - DispatchOutput, - ) - from sglang.srt.layers.quantization.base_config import QuantizationConfig - - -def npu_fused_experts_w4a4( - hidden_states: torch.Tensor, - w13: torch.Tensor, - w13_scale: torch.Tensor, - w2: torch.Tensor, - w2_scale: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - top_k: int, -): - original_shape = hidden_states.shape - original_dtype = hidden_states.dtype - scale_dtype = original_dtype if original_dtype == torch.bfloat16 else torch.float32 - if len(original_shape) == 3: - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - num_tokens = hidden_states.shape[0] - num_experts = w13.shape[0] - - hidden_states, expanded_row_idx, expert_tokens, _ = ( - torch.ops.npu.npu_moe_init_routing_v2( - hidden_states, - topk_ids, - active_num=num_tokens * top_k, - expert_num=num_experts, - expert_tokens_num_type=1, - expert_tokens_num_flag=True, - active_expert_range=[0, num_experts], - quant_mode=-1, - ) - ) - expert_tokens = expert_tokens.to(torch.int64) - - # gmm1: gate_up_proj - hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant( - hidden_states, dst_type=torch.quint4x2 - ) - scale_args13 = { - "scale": [w13_scale], - "per_token_scale": [pertoken_scale], - } - - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[w13], - **scale_args13, - split_item=2, - group_list_type=1, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - # act_fn: swiglu - hidden_states = torch.ops.npu.npu_swiglu(hidden_states) - hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(hidden_states) - - scale_args2 = { - "scale": [w2_scale.to(scale_dtype)], - "per_token_scale": [pertoken_scale], - } - # gmm2: down_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[w2], - **scale_args2, - split_item=2, - group_list_type=1, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - - final_hidden_states = torch.ops.npu.npu_moe_finalize_routing( - hidden_states, - skip1=None, - skip2=None, - bias=None, - scales=topk_weights, - expanded_src_to_dst_row=expanded_row_idx, - export_for_source_row=topk_ids, - drop_pad_mode=2, - ) - if len(original_shape) == 3: - final_hidden_states = final_hidden_states.view(original_shape) - return final_hidden_states - - -def npu_fused_experts( - hidden_states: torch.Tensor, - w13: torch.Tensor, - w13_scale: torch.Tensor, - w2: torch.Tensor, - w2_scale: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - top_k: int, - **kwargs, -): - w13_offset = kwargs.get("w13_offset", None) - w2_offset = kwargs.get("w2_offset", None) - use_wna16 = kwargs.get("use_wna16", False) - - original_shape = hidden_states.shape - original_dtype = hidden_states.dtype - scale_dtype = original_dtype if original_dtype == torch.bfloat16 else torch.float32 - if len(original_shape) == 3: - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - num_tokens = hidden_states.shape[0] - num_experts = w13.shape[0] - row_idx_len = num_tokens * top_k - row_idx = ( - torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device) - .view(top_k, -1) - .permute(1, 0) - .contiguous() - ) - hidden_states, expanded_row_idx, expanded_expert_idx = ( - torch.ops.npu.npu_moe_init_routing( - hidden_states, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens - ) - ) - expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens( - expanded_expert_idx, num_experts - ) - expert_tokens = expert_tokens.to(torch.int64) - # gmm1: gate_up_proj - if not use_wna16: - hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(hidden_states) - scale_args13 = { - "scale": [w13_scale.to(scale_dtype)], - "per_token_scale": [pertoken_scale], - } - else: - scale_args13 = { - "antiquant_scale": [w13_scale], - "antiquant_offset": [w13_offset], - } - - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[w13], - **scale_args13, - split_item=2, - group_list_type=0, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - # act_fn: swiglu - if not use_wna16: - hidden_states, pertoken_scale = torch.ops.npu.npu_dequant_swiglu_quant( - hidden_states, - activate_left=True, - quant_mode=1, - ) - - scale_args2 = { - "scale": [w2_scale.to(scale_dtype)], - "per_token_scale": [pertoken_scale], - } - else: - hidden_states = torch.ops.npu.npu_swiglu(hidden_states) - scale_args2 = {"antiquant_scale": [w2_scale], "antiquant_offset": [w2_offset]} - # gmm2: down_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[w2], - **scale_args2, - split_item=2, - group_list_type=0, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - - final_hidden_states = torch.ops.npu.npu_moe_finalize_routing( - hidden_states, - skip1=None, - skip2=None, - bias=None, - scales=topk_weights, - expanded_src_to_dst_row=expanded_row_idx, - export_for_source_row=topk_ids, - ) - if len(original_shape) == 3: - final_hidden_states = final_hidden_states.view(original_shape) - return final_hidden_states - - -def npu_fused_experts_w8a8_decode( - hidden_states: torch.Tensor, - w13: torch.Tensor, - w13_scale: torch.Tensor, - w2: torch.Tensor, - w2_scale: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - top_k: int, - **kwargs, -): - num_tokens = hidden_states.shape[:-1].numel() - first_expert_idx = 0 - last_expert_idx = w13.shape[0] - global_num_experts = w13.shape[0] - original_shape = hidden_states.shape - group_list_type = 1 - - sorted_hidden_states, expanded_row_idx, expert_tokens, pertoken_scale = ( - torch.ops.npu.npu_moe_init_routing_v2( - hidden_states, - topk_ids, - active_num=num_tokens * top_k, - expert_num=global_num_experts, - expert_tokens_num_type=group_list_type, - expert_tokens_num_flag=True, - active_expert_range=[first_expert_idx, last_expert_idx], - quant_mode=1, - ) - ) - - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[sorted_hidden_states], - weight=[w13], - scale=[w13_scale], - per_token_scale=[pertoken_scale], - group_list=expert_tokens, - split_item=2, - group_type=0, - group_list_type=group_list_type, - output_dtype=torch.bfloat16, - )[0] - - # act_fn: swiglu - hidden_states, swiglu_out_scale = torch.ops.npu.npu_dequant_swiglu_quant( - hidden_states, quant_mode=1, activate_left=True - ) - - output = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[w2], - scale=[w2_scale], - per_token_scale=[swiglu_out_scale], - group_list=expert_tokens, - split_item=2, - group_type=0, - group_list_type=group_list_type, - output_dtype=torch.bfloat16, - )[0] - - assert original_shape is not None - final_hidden_states = torch.ops.npu.npu_moe_token_unpermute( - permuted_tokens=output, - sorted_indices=torch.abs(expanded_row_idx), - probs=topk_weights, - ) - if len(original_shape) == 3: - final_hidden_states = final_hidden_states.view(original_shape) - - return final_hidden_states - - -def npu_fused_moe_without_routing_weights_bf16( - layer, hidden_states, group_list_type, group_list, output_dtype -): - from sgl_kernel_npu.activation.swiglu_quant import swiglu_quant - - # gmm1: gate_up_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w13_weight.transpose(1, 2)], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - hidden_states, _ = swiglu_quant( - hidden_states, group_list, group_list_type, need_quant=False - ) - # gmm2: down_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w2_weight.transpose(1, 2)], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - return hidden_states - - -def fused_moe_npu( - x, - w1, - w2, - topk_output, - moe_runner_config, -): - # TODO: reuse the codes of UnquantizedFusedMoEMethod-forward_npu - topk_weights, topk_ids, _ = topk_output - original_dtype = x.dtype - num_tokens = x.shape[0] - topk_weights = topk_weights.to(x.dtype) - topk_ids = topk_ids.to(torch.int32) - num_experts = w1.shape[0] - top_k = topk_weights.shape[-1] - row_idx_len = num_tokens * top_k - row_idx = ( - torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device) - .view(top_k, -1) - .permute(1, 0) - .contiguous() - ) - - hidden_states, expanded_row_idx, expanded_expert_idx = ( - torch.ops.npu.npu_moe_init_routing( - x, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens - ) - ) - - expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens( - expanded_expert_idx, num_experts - ) - - expert_tokens = expert_tokens.to(torch.int64) - - # gmm1: gate_up_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[w1.permute(0, 2, 1)], - bias=None, - split_item=2, - group_list_type=0, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - - # act_fn: - if moe_runner_config.activation == "silu": - hidden_states = torch.ops.npu.npu_swiglu(hidden_states) - else: - from sglang.srt.layers.activation import GeluAndMul - - hidden_states = GeluAndMul()(hidden_states) - - # gmm2: down_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[w2.permute(0, 2, 1)], - bias=None, - split_item=2, - group_list_type=0, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - - final_hidden_states = torch.ops.npu.npu_moe_finalize_routing( - hidden_states, - skip1=None, - skip2=None, - bias=None, - scales=topk_weights, - expanded_src_to_dst_row=expanded_row_idx, - export_for_source_row=topk_ids, - ) - return final_hidden_states - - -def maybe_apply_deepep_npu( - quant_method, - layer: torch.nn.Module, - dispatch_output: "DispatchOutput", -) -> Optional["CombineInput"]: - """Route DeepEP dispatch outputs through the NPU compute path. - - Replaces the deprecated DeepEPMoE.forward_npu wrapper: detects DeepEP - normal/LL formats, calls ``quant_method.apply_without_routing_weights``, - and wraps the result in the matching CombineInput. Returns None for - non-DeepEP formats so the caller falls through to its standard path. - """ - from sglang.srt.layers.moe.token_dispatcher import ( - DeepEPLLCombineInput, - DeepEPNormalCombineInput, - ) - from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker - - if not dispatch_output.format.is_deepep(): - return None - - # NOTE: Ascend's Dispatch & Combine does not support FP16 - output_dtype = torch.bfloat16 - group_list_type = 1 - - if DispatchOutputChecker.format_is_deepep_normal(dispatch_output): - if TYPE_CHECKING: - assert isinstance(dispatch_output, DeepEPNormalDispatchOutput) - ( - hidden_states, - hidden_states_scale, - _, - _, - num_recv_tokens_per_expert, - ) = dispatch_output - group_list = torch.tensor( - num_recv_tokens_per_expert, - dtype=torch.int64, - device=hidden_states.device, - ) - combine_cls = DeepEPNormalCombineInput - else: - if TYPE_CHECKING: - assert isinstance(dispatch_output, DeepEPLLDispatchOutput) - ( - hidden_states, - hidden_states_scale, - _, - _, - group_list, - _, - ) = dispatch_output - group_list = group_list.to(torch.int64) - combine_cls = DeepEPLLCombineInput - - hidden_states = quant_method.apply_without_routing_weights( - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ) - - return combine_cls( - hidden_states=hidden_states, - topk_ids=dispatch_output.topk_ids, - topk_weights=dispatch_output.topk_weights, - ) - - -def maybe_apply_fuseep_weights(layer: torch.nn.Module) -> bool: - """Apply the FuseEP weight layout if --moe-a2a-backend is ascend_fuseep. - - Returns True when the FuseEP layout was applied and the caller should - skip its own ``process_weights_after_loading`` body. - """ - from sglang.srt.layers.moe import get_moe_a2a_backend - - if not get_moe_a2a_backend().is_ascend_fuseep(): - return False - from sglang.srt.hardware_backend.npu.moe.fuseep import process_fuseep_weights - - process_fuseep_weights(layer) - return True - - -class _NPUFusedMoEMethodBase(FusedMoEMethodBase): - - def __init__( - self, - quant_config: Optional["QuantizationConfig"] = None, - ): - self.quant_config = quant_config - - def _maybe_apply_deepep( - self, - layer: torch.nn.Module, - dispatch_output: "DispatchOutput", - ) -> Optional["CombineInput"]: - return maybe_apply_deepep_npu(self, layer, dispatch_output) - - @staticmethod - def _maybe_apply_fuseep_weights(layer: torch.nn.Module) -> bool: - return maybe_apply_fuseep_weights(layer) - - -class NPUW4A4Int4DynamicMoEMethod(_NPUFusedMoEMethodBase): - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - layer.w13_weight.data = npu_format_cast( - layer.w13_weight.data.transpose(1, 2).contiguous() - ) - layer.w13_weight.data = self._pack_to_int32( - layer.w13_weight.data.to(torch.int32) - ) - - layer.w2_weight.data = npu_format_cast( - layer.w2_weight.data.transpose(1, 2).contiguous() - ) - - scale_np = layer.w13_weight_scale.data.cpu().numpy() - scale_np.dtype = np.uint32 - scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu() - - layer.w13_weight_scale = torch.nn.Parameter( - scale_uint64_tensor.squeeze(-1), requires_grad=False - ) - layer.w2_weight_scale = torch.nn.Parameter( - layer.w2_weight_scale.data.squeeze(-1), requires_grad=False - ) - - # Compressed-tensors format doesn't have this field - if hasattr(layer, "w13_weight_offset"): - layer.w13_weight_offset = torch.nn.Parameter( - layer.w13_weight_offset.data.squeeze(-1), - requires_grad=False, - ) - if hasattr(layer, "w2_weight_offset"): - layer.w2_weight_offset = torch.nn.Parameter( - layer.w2_weight_offset.data.squeeze(-1), - requires_grad=False, - ) - - # Quantizes in int4 separately from the dispatcher - # since deep_ep does not support quantization in int4 - # dispatching works in bf16 - if hasattr(layer, "dispatcher"): - layer.dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"}) - - def _pack_to_int32(self, weight: torch.Tensor): - # pack 8 int4 to int32, we use a int32 to represent a int4 - assert ( - weight.shape[-1] % 8 == 0 - ), "the last dim of weight needs to be divided by 8" - new_weight = torch.ops.npu.npu_convert_weight_to_int4pack(weight.flatten(0, 1)) - new_weight = new_weight.view(weight.shape[0], weight.shape[1], -1) - return new_weight - - def apply( - self, - layer, - dispatch_output: "DispatchOutput", - ) -> "CombineInput": - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - combine_input = self._maybe_apply_deepep(layer, dispatch_output) - if combine_input is not None: - return combine_input - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - - topk_weights, topk_ids, _ = topk_output - topk_ids = topk_ids.to(torch.int32) - topk_weights = topk_weights.to(x.dtype) - output = npu_fused_experts_w4a4( - hidden_states=x, - w13=layer.w13_weight, - w13_scale=layer.w13_weight_scale, - w2=layer.w2_weight, - w2_scale=layer.w2_weight_scale, - topk_weights=topk_weights, - topk_ids=topk_ids, - top_k=topk_ids.shape[1], - ) - return StandardCombineInput(hidden_states=output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - hidden_states, hidden_states_scale = torch.ops.npu.npu_dynamic_quant( - hidden_states, dst_type=torch.quint4x2 - ) - # gmm1: up_gate_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w13_weight], - scale=[layer.w13_weight_scale], - per_token_scale=[hidden_states_scale], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - # act_fn: swiglu - hidden_states = torch.ops.npu.npu_swiglu(hidden_states) - hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(hidden_states) - - # gmm2: down_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w2_weight], - scale=[layer.w2_weight_scale.to(output_dtype)], - per_token_scale=[pertoken_scale], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - return hidden_states - - -class NPUW8A8Int8DynamicMoEMethod(_NPUFusedMoEMethodBase): - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - if self._maybe_apply_fuseep_weights(layer): - return - layer.w13_weight.data = npu_format_cast( - layer.w13_weight.data.transpose(1, 2).contiguous() - ) - layer.w2_weight.data = npu_format_cast( - layer.w2_weight.data.transpose(1, 2).contiguous() - ) - layer.w13_weight_scale = torch.nn.Parameter( - layer.w13_weight_scale.data.squeeze(-1), requires_grad=False - ) - layer.w2_weight_scale = torch.nn.Parameter( - layer.w2_weight_scale.data.squeeze(-1), requires_grad=False - ) - layer.w13_weight_scale_bf16 = torch.nn.Parameter( - layer.w13_weight_scale.data.to(dtype=torch.bfloat16), requires_grad=False - ) - layer.w2_weight_scale_bf16 = torch.nn.Parameter( - layer.w2_weight_scale.data.to(dtype=torch.bfloat16), requires_grad=False - ) - # Compressed-tensors format doesn't have this field - if hasattr(layer, "w13_weight_offset"): - layer.w13_weight_offset = torch.nn.Parameter( - layer.w13_weight_offset.data.squeeze(-1), - requires_grad=False, - ) - if hasattr(layer, "w2_weight_offset"): - layer.w2_weight_offset = torch.nn.Parameter( - layer.w2_weight_offset.data.squeeze(-1), - requires_grad=False, - ) - - if hasattr(layer, "dispatcher"): - layer.dispatcher.set_quant_config({"dispatcher_output_dtype": "int8"}) - - def apply( - self, - layer, - dispatch_output: "DispatchOutput", - ) -> "CombineInput": - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - combine_input = self._maybe_apply_deepep(layer, dispatch_output) - if combine_input is not None: - return combine_input - - # release fp32 scale to save memory - layer.w13_weight_scale = None - layer.w2_weight_scale = None - - hidden_states = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - - topk_weights, topk_ids, _ = topk_output - topk_ids = topk_ids.to(torch.int32) - topk_weights = topk_weights.to(hidden_states.dtype) - - # prefill - if not torch.npu.is_current_stream_capturing(): - output = npu_fused_experts( - hidden_states=hidden_states, - w13=layer.w13_weight, - w13_scale=layer.w13_weight_scale_bf16, - w2=layer.w2_weight, - w2_scale=layer.w2_weight_scale_bf16, - topk_weights=topk_weights, - topk_ids=topk_ids, - top_k=topk_ids.shape[1], - ) - # decode - else: - output = npu_fused_experts_w8a8_decode( - hidden_states=hidden_states, - w13=layer.w13_weight, - w13_scale=layer.w13_weight_scale_bf16, - w2=layer.w2_weight, - w2_scale=layer.w2_weight_scale_bf16, - topk_weights=topk_weights, - topk_ids=topk_ids, - top_k=topk_ids.shape[1], - ) - - return StandardCombineInput(hidden_states=output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - # gmm1: gate_up_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w13_weight], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=torch.int32, - )[0] - - # act_fn: swiglu - hidden_states, swiglu_out_scale = torch.ops.npu.npu_dequant_swiglu_quant( - x=hidden_states, - weight_scale=layer.w13_weight_scale, - activation_scale=hidden_states_scale, - bias=None, - quant_scale=None, - quant_offset=None, - group_index=group_list, - activate_left=True, - quant_mode=1, - ) - - # gmm2: down_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w2_weight], - scale=[layer.w2_weight_scale_bf16], - per_token_scale=[swiglu_out_scale], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - return hidden_states - - -class NPUW4A8Int8DynamicMoEMethod(_NPUFusedMoEMethodBase): - - def _process_scale( - self, weight: torch.Tensor, scale, per_group_scale, is_per_channel_weight - ): - scale = scale.transpose(1, 2).contiguous() - - if is_per_channel_weight: - scale_np = scale.cpu().numpy() - scale_np.dtype = np.uint32 - scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu() - return scale_uint64_tensor, None - - per_group_scale = per_group_scale.transpose(1, 2).contiguous() - group_num, k, n = weight.shape - # the weight of the new version is reduced by half by pack n, so it needs to be restored - n = n * 2 - per_group_scale = per_group_scale.reshape(group_num, -1, n) - group_num, quantgroup_num, n = per_group_scale.shape - bias = None - - scale_fp32 = (scale * per_group_scale).to(torch.float16).to(torch.float32) - scale_fp32_np = scale_fp32.cpu().numpy() - scale_fp32_np.dtype = np.uint32 - sscale_uint64 = np.zeros((group_num, quantgroup_num, n * 2), dtype=np.uint32) - - sscale_uint64[..., ::2] = scale_fp32_np - - sscale_uint64_buffer = np.frombuffer( - sscale_uint64.tobytes(), dtype=np.int64 - ).copy() - sscale_uint64_tensor = torch.from_numpy(sscale_uint64_buffer).reshape( - group_num, quantgroup_num, n - ) - sscale_uint64_tensor = sscale_uint64_tensor.npu() - return sscale_uint64_tensor, bias - - def _update_bias(self, layer, w13_bias, w2_bias): - layer.w13_scale_bias.data = ( - layer.w13_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1) - ) - layer.w2_scale_bias.data = ( - layer.w2_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1) - ) - - def _pack_to_int32(self, weight: torch.Tensor): - # pack 4 int8(int4*2) to int32, because in pytorch, we need to use int32 to represent int4 - assert ( - weight.shape[-1] % 4 == 0 - ), "the last dim of weight needs to be divided by 4" - return weight.view(torch.int32).contiguous() - - def process_weights_after_loading( - self, layer: torch.nn.Module, is_per_channel_weight, activation_use_clip - ) -> None: - if not activation_use_clip: - self._process_weights_without_clip(layer, is_per_channel_weight) - else: - self._process_weights_with_clip(layer) - - layer.w13_weight = torch.nn.Parameter( - layer.w13_weight.data.transpose(1, 2).contiguous(), requires_grad=False - ) - layer.w2_weight = torch.nn.Parameter( - layer.w2_weight.data.transpose(1, 2).contiguous(), requires_grad=False - ) - - layer.w13_weight.data = npu_format_cast(layer.w13_weight.data) - layer.w2_weight.data = npu_format_cast(layer.w2_weight.data) - - layer.w13_weight.data = self._pack_to_int32(layer.w13_weight.data) - layer.w2_weight.data = self._pack_to_int32(layer.w2_weight.data) - - if hasattr(layer, "dispatcher"): - layer.dispatcher.set_quant_config({"dispatcher_output_dtype": "int8"}) - - def _process_weights_without_clip( - self, layer: torch.nn.Module, is_per_channel_weight - ) -> None: - w13_weight_scale_second = ( - layer.w13_weight_scale_second.data - if hasattr(layer, "w13_weight_scale_second") - else None - ) - w2_weight_scale_second = ( - layer.w2_weight_scale_second.data - if hasattr(layer, "w2_weight_scale_second") - else None - ) - layer.w13_weight_scale.data, w13_bias = self._process_scale( - layer.w13_weight, - layer.w13_weight_scale.data, - w13_weight_scale_second, - is_per_channel_weight, - ) - layer.w2_weight_scale.data, w2_bias = self._process_scale( - layer.w2_weight, - layer.w2_weight_scale.data, - w2_weight_scale_second, - is_per_channel_weight, - ) - if hasattr(layer, "w13_weight_scale_second"): - # scale_second is no longer used, release this part of the memory - del layer.w13_weight_scale_second - del layer.w2_weight_scale_second - del layer.w13_weight_offset_second - del layer.w2_weight_offset_second - - self._update_bias(layer, w13_bias, w2_bias) - - def _process_weights_with_clip(self, layer: torch.nn.Module) -> None: - w13_weight_scale = ( - layer.w13_weight_scale.data.squeeze(-1).contiguous().unsqueeze(1) - ) - w2_weight_scale = ( - layer.w2_weight_scale.data.squeeze(-1).contiguous().unsqueeze(1) - ) - layer.w13_weight_scale = torch.nn.Parameter( - w13_weight_scale, requires_grad=False - ) - layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False) - layer.w13_scale_bias = layer.w13_bias - layer.w2_scale_bias = layer.w2_bias - - def apply( - self, - layer, - dispatch_output: "DispatchOutput", - ) -> "CombineInput": - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - combine_input = self._maybe_apply_deepep(layer, dispatch_output) - if combine_input is not None: - return combine_input - - hidden_states = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - - topk_weights, topk_ids, _ = topk_output - top_k = topk_ids.shape[1] - group_list_type = 1 - original_shape = hidden_states.shape - topk_weights = topk_weights - - num_tokens = hidden_states.shape[:-1].numel() - - first_expert_idx = 0 - last_expert_idx = layer.num_experts - global_num_experts = layer.num_experts - - sorted_hidden_states, expanded_row_idx, expert_tokens, pertoken_scale = ( - torch.ops.npu.npu_moe_init_routing_v2( - hidden_states, - topk_ids, - active_num=num_tokens * top_k, - expert_num=global_num_experts, - expert_tokens_num_type=1, - expert_tokens_num_flag=True, - active_expert_range=[first_expert_idx, last_expert_idx], - quant_mode=1, - ) - ) - - expert_tokens = expert_tokens.to(torch.int64) - - bias1 = [layer.w13_scale_bias] - bias2 = [layer.w2_scale_bias] - w1_scale = [layer.w13_weight_scale] - w2_scale = [layer.w2_weight_scale] - _output_dtype = torch.bfloat16 - - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[sorted_hidden_states], - weight=[layer.w13_weight], - scale=w1_scale, - bias=bias1, - per_token_scale=[pertoken_scale], - group_list=expert_tokens, - split_item=2, - group_type=0, - group_list_type=group_list_type, - output_dtype=_output_dtype, - )[0] - - # act_fn: swiglu - hidden_states = torch.ops.npu.npu_swiglu(hidden_states) - hidden_states, swiglu_out_scale = torch.ops.npu.npu_dynamic_quant(hidden_states) - - output = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w2_weight], - scale=w2_scale, - bias=bias2, - per_token_scale=[swiglu_out_scale], - group_list=expert_tokens, - split_item=2, - group_type=0, - group_list_type=group_list_type, - output_dtype=_output_dtype, - )[0] - - assert original_shape is not None - final_hidden_states = torch.ops.npu.npu_moe_token_unpermute( - permuted_tokens=output, - sorted_indices=torch.abs(expanded_row_idx), - probs=topk_weights, - ) - if len(original_shape) == 3: - final_hidden_states = final_hidden_states.view(original_shape) - - return StandardCombineInput(hidden_states=final_hidden_states) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - from sgl_kernel_npu.activation.swiglu_quant import swiglu_quant - - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w13_weight], - scale=[layer.w13_weight_scale], - bias=[layer.w13_scale_bias], - per_token_scale=[hidden_states_scale], - group_list=group_list, - split_item=2, - group_type=0, - group_list_type=group_list_type, - output_dtype=output_dtype, - )[0] - - hidden_states, swiglu_out_scale = swiglu_quant( - hidden_states, group_list, group_list_type - ) - - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w2_weight], - scale=[layer.w2_weight_scale], - bias=[layer.w2_scale_bias], - per_token_scale=[swiglu_out_scale], - group_list=group_list, - split_item=2, - group_type=0, - group_list_type=group_list_type, - output_dtype=output_dtype, - )[0] - - return hidden_states - - -class NPUW4A16Int4DynamicMoEMethod(_NPUFusedMoEMethodBase): - - def _pack_to_int32(self, weight: torch.Tensor): - assert weight.dim() == 3 - if weight.dtype == torch.int32: - # pack 8 int4 to int32, we use a int32 to represent a int4 - assert ( - weight.shape[-1] % 8 == 0 - ), "the last dim of weight needs to be divided by 8" - new_weight = torch.ops.npu.npu_convert_weight_to_int4pack( - weight.flatten(0, 1) - ) - new_weight = new_weight.view(weight.shape[0], weight.shape[1], -1) - elif weight.dtype == torch.int8: - # pack 4 int8(int4*2) to int32, because in pytorch, we need to use int32 to represent int4 - assert ( - weight.shape[-1] % 4 == 0 - ), "the last dim of weight needs to be divided by 4" - new_weight = weight.view(torch.int32).contiguous() - else: - raise ValueError(f"{weight.dtype=} is not supported !") - return new_weight - - def _unpack_from_int32( - self, - value: torch.Tensor, - num_bits: int, - shape: torch.Size = None, - packed_dim=1, - ) -> torch.Tensor: - """ - Unpacks a tensor of packed int32 weights into individual int8s, maintaining the - original bit range. - - Return tensors in int8 - - :param value: tensor to unpack - :param num_bits: number of bits to unpack each data point into - :param shape: shape to unpack into, used to remove padding - :returns: unpacked int8 tensor - """ - if value.dtype is not torch.int32: - raise ValueError( - f"Expected {torch.int32} but got {value.dtype}, Aborting unpack." - ) - - if num_bits > 8: - raise ValueError("Unpacking is only supported for less than 8 bits") - - pack_factor = 32 // num_bits - - # unpack - mask = (1 << num_bits) - 1 - - if packed_dim == 1: - unpacked = torch.zeros( - (value.shape[0], value.shape[1] * pack_factor), - device=value.device, - dtype=torch.int32, - ) - for i in range(pack_factor): - unpacked[:, i::pack_factor] = (value >> (num_bits * i)) & mask - - # remove padding - if shape is not None: - original_row_size = int(shape[1]) - unpacked = unpacked[:, :original_row_size] - else: - unpacked = torch.zeros( - (value.shape[0] * pack_factor, value.shape[1]), - device=value.device, - dtype=torch.int32, - ) - for i in range(pack_factor): - unpacked[i::pack_factor, :] = (value >> (num_bits * i)) & mask - - # remove padding - original_row_size = int(shape[0]) - unpacked = unpacked[:original_row_size, :] - - # bits are packed in unsigned format, reformat to signed - # update the value range from unsigned to signed - offset = pow(2, num_bits) // 2 - unpacked = (unpacked - offset).to(torch.int8) - - return unpacked - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - w13_weight_scale = layer.w13_weight_scale.data.transpose(-1, -2).contiguous() - w2_weight_scale = layer.w2_weight_scale.data.transpose(-1, -2).contiguous() - layer.w13_weight_scale = torch.nn.Parameter( - w13_weight_scale, requires_grad=False - ) - layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False) - - layer.w13_weight_offset = torch.nn.Parameter( - layer.w13_weight_offset.data.transpose(-1, -2).contiguous(), - requires_grad=False, - ) - layer.w2_weight_offset = torch.nn.Parameter( - layer.w2_weight_offset.data.transpose(-1, -2).contiguous(), - requires_grad=False, - ) - - # w = [n, k // 8] --> [k, n // 8] - # w13_weight = layer.w13_weight.data.transpose(1, 2).contiguous() - # w2_weight = layer.w2_weight.data.transpose(1, 2).contiguous() - unpacked_w13_weight = ( - self._unpack_from_int32(layer.w13_weight.data.flatten(0, 1), 4) - .view(layer.w13_weight.data.shape[0], layer.w13_weight.data.shape[1], -1) - .transpose(1, 2) - .contiguous() - .int() - ) - unpacked_w2_weight = ( - self._unpack_from_int32(layer.w2_weight.data.flatten(0, 1), 4) - .view(layer.w2_weight.data.shape[0], layer.w2_weight.data.shape[1], -1) - .transpose(1, 2) - .contiguous() - .int() - ) - - w13_weight = self._pack_to_int32(unpacked_w13_weight) - w2_weight = self._pack_to_int32(unpacked_w2_weight) - - layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) - layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) - - if hasattr(layer, "dispatcher"): - layer.dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"}) - - def apply( - self, - layer, - dispatch_output: "DispatchOutput", - ) -> "CombineInput": - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - combine_input = self._maybe_apply_deepep(layer, dispatch_output) - if combine_input is not None: - return combine_input - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - - topk_weights, topk_ids, _ = topk_output - topk_ids = topk_ids.to(torch.int32) - topk_weights = topk_weights.to(x.dtype) - output = npu_fused_experts( - hidden_states=x, - w13=layer.w13_weight, - w13_scale=layer.w13_weight_scale, - w13_offset=layer.w13_weight_offset, - w2=layer.w2_weight, - w2_scale=layer.w2_weight_scale, - w2_offset=layer.w2_weight_offset, - topk_weights=topk_weights, - topk_ids=topk_ids, - top_k=topk_ids.shape[1], - use_wna16=True, - ) - return StandardCombineInput(hidden_states=output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - if hidden_states_scale is None: - # gmm1: gate_up_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w13_weight], - antiquant_scale=[layer.w13_weight_scale], - antiquant_offset=[layer.w13_weight_offset], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - - # act_fn: swiglu - hidden_states = torch.ops.npu.npu_swiglu(hidden_states) - - # gmm2: down_proj - out_hidden = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w2_weight], - antiquant_scale=[layer.w2_weight_scale], - antiquant_offset=[layer.w2_weight_offset], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - else: - raise ValueError( - "when weight is int4, hidden_states only supports non-quant dtype!" - ) - - return out_hidden diff --git a/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py b/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py index b3ed395d0..3a4cb18fa 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py @@ -3,17 +3,14 @@ from __future__ import annotations from typing import TYPE_CHECKING, Optional import torch -import torch_npu - -from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( - npu_fused_experts, -) if TYPE_CHECKING: - from sglang.srt.layers.moe import MoeRunnerConfig - from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput from sglang.srt.layers.quantization.base_config import QuantizationConfig +import logging + +logger = logging.getLogger(__name__) + def unpack_from_int32( weight: torch.Tensor, @@ -90,7 +87,7 @@ class GPTQLinearAscendKernel: # for 4bit case we need to pack 4bit weight to int32 to save memory layer.qweight = torch.nn.Parameter( - torch_npu.npu_convert_weight_to_int4pack(qweight_tmp.to(torch.int32)), + torch.ops.npu.npu_convert_weight_to_int4pack(qweight_tmp.to(torch.int32)), requires_grad=False, ) @@ -115,7 +112,7 @@ class GPTQLinearAscendKernel: else: out_shape = x.shape[:-1] + (qweight.shape[-1],) - out = torch_npu.npu_weight_quant_batchmatmul( + out = torch.ops.npu.npu_weight_quant_batchmatmul( reshaped_x, qweight, antiquant_scale=scales, @@ -131,17 +128,9 @@ class GPTQMoEAscendKernel: def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config self.use_v2_format = quant_config.checkpoint_format == "gptq_v2" - self.moe_runner_config: Optional[MoeRunnerConfig] = None - - def create_moe_runner( - self, - layer: torch.nn.Module, - moe_runner_config: MoeRunnerConfig, - **extra_weight_attrs, - ): - self.moe_runner_config = moe_runner_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # ----- zero‑points (unchanged) ----- w13_qzeros_2d = layer.w13_qzeros.data.contiguous().reshape( -1, layer.w13_qzeros.shape[-1] ) @@ -174,6 +163,7 @@ class GPTQMoEAscendKernel: if not self.use_v2_format: layer.w2_qzeros += 1 + # ----- w13 ----- w13_qweight_2d = ( layer.w13_qweight.data.transpose(-1, -2) .contiguous() @@ -185,35 +175,61 @@ class GPTQMoEAscendKernel: if self.quant_config.weight_bits == 4: group_size = self.quant_config.group_size - scale_expanded = layer.w13_scales.data.repeat_interleave(group_size, dim=1) + k_shard_w13 = w13_qweight_tmp.shape[1] - neg_mask = scale_expanded < 0 + # Check if the scales are compatible (expanded size must equal K_shard) + if layer.w13_scales.shape[1] * group_size != k_shard_w13: + logger.warning_once( + f"w13 scales expanded size {layer.w13_scales.shape[1] * group_size} " + f"does not match K_shard {k_shard_w13}. Skipping negative-scale correction." + f"This may break the accuracy, please try another TP-size or use DeepEP." + ) + # pack directly + layer.w13_qweight = torch.nn.Parameter( + torch.ops.npu.npu_convert_weight_to_int4pack( + w13_qweight_tmp.reshape( + layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1 + ) + .transpose(-1, -2) + .contiguous() + .reshape(-1, layer.w13_qweight.shape[2]) + .to(torch.int32) + ) + .reshape( + layer.w13_qweight.shape[0], layer.w13_qweight.shape[1] * 8, -1 + ) + .contiguous(), + requires_grad=False, + ) + else: + scale_expanded = layer.w13_scales.data.repeat_interleave( + group_size, dim=1 + ) + neg_mask = scale_expanded < 0 + if neg_mask.any(): + neg_mask = neg_mask.transpose(-1, -2) + neg_mask = neg_mask.contiguous().reshape(w13_qweight_tmp.shape) + w13_qweight_tmp[neg_mask] = -w13_qweight_tmp[neg_mask] + if w13_qweight_tmp.max() > 7: + w13_qweight_tmp.clamp_(max=7) + layer.w13_scales.data.abs_() - if neg_mask.any(): - neg_mask = neg_mask.transpose(-1, -2) - neg_mask = neg_mask.contiguous().reshape(w13_qweight_tmp.shape) - w13_qweight_tmp[neg_mask] = -w13_qweight_tmp[neg_mask] - - if w13_qweight_tmp.max() > 7: - w13_qweight_tmp.clamp_(max=7) - - layer.w13_scales.data.abs_() - - layer.w13_qweight = torch.nn.Parameter( - torch_npu.npu_convert_weight_to_int4pack( - w13_qweight_tmp.reshape( - layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1 + layer.w13_qweight = torch.nn.Parameter( + torch.ops.npu.npu_convert_weight_to_int4pack( + w13_qweight_tmp.reshape( + layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1 + ) + .transpose(-1, -2) + .contiguous() + .reshape(-1, layer.w13_qweight.shape[2]) + .to(torch.int32) + ) + .reshape( + layer.w13_qweight.shape[0], layer.w13_qweight.shape[1] * 8, -1 ) - .transpose(-1, -2) - .contiguous() - .reshape(-1, layer.w13_qweight.shape[2]) - .to(torch.int32) + .contiguous(), + requires_grad=False, ) - .reshape(layer.w13_qweight.shape[0], layer.w13_qweight.shape[1] * 8, -1) - .contiguous(), - requires_grad=False, - ) - # use int8 to store weight by default else: layer.w13_qweight = torch.nn.Parameter( w13_qweight_tmp.reshape( @@ -224,6 +240,7 @@ class GPTQMoEAscendKernel: requires_grad=False, ) + # ----- w2 ----- w2_qweight_2d = ( layer.w2_qweight.data.transpose(-1, -2) .contiguous() @@ -235,35 +252,61 @@ class GPTQMoEAscendKernel: if self.quant_config.weight_bits == 4: group_size = self.quant_config.group_size - scale_expanded = layer.w2_scales.data.repeat_interleave(group_size, dim=1) + k_shard_w2 = w2_qweight_tmp.shape[1] - neg_mask = scale_expanded < 0 + # Check if the scales are compatible + if layer.w2_scales.shape[1] * group_size != k_shard_w2: + logger.warning_once( + f"w2 scales expanded size {layer.w2_scales.shape[1] * group_size} " + f"does not match K_shard {k_shard_w2}. Skipping negative-scale correction." + f"This may break the accuracy, please try another TP-size or use DeepEP." + ) + # pack directly + layer.w2_qweight = torch.nn.Parameter( + torch.ops.npu.npu_convert_weight_to_int4pack( + w2_qweight_tmp.reshape( + layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1 + ) + .transpose(-1, -2) + .contiguous() + .reshape(-1, layer.w2_qweight.shape[2]) + .to(torch.int32) + ) + .reshape( + layer.w2_qweight.shape[0], layer.w2_qweight.shape[1] * 8, -1 + ) + .contiguous(), + requires_grad=False, + ) + else: + scale_expanded = layer.w2_scales.data.repeat_interleave( + group_size, dim=1 + ) + neg_mask = scale_expanded < 0 + if neg_mask.any(): + neg_mask = neg_mask.transpose(-1, -2) + neg_mask = neg_mask.contiguous().reshape(w2_qweight_tmp.shape) + w2_qweight_tmp[neg_mask] = -w2_qweight_tmp[neg_mask] + if w2_qweight_tmp.max() > 7: + w2_qweight_tmp.clamp_(max=7) + layer.w2_scales.data.abs_() - if neg_mask.any(): - neg_mask = neg_mask.transpose(-1, -2) - neg_mask = neg_mask.contiguous().reshape(w2_qweight_tmp.shape) - w2_qweight_tmp[neg_mask] = -w2_qweight_tmp[neg_mask] - - if w2_qweight_tmp.max() > 7: - w2_qweight_tmp.clamp_(max=7) - - layer.w2_scales.data.abs_() - - layer.w2_qweight = torch.nn.Parameter( - torch_npu.npu_convert_weight_to_int4pack( - w2_qweight_tmp.reshape( - layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1 + layer.w2_qweight = torch.nn.Parameter( + torch.ops.npu.npu_convert_weight_to_int4pack( + w2_qweight_tmp.reshape( + layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1 + ) + .transpose(-1, -2) + .contiguous() + .reshape(-1, layer.w2_qweight.shape[2]) + .to(torch.int32) + ) + .reshape( + layer.w2_qweight.shape[0], layer.w2_qweight.shape[1] * 8, -1 ) - .transpose(-1, -2) - .contiguous() - .reshape(-1, layer.w2_qweight.shape[2]) - .to(torch.int32) + .contiguous(), + requires_grad=False, ) - .reshape(layer.w2_qweight.shape[0], layer.w2_qweight.shape[1] * 8, -1) - .contiguous(), - requires_grad=False, - ) - # use int8 to store weight by default else: layer.w2_qweight = torch.nn.Parameter( w2_qweight_tmp.reshape( @@ -273,43 +316,3 @@ class GPTQMoEAscendKernel: .contiguous(), requires_grad=False, ) - - def apply( - self, - layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, - ) -> torch.Tensor: - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - assert self.moe_runner_config is not None, ( - "moe_runner_config is not set. " - "Did you forget to call create_weights/create_moe_runner?" - ) - - assert self.moe_runner_config.activation in ("silu", "swiglu"), ( - f"Only SiLU/Swiglu activation is supported, " - f"got {self.moe_runner_config.activation!r}." - ) - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - topk_weights, topk_ids, _ = topk_output - - topk_ids = topk_ids.to(torch.int32) - topk_weights = topk_weights.to(x.dtype) - - output = npu_fused_experts( - hidden_states=x, - w13=layer.w13_qweight, - w13_scale=layer.w13_scales, - w13_offset=layer.w13_qzeros, - w2=layer.w2_qweight, - w2_scale=layer.w2_scales, - w2_offset=layer.w2_qzeros, - topk_weights=topk_weights, - topk_ids=topk_ids, - top_k=topk_ids.shape[1], - use_wna16=True, - ) - - return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py index 2bd602f1d..ffd25114d 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py @@ -10,6 +10,8 @@ from sglang.srt.layers.quantization.base_config import LinearMethodBase if TYPE_CHECKING: from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.environ import envs + logger = logging.getLogger(__name__) MXFP8_BLOCK_SIZE = 32 @@ -313,9 +315,12 @@ class NPU_W4A4DynamicLinearMethod(_NPULinearMethodBase): layer.weight_scale.data = layer.weight_scale.data.flatten() layer.weight_scale_fp32 = layer.weight_scale.data.to(torch.float32) layer.weight_offset.data = layer.weight_offset.data.flatten() - layer.weight.data = torch.ops.npu.npu_convert_weight_to_int4pack( - layer.weight.data.to(torch.int32) - ) + if envs.SGLANG_NPU_W4A4_NEW_PACKING.get(): + layer.weight.data = layer.weight.data.view(torch.int32).contiguous() + else: + layer.weight.data = torch.ops.npu.npu_convert_weight_to_int4pack( + layer.weight.data.to(torch.int32) + ) def apply( self, diff --git a/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py b/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py new file mode 100644 index 000000000..f9b33e473 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py @@ -0,0 +1,724 @@ +from typing import TYPE_CHECKING, Any, Dict, Optional + +import numpy as np +import torch + +from sglang.srt.environ import envs +from sglang.srt.hardware_backend.npu.utils import npu_format_cast +from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase + +if TYPE_CHECKING: + from sglang.srt.layers.quantization.base_config import QuantizationConfig + from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo + +import logging + +from sglang.srt.hardware_backend.npu.moe.hidden_states_quant import ( + HiddenStatesDynamicQuant, +) +from sglang.srt.hardware_backend.npu.moe.matmul import GroupedMatmul + +logger = logging.getLogger(__name__) + + +# DEPRECATED METHOD +# TODO: Remove in future realeses +def fused_moe_npu( + x, + w1, + w2, + topk_output, + moe_runner_config, +): + logger.warning_once( + f"The fused_moe_npu method deprecated and will be removed in future releases" + ) + topk_weights, topk_ids, _ = topk_output + original_dtype = x.dtype + num_tokens = x.shape[0] + topk_weights = topk_weights.to(x.dtype) + topk_ids = topk_ids.to(torch.int32) + num_experts = w1.shape[0] + top_k = topk_weights.shape[-1] + row_idx_len = num_tokens * top_k + row_idx = ( + torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device) + .view(top_k, -1) + .permute(1, 0) + .contiguous() + ) + + hidden_states, expanded_row_idx, expanded_expert_idx = ( + torch.ops.npu.npu_moe_init_routing( + x, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens + ) + ) + + expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens( + expanded_expert_idx, num_experts + ) + + expert_tokens = expert_tokens.to(torch.int64) + + # gmm1: gate_up_proj + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[w1.permute(0, 2, 1)], + bias=None, + split_item=2, + group_list_type=0, + group_type=0, + group_list=expert_tokens, + output_dtype=original_dtype, + )[0] + + # act_fn: + if moe_runner_config.activation == "silu": + hidden_states = torch.ops.npu.npu_swiglu(hidden_states) + else: + from sglang.srt.layers.activation import GeluAndMul + + hidden_states = GeluAndMul()(hidden_states) + + # gmm2: down_proj + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[w2.permute(0, 2, 1)], + bias=None, + split_item=2, + group_list_type=0, + group_type=0, + group_list=expert_tokens, + output_dtype=original_dtype, + )[0] + + final_hidden_states = torch.ops.npu.npu_moe_finalize_routing( + hidden_states, + skip1=None, + skip2=None, + bias=None, + scales=topk_weights, + expanded_src_to_dst_row=expanded_row_idx, + export_for_source_row=topk_ids, + ) + return final_hidden_states + + +class _NPUMoEMethodBase(FusedMoEMethodBase): + """Base class for NPU MoE methods with common helpers.""" + + def __init__( + self, + quant_config: Optional["QuantizationConfig"] = None, + ): + super().__init__() + self.quant_config = quant_config + + @staticmethod + def _set_dispatcher_output_dtype(layer: torch.nn.Module, dtype) -> None: + """Set dispatcher output dtype if the layer has a dispatcher.""" + if hasattr(layer, "dispatcher"): + layer.dispatcher.set_quant_config({"dispatcher_output_dtype": dtype}) + + @staticmethod + def _validate_weight_prefix(layer: torch.nn.Module, weight_prefix: str) -> None: + """Ensure the required attributes exist on the layer for the given prefix.""" + required = [f"{weight_prefix}_weight"] + for attr in required: + if not hasattr(layer, attr): + raise AttributeError( + f"Layer {layer} is missing required attribute '{attr}' for " + f"weight_prefix='{weight_prefix}'" + ) + + @staticmethod + def _get_bias_args( + quant_info: "AscendQuantInfo", weight_prefix: str + ) -> Dict[str, Any]: + bias = getattr(quant_info, f"{weight_prefix}_scale_bias", None) + if bias is None: + bias = getattr(quant_info, f"{weight_prefix}_weight_bias", None) + return {"bias": [bias]} if bias is not None else {} + + +# --------------------------------------------------------------------------- +# NPUW4A4Int4DynamicMoEMethod +# --------------------------------------------------------------------------- +class NPUW4A4Int4MoEMethod(_NPUMoEMethodBase): + """W4A4 dynamic MoE – weights are int4, activations are int4.""" + + def __init__(self): + super().__init__(quant_config=None) + self.matmul = GroupedMatmul() + self.hidden_states_quantizer = HiddenStatesDynamicQuant( + quant_dtype=torch.quint4x2 + ) + + def process_weights_after_loading( + self, layer: torch.nn.Module, weight_prefix: str + ) -> None: + self._validate_weight_prefix(layer, weight_prefix) + + # Process scale + scale: torch.Tensor = getattr(layer, f"{weight_prefix}_weight_scale") + scale_np = scale.data.cpu().contiguous().numpy() + scale_np.dtype = np.uint32 + scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu() + processed_scale = torch.nn.Parameter( + scale_uint64_tensor.squeeze(-1), requires_grad=False + ) + setattr(layer, f"{weight_prefix}_weight_scale", processed_scale) + + # Process offset + offset: Optional[torch.Tensor] = getattr( + layer, f"{weight_prefix}_weight_offset", None + ) + if offset is not None: + processed_offset = torch.nn.Parameter( + offset.data.squeeze(-1), requires_grad=False + ) + setattr(layer, f"{weight_prefix}_weight_offset", processed_offset) + + # Process weight + weight: torch.Tensor = getattr(layer, f"{weight_prefix}_weight") + if not envs.SGLANG_NPU_W4A4_NEW_PACKING.get(): + weight.data = self._w4a4_pack_int4(weight.data) + weight.data = weight.data.transpose(-2, -1).contiguous() + weight.data = npu_format_cast(weight.data) + weight.data = self._pack_to_int32(weight.data) + + # Set DeepEP dispatcher output dtype + if weight_prefix == "w13": + self._set_dispatcher_output_dtype(layer, "bf16") + + def _pack_int4(self, weight) -> torch.Tensor: + """ + Pack int4 weight to int8 weight + @param weight: torch.Tensor, int4 weight + @return: torch.Tensor, int8 weight + """ + weight = weight.to(torch.int8) + e = 0 # number of experts + if len(weight.shape) == 2: + k, n = weight.shape + elif len(weight.shape) == 3: + e, k, n = weight.shape + n_new = n // 2 + n % 2 + + if n_new != n // 2: + raise AssertionError("n dimension should be even") + + weight = weight.reshape(-1, 2) + weight0 = weight[:, :1] + weight1 = weight[:, 1:] + + weight1_4 = torch.bitwise_left_shift(weight1, 4) + weight2_4 = weight0 & 0b00001111 + + weight_add = torch.bitwise_or(weight1_4, weight2_4) + if e == 0: + weight_res = weight_add.reshape(k, n_new) + else: + weight_res = weight_add.reshape(e, k, n_new) + return weight_res + + def _w4a4_pack_int4(self, save_quant_weight): + """ + Pack int4 weight to int8 weight + @param save_quant_weight: torch.Tensor, int4 weight + @return: torch.Tensor, int8 weight + """ + weight = save_quant_weight.transpose(-1, -2).contiguous() + packed_weight_tensor = self._pack_int4(weight) + packed_weight_tensor = packed_weight_tensor.transpose(-1, -2).contiguous() + return packed_weight_tensor + + def _pack_to_int32(self, weight: torch.Tensor): + # pack 4 int8(int4*2) to int32 + return weight.contiguous().view(torch.int32) + + def apply( + self, + quant_info: "AscendQuantInfo", + hidden_states: torch.Tensor, + expert_tokens: torch.Tensor, + pertoken_scale: torch.Tensor, + output_dtype: torch.dtype, + weight_prefix: str, + group_list_type, + ) -> torch.Tensor: + scale = getattr(quant_info, f"{weight_prefix}_weight_scale", None) + if pertoken_scale is None: + hidden_states, pertoken_scale = self.hidden_states_quantizer.__call__( + hidden_states + ) + scale_args: Dict[str, Any] = { + "scale": [scale], + "per_token_scale": [pertoken_scale], + } + scale_args.update(self._get_bias_args(quant_info, weight_prefix)) + return self.matmul.forward( + quant_info, + weight_prefix, + hidden_states, + expert_tokens, + output_dtype, + group_list_type=group_list_type, + transposed=True, + **scale_args, + ) + + +# --------------------------------------------------------------------------- +# NPUW8A8Int8MoEMethod +# --------------------------------------------------------------------------- +class NPUW8A8Int8MoEMethod(_NPUMoEMethodBase): + """W8A8 MoE – weights are int8, activations in int8.""" + + def __init__(self): + super().__init__(quant_config=None) + self.matmul = GroupedMatmul() + self.hidden_states_quantizer = HiddenStatesDynamicQuant(quant_dtype=torch.int8) + + @staticmethod + def maybe_process_fuseep_weights(layer: torch.nn.Module) -> bool: + """Apply the FuseEP weight layout if --moe-a2a-backend is ascend_fuseep. + + Returns True when the FuseEP layout was (or has already been) applied, + so that the caller can skip its own ``process_weights_after_loading`` body. + """ + from sglang.srt.layers.moe import get_moe_a2a_backend + + if not get_moe_a2a_backend().is_ascend_fuseep(): + return False + + # Guard against double processing when called for multiple prefixes. + if getattr(layer, "_fuseep_weights_processed", False): + return True + + from sglang.srt.hardware_backend.npu.moe.fuseep import process_fuseep_weights + + for prefix in ("w13", "w2"): + process_fuseep_weights(layer, prefix) + layer._fuseep_weights_processed = True + return True + + def process_weights_after_loading( + self, layer: torch.nn.Module, weight_prefix: str + ) -> None: + # If the FuseEP weight layout is used, process weights via + # maybe_apply_fuseep_weights and skip the rest of this method. + if self.maybe_process_fuseep_weights(layer): + return + + self._validate_weight_prefix(layer, weight_prefix) + + # Process scale + scale: torch.Tensor = getattr(layer, f"{weight_prefix}_weight_scale") + processed_scale = torch.nn.Parameter( + scale.data.squeeze(-1).to(dtype=torch.bfloat16), requires_grad=False + ) + setattr(layer, f"{weight_prefix}_weight_scale", processed_scale) + + # Process offset + offset: Optional[torch.Tensor] = getattr( + layer, f"{weight_prefix}_weight_offset", None + ) + if offset is not None: + processed_offset = torch.nn.Parameter( + offset.data.squeeze(-1), requires_grad=False + ) + setattr(layer, f"{weight_prefix}_weight_offset", processed_offset) + + # Process weight + weight: torch.Tensor = getattr(layer, f"{weight_prefix}_weight") + weight.data = npu_format_cast(weight.data.transpose(1, 2)) + setattr( + layer, + f"{weight_prefix}_weight", + torch.nn.Parameter(weight, requires_grad=False), + ) + + # Set dispatcher output dtype + if weight_prefix == "w13": + self._set_dispatcher_output_dtype(layer, "int8") + + def apply( + self, + quant_info: "AscendQuantInfo", + hidden_states: torch.Tensor, + expert_tokens: torch.Tensor, + pertoken_scale: torch.Tensor, + output_dtype: torch.dtype, + weight_prefix: str, + group_list_type, + ) -> torch.Tensor: + scale = getattr(quant_info, f"{weight_prefix}_weight_scale", None) + if pertoken_scale is None: + hidden_states, pertoken_scale = self.hidden_states_quantizer.__call__( + hidden_states + ) + scale_args: Dict[str, Any] = { + "scale": [scale], + "per_token_scale": [pertoken_scale], + } + scale_args.update(self._get_bias_args(quant_info, weight_prefix)) + return self.matmul.forward( + quant_info, + weight_prefix, + hidden_states, + expert_tokens, + output_dtype, + group_list_type=group_list_type, + transposed=True, + **scale_args, + ) + + +# --------------------------------------------------------------------------- +# NPUW4A8Int8MoEMethod +# --------------------------------------------------------------------------- +class NPUW4A8Int8MoEMethod(_NPUMoEMethodBase): + """W4A8 MoE – weights are int4, activations quantized to int8.""" + + def __init__( + self, + quant_config: Optional["QuantizationConfig"] = None, + is_per_channel_weight: bool = False, + activation_use_clip: bool = False, + ): + super().__init__(quant_config) + self.is_per_channel_weight = is_per_channel_weight + self.activation_use_clip = activation_use_clip + self.matmul = GroupedMatmul() + self.hidden_states_quantizer = HiddenStatesDynamicQuant(quant_dtype=torch.int8) + + def process_weights_after_loading( + self, layer: torch.nn.Module, weight_prefix: str + ) -> None: + self._validate_weight_prefix(layer, weight_prefix) + + # Process scale (and bias if needed) + scale = getattr(layer, f"{weight_prefix}_weight_scale") + scale_second = getattr(layer, f"{weight_prefix}_weight_scale_second", None) + bias = getattr(layer, f"{weight_prefix}_bias", None) + + if not self.activation_use_clip: + # Process scale according to per-channel or per-group + processed_scale = self._process_scale( + getattr(layer, f"{weight_prefix}_weight"), + scale, + scale_second, + self.is_per_channel_weight, + ) + setattr( + layer, + f"{weight_prefix}_weight_scale", + torch.nn.Parameter(processed_scale.squeeze(-1), requires_grad=False), + ) + if scale_second is not None: + delattr(layer, f"{weight_prefix}_weight_scale_second") + delattr(layer, f"{weight_prefix}_weight_offset_second") + else: + # With clip: simple squeeze + unsqueeze + processed_scale = scale.data.squeeze(-1).unsqueeze(1).contiguous() + setattr( + layer, + f"{weight_prefix}_weight_scale", + torch.nn.Parameter(processed_scale, requires_grad=False), + ) + if bias is not None: + setattr( + layer, + f"{weight_prefix}_scale_bias", + torch.nn.Parameter( + bias.data.transpose(1, 2).sum(dim=1).contiguous(), + requires_grad=False, + ), + ) + + # Process weight + weight = getattr(layer, f"{weight_prefix}_weight") + weight.data = npu_format_cast(weight.data.transpose(1, 2)) + weight.data = self._pack_to_int32(weight.data) + setattr( + layer, + f"{weight_prefix}_weight", + torch.nn.Parameter(weight, requires_grad=False), + ) + + # Set dispatcher output dtype + if weight_prefix == "w13": + self._set_dispatcher_output_dtype(layer, "int8") + + def _process_scale( + self, + weight: torch.Tensor, + scale: torch.Tensor, + per_group_scale: Optional[torch.Tensor], + is_per_channel: bool, + ) -> torch.Tensor: + scale = scale.transpose(1, 2).contiguous() + if is_per_channel: + scale_np = scale.cpu().contiguous().numpy() + scale_np.dtype = np.uint32 + scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu() + return scale_uint64_tensor + + # Per‑group: multiply channel and group scales, then pack into uint64 + per_group_scale = per_group_scale.transpose(1, 2).contiguous() + group_num, k, n = weight.shape + n = n * 2 # packed weight halves the column dimension + per_group_scale = per_group_scale.reshape(group_num, -1, n) + group_num, quantgroup_num, n = per_group_scale.shape + + scale_fp32 = (scale * per_group_scale).to(torch.float16).to(torch.float32) + scale_fp32_np = scale_fp32.cpu().numpy() + scale_fp32_np.dtype = np.uint32 + sscale_uint64 = np.zeros((group_num, quantgroup_num, n * 2), dtype=np.uint32) + sscale_uint64[..., ::2] = scale_fp32_np + sscale_uint64_tensor = ( + torch.from_numpy(sscale_uint64.view(np.int64).copy()) + .reshape(group_num, quantgroup_num, n) + .npu() + ) + return sscale_uint64_tensor + + def _pack_to_int32(self, weight: torch.Tensor) -> torch.Tensor: + # pack 4 int8 (representing 8 int4) into int32 + assert weight.shape[-1] % 4 == 0, ( + f"Last dimension of weight must be divisible by 4 for int8→int32 packing, " + f"got shape {weight.shape}" + ) + return weight.contiguous().view(torch.int32) + + def apply( + self, + quant_info: "AscendQuantInfo", + hidden_states: torch.Tensor, + expert_tokens: torch.Tensor, + pertoken_scale: torch.Tensor, + output_dtype: torch.dtype, + weight_prefix: str, + group_list_type, + ) -> torch.Tensor: + scale = getattr(quant_info, f"{weight_prefix}_weight_scale", None) + if pertoken_scale is None: + hidden_states, pertoken_scale = self.hidden_states_quantizer.__call__( + hidden_states + ) + scale_args: Dict[str, Any] = { + "scale": [scale], + "per_token_scale": [pertoken_scale], + } + scale_args.update(self._get_bias_args(quant_info, weight_prefix)) + + return self.matmul.forward( + quant_info, + weight_prefix, + hidden_states, + expert_tokens, + output_dtype, + group_list_type=group_list_type, + transposed=True, + **scale_args, + ) + + +# --------------------------------------------------------------------------- +# NPUWNA16Int4MoEMethod +# --------------------------------------------------------------------------- +class NPUWNA16Int4MoEMethod(_NPUMoEMethodBase): + """W4A16 MoE – weights are int4, activations stay in BF16.""" + + def __init__(self): + super().__init__(quant_config=None) + self.matmul = GroupedMatmul() + + def process_weights_after_loading( + self, layer: torch.nn.Module, weight_prefix: str + ) -> None: + self._validate_weight_prefix(layer, weight_prefix) + + # Process scale + scale = getattr(layer, f"{weight_prefix}_weight_scale") # shape [E, N, 1] + scale = scale.data.transpose(-1, -2).contiguous() # [E, N, 1] -> [E, 1, N] + setattr( + layer, + f"{weight_prefix}_weight_scale", + torch.nn.Parameter(scale, requires_grad=False), + ) + + # Process offset + offset = getattr(layer, f"{weight_prefix}_weight_offset", None) + if offset is not None: + offset = offset.data.transpose(-1, -2).contiguous() + setattr( + layer, + f"{weight_prefix}_weight_offset", + torch.nn.Parameter(offset, requires_grad=False), + ) + + # Process weight: unpack, transpose, repack + weight: torch.Tensor = getattr(layer, f"{weight_prefix}_weight") + unpacked_weight = ( + self._unpack_from_int32(weight.data.flatten(0, 1), 4) + .view(weight.shape[0], weight.shape[1], -1) + .transpose(1, 2) + .int() + ) + weight.data = self._pack_to_int32(unpacked_weight) + setattr( + layer, + f"{weight_prefix}_weight", + torch.nn.Parameter(weight, requires_grad=False), + ) + + # Set dispatcher output dtype + if weight_prefix == "w13": + self._set_dispatcher_output_dtype(layer, "bf16") + + def _pack_to_int32(self, weight: torch.Tensor) -> torch.Tensor: + assert weight.dim() == 3 + if weight.dtype == torch.int32: + assert weight.shape[-1] % 8 == 0, ( + f"Last dimension of int32 weight must be divisible by 8 for int4 packing, " + f"got {weight.shape}" + ) + new_weight = torch.ops.npu.npu_convert_weight_to_int4pack( + weight.flatten(0, 1) + ) + new_weight = new_weight.view(weight.shape[0], weight.shape[1], -1) + elif weight.dtype == torch.int8: + assert weight.shape[-1] % 4 == 0, ( + f"Last dimension of int8 weight must be divisible by 4 for int32 packing, " + f"got {weight.shape}" + ) + new_weight = weight.contiguous().view(torch.int32) + else: + raise ValueError(f"Unsupported weight dtype for packing: {weight.dtype}") + return new_weight.contiguous() + + def _unpack_from_int32( + self, + value: torch.Tensor, + num_bits: int, + shape: Optional[torch.Size] = None, + packed_dim: int = 1, + ) -> torch.Tensor: + """ + Unpacks a tensor of packed int32 weights into individual int8s, + maintaining the original bit range. + """ + if value.dtype is not torch.int32: + raise ValueError( + f"Expected {torch.int32} but got {value.dtype}, Aborting unpack." + ) + if num_bits > 8: + raise ValueError("Unpacking is only supported for less than 8 bits") + + pack_factor = 32 // num_bits + mask = (1 << num_bits) - 1 + + if packed_dim == 1: + unpacked = torch.zeros( + (value.shape[0], value.shape[1] * pack_factor), + device=value.device, + dtype=torch.int32, + ) + for i in range(pack_factor): + unpacked[:, i::pack_factor] = (value >> (num_bits * i)) & mask + if shape is not None: + original_row_size = int(shape[1]) + unpacked = unpacked[:, :original_row_size] + else: + unpacked = torch.zeros( + (value.shape[0] * pack_factor, value.shape[1]), + device=value.device, + dtype=torch.int32, + ) + for i in range(pack_factor): + unpacked[i::pack_factor, :] = (value >> (num_bits * i)) & mask + if shape is not None: + original_row_size = int(shape[0]) + unpacked = unpacked[:original_row_size, :] + + offset = pow(2, num_bits) // 2 + unpacked = (unpacked - offset).to(torch.int8) + return unpacked + + def apply( + self, + quant_info: "AscendQuantInfo", + hidden_states: torch.Tensor, + expert_tokens: torch.Tensor, + pertoken_scale: torch.Tensor, # not used, but kept for interface consistency + output_dtype: torch.dtype, + weight_prefix: str, + group_list_type, + ) -> torch.Tensor: + scale = getattr(quant_info, f"{weight_prefix}_weight_scale", None) + offset = getattr(quant_info, f"{weight_prefix}_weight_offset", None) + scale_args: Dict[str, Any] = { + "antiquant_scale": [scale], + "antiquant_offset": [offset] if offset is not None else [], + } + scale_args.update(self._get_bias_args(quant_info, weight_prefix)) + return self.matmul.forward( + quant_info, + weight_prefix, + hidden_states, + expert_tokens, + output_dtype, + group_list_type=group_list_type, + transposed=True, + **scale_args, + ) + + +# --------------------------------------------------------------------------- +# NPUWUnquantMoEMethod +# --------------------------------------------------------------------------- +class NPUUnquantMoEMethod(_NPUMoEMethodBase): + """Unquant MoE – all computations in BF16, no quantization.""" + + def __init__(self): + super().__init__(quant_config=None) + self.matmul = GroupedMatmul() + + def process_weights_after_loading( + self, layer: torch.nn.Module, weight_prefix: str + ) -> None: + self._validate_weight_prefix(layer, weight_prefix) + + weight: torch.Tensor = getattr(layer, f"{weight_prefix}_weight") + weight.data = npu_format_cast(weight) + setattr( + layer, + f"{weight_prefix}_weight", + torch.nn.Parameter(weight, requires_grad=False), + ) + + if weight_prefix == "w13": + self._set_dispatcher_output_dtype(layer, "bf16") + + def apply( + self, + quant_info: "AscendQuantInfo", + hidden_states: torch.Tensor, + expert_tokens: torch.Tensor, + pertoken_scale: torch.Tensor, # ignored + output_dtype: torch.dtype, + weight_prefix: str, + group_list_type, + ) -> torch.Tensor: + return self.matmul.forward( + quant_info, + weight_prefix, + hidden_states, + expert_tokens, + output_dtype, + group_list_type=group_list_type, + transposed=False, + **self._get_bias_args(quant_info, weight_prefix), + ) diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py index 15d117878..54389f39a 100644 --- a/python/sglang/srt/hardware_backend/npu/utils.py +++ b/python/sglang/srt/hardware_backend/npu/utils.py @@ -24,11 +24,6 @@ class NPUACLFormat(IntEnum): ACL_FORMAT_FRACTAL_NZ = 29 -class FusedMoEMode(IntEnum): - FUSED_DEEP_MOE = 1 - DISPATCH_FFN_COMBINE = 2 - - def _call_once(fn: Callable): @functools.wraps(fn) diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 7c775e9b0..5ea9e93dd 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -24,6 +24,7 @@ from sglang.srt.distributed import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) +from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import ( is_allocation_symmetric, ) @@ -670,6 +671,14 @@ class MergedColumnParallelLinear(ColumnParallelLinear): shard_size = loaded_weight.shape[output_dim] shard_offset = loaded_weight.shape[output_dim] * loaded_shard_id + # Needed for experimental ModelSlim W4A4 int4x2 packing support + # TODO: remove env variable once new packing is fully released + if envs.SGLANG_NPU_W4A4_NEW_PACKING.get(): + pack_factor = getattr(param, "pack_factor", None) + if pack_factor is not None: + shard_size = shard_size // pack_factor + shard_offset = shard_offset // pack_factor + param_data = param_data.narrow(output_dim, shard_offset, shard_size) start_idx = self.tp_rank * shard_size diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index d1e5cc2bd..88028dd87 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -285,9 +285,4 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]): or get_moe_a2a_backend().is_nixl() ): return DeepEPMoE - if get_moe_a2a_backend().is_ascend_fuseep(): - # ascend_fuseep bypasses dispatch/combine inside FusedMoE.forward - # (see forward_fuseep in hardware_backend/npu/moe/fuseep.py). - return FusedMoE - return FusedMoE 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 c1b432a38..b326663e3 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -34,6 +34,9 @@ from sglang.srt.layers.moe.kt_ep_wrapper import ( create_kt_config_from_server_args, ) from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput +from sglang.srt.layers.moe.token_dispatcher.ascend_tp import ( + AscendTPDispatcher, +) from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher from sglang.srt.layers.moe.token_dispatcher.standard import ( @@ -100,7 +103,9 @@ def _get_deepep_comm_group(a2a_backend): def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher: a2a_backend = get_moe_a2a_backend() - if ( + if a2a_backend.is_none() and is_npu(): + return AscendTPDispatcher(moe_runner_config) + elif ( a2a_backend.is_none() or a2a_backend.is_megamoe() or a2a_backend.is_ascend_fuseep() diff --git a/python/sglang/srt/layers/moe/moe_runner/ascend.py b/python/sglang/srt/layers/moe/moe_runner/ascend.py new file mode 100644 index 000000000..7b15281c3 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/ascend.py @@ -0,0 +1,310 @@ +"""Ascend MoE runner backend with NPU‑specific ops.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Optional + +import torch + +from sglang.srt.hardware_backend.npu.moe.activation import ( + AllGatherActivationWrapper, + NPUGeluAndMul, + NPUSwiglu, + NPUSwigluDeepEPKernel, + NPUSwigluOAI, + NPUSwigluQuant, + NPUSwigluStepAndMul, +) +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUW4A8Int8MoEMethod, + NPUW8A8Int8MoEMethod, +) +from sglang.srt.layers.moe.moe_runner.base import ( + MoeQuantInfo, + MoeRunnerConfig, + MoeRunnerCore, + RunnerInput, + RunnerOutput, + register_post_permute, + register_pre_permute, +) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher.deepep import ( + DeepEPLLCombineInput, + DeepEPLLDispatchOutput, + DeepEPNormalCombineInput, + DeepEPNormalDispatchOutput, + ) + from sglang.srt.layers.moe.token_dispatcher.ascend_tp import ( + AscendTPDispatchOutput, + AscendTPCombineInput, + ) + +from sglang.srt.layers.moe.utils import ( + MoeRunnerBackend, + get_moe_a2a_backend, +) + + +# --------------------------------------------------------------------------- +# Runner IO dataclasses +# --------------------------------------------------------------------------- +@dataclass +class AscendRunnerInput(RunnerInput): + """Input bundle for the NPU runner.""" + + hidden_states: torch.Tensor + hidden_states_scale: Optional[torch.Tensor] # None for unquant + expert_tokens: torch.Tensor + group_list_type: int # 0 or 1 (passed to NPU ops) + + @property + def runner_backend(self) -> MoeRunnerBackend: + return MoeRunnerBackend.ASCEND + + +@dataclass +class AscendRunnerOutput(RunnerOutput): + """Output bundle from the NPU runner.""" + + hidden_states: torch.Tensor + + @property + def runner_backend(self) -> MoeRunnerBackend: + return MoeRunnerBackend.ASCEND + + +# --------------------------------------------------------------------------- +# Main runner core +# --------------------------------------------------------------------------- +class AscendRunnerCore(MoeRunnerCore): + runner_backend = MoeRunnerBackend.ASCEND + + def __init__(self, config: MoeRunnerConfig): + super().__init__(config) + + kernel = config.layer.w2_kernel + + if get_moe_a2a_backend().is_deepep(): + # DeepEP path: use a unified kernel that decides quantisation + is_quant_kernel = isinstance( + kernel, (NPUW4A8Int8MoEMethod, NPUW8A8Int8MoEMethod) + ) + self.activation = NPUSwigluDeepEPKernel(need_quant=is_quant_kernel) + else: + # Non‑DeepEP (ascend_tp) path + # 1. Choose the base activation according to the quant method + if isinstance(kernel, (NPUW4A8Int8MoEMethod, NPUW8A8Int8MoEMethod)): + inner = NPUSwigluQuant() + else: + if config.activation == "npu_swiglu_oai": + # NPUSwigluOAI requires the runner config to pass + # gemm1_alpha and gemm1_clamp_limit to the triton kernel. + inner = NPUSwigluOAI(moe_runner_config=config) + elif config.activation == "silu": + if config.gemm1_clamp_limit is not None: + inner = NPUSwigluStepAndMul( + clamp_limit=config.gemm1_clamp_limit + ) + else: + inner = NPUSwiglu() + else: + inner = NPUGeluAndMul() + + # 2. If the quant method (GGUF) needs TP all‑gather, wrap the activation + if getattr(config, "use_tp_all_gather_activation", False): + self.activation = AllGatherActivationWrapper(inner, dim=-1) + else: + self.activation = inner + + def run( + self, + runner_input: AscendRunnerInput, + quant_info: AscendQuantInfo, + running_state: dict, + hooks: Optional[Any] = None, + ) -> AscendRunnerOutput: + """ + Execute the MoE layer using NPU‑specific grouped matmul ops. + """ + x = runner_input.hidden_states + original_dtype = torch.float16 if x.dtype == torch.float16 else torch.bfloat16 + expert_tokens = runner_input.expert_tokens + group_list_type = runner_input.group_list_type + + # --- w13 (gate & up) projection --- + hidden_states = self.config.layer.w13_kernel.apply( + quant_info, + x, + expert_tokens, + pertoken_scale=runner_input.hidden_states_scale, + output_dtype=original_dtype, + weight_prefix="w13", + group_list_type=group_list_type, + ) + + # --- Activation --- + # The DeepEP kernel expects extra dispatch metadata + if isinstance(self.activation, NPUSwigluDeepEPKernel): + hidden_states, pertoken_scale = self.activation._apply_activation( + hidden_states, + group_list=expert_tokens, + group_list_type=group_list_type, + ) + else: + hidden_states, pertoken_scale = self.activation._apply_activation( + hidden_states + ) + + # --- w2 (down) projection --- + hidden_states = self.config.layer.w2_kernel.apply( + quant_info, + hidden_states, + expert_tokens, + pertoken_scale=pertoken_scale, + output_dtype=original_dtype, + weight_prefix="w2", + group_list_type=group_list_type, + ) + return AscendRunnerOutput(hidden_states=hidden_states) + + +# --------------------------------------------------------------------------- +# QuantInfo +# --------------------------------------------------------------------------- +@dataclass +class AscendQuantInfo(MoeQuantInfo): + """Quantization payload for Ascend.""" + + w13_weight: torch.Tensor + w2_weight: torch.Tensor + w13_weight_scale: Optional[torch.Tensor] = None + w2_weight_scale: Optional[torch.Tensor] = None + w13_weight_offset: Optional[torch.Tensor] = None + w2_weight_offset: Optional[torch.Tensor] = None + w13_weight_bias: Optional[torch.Tensor] = None + w2_weight_bias: Optional[torch.Tensor] = None + w13_scale_bias: Optional[torch.Tensor] = None + w2_scale_bias: Optional[torch.Tensor] = None + + +# --------------------------------------------------------------------------- +# Pre/Post permute hooks +# --------------------------------------------------------------------------- + + +@register_pre_permute("ascend_tp", "ascend") +def pre_permute_ascend_tp_to_ascend( + dispatch_output: AscendTPDispatchOutput, + quant_info: AscendQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> AscendRunnerInput: + return AscendRunnerInput( + hidden_states=dispatch_output.hidden_states, + hidden_states_scale=dispatch_output.hidden_states_scale, + expert_tokens=dispatch_output.expert_tokens, + group_list_type=dispatch_output.group_list_type, + ) + + +@register_pre_permute("deepep_normal", "ascend") +def pre_permute_deepep_normal_to_ascend( + dispatch_output: DeepEPNormalDispatchOutput, + quant_info: AscendQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> AscendRunnerInput: + ( + hidden_states, + hidden_states_scale, + topk_ids, + topk_weights, + num_recv_tokens_per_expert, + ) = dispatch_output + group_list = torch.tensor( + num_recv_tokens_per_expert, + dtype=torch.int64, + device=hidden_states.device, + ) + running_state["topk_ids"] = topk_ids + running_state["topk_weights"] = topk_weights + + return AscendRunnerInput( + hidden_states=hidden_states, + hidden_states_scale=hidden_states_scale, + expert_tokens=group_list, + group_list_type=1, + ) + + +@register_pre_permute("deepep_ll", "ascend") +def pre_permute_deepep_ll_to_ascend( + dispatch_output: DeepEPLLDispatchOutput, + quant_info: AscendQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> AscendRunnerInput: + ( + hidden_states, + hidden_states_scale, + topk_ids, + topk_weights, + group_list, + _, + ) = dispatch_output + group_list = group_list.to(torch.int64) + running_state["topk_ids"] = topk_ids + running_state["topk_weights"] = topk_weights + return AscendRunnerInput( + hidden_states=hidden_states, + hidden_states_scale=hidden_states_scale, + expert_tokens=group_list, + group_list_type=1, + ) + + +@register_post_permute("ascend", "ascend_tp") +def post_permute_ascend_to_ascend_tp( + runner_output: AscendRunnerOutput, + quant_info: AscendQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> AscendTPCombineInput: + from sglang.srt.layers.moe.token_dispatcher.ascend_tp import AscendTPCombineInput + + return AscendTPCombineInput(hidden_states=runner_output.hidden_states) + + +@register_post_permute("ascend", "deepep_normal") +def post_permute_ascend_to_deepep_normal( + runner_output: AscendRunnerOutput, + quant_info: AscendQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> DeepEPNormalCombineInput: + from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPNormalCombineInput + + return DeepEPNormalCombineInput( + hidden_states=runner_output.hidden_states, + topk_ids=running_state["topk_ids"], + topk_weights=running_state["topk_weights"], + ) + + +@register_post_permute("ascend", "deepep_ll") +def post_permute_ascend_to_deepep_ll( + runner_output: AscendRunnerOutput, + quant_info: AscendQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> DeepEPLLCombineInput: + from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPLLCombineInput + + return DeepEPLLCombineInput( + hidden_states=runner_output.hidden_states, + topk_ids=running_state["topk_ids"], + topk_weights=running_state["topk_weights"], + ) diff --git a/python/sglang/srt/layers/moe/moe_runner/base.py b/python/sglang/srt/layers/moe/moe_runner/base.py index ef913a75f..49f3ed5c7 100644 --- a/python/sglang/srt/layers/moe/moe_runner/base.py +++ b/python/sglang/srt/layers/moe/moe_runner/base.py @@ -60,6 +60,8 @@ class MoeRunnerConfig: # silu+is_gated swiglu path consumes it (interleaved -> swiglu_gpt_oss_*, # otherwise chunk gate/up then apply alpha/limit). gate_up_interleaved: bool = True + layer: Optional[torch.nn.Module] = None + use_tp_all_gather_activation: bool = False @dataclass diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index c96b4674a..a02a46d79 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -39,6 +39,10 @@ class MoeRunner: if runner_backend.is_triton(): self.runner_core = TritonRunnerCore(config) + elif runner_backend.is_ascend(): + from sglang.srt.layers.moe.moe_runner.ascend import AscendRunnerCore + + self.runner_core = AscendRunnerCore(config) elif runner_backend.is_triton_kernels(): self.runner_core = TritonKernelsRunnerCore(config) elif runner_backend.is_deep_gemm(): diff --git a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py index f1ebac970..b605f74c0 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py @@ -1,3 +1,8 @@ +from sglang.srt.layers.moe.token_dispatcher.ascend_tp import ( + AscendTPCombineInput, + AscendTPDispatcher, + AscendTPDispatchOutput, +) from sglang.srt.layers.moe.token_dispatcher.base import ( BaseDispatcher, BaseDispatcherConfig, @@ -74,4 +79,7 @@ __all__ = [ "DeepEPLLDispatchOutput", "DeepEPLLCombineInput", "DeepEPNormalCombineInput", + "AscendTPDispatcher", + "AscendTPDispatchOutput", + "AscendTPCombineInput", ] diff --git a/python/sglang/srt/layers/moe/token_dispatcher/ascend_tp.py b/python/sglang/srt/layers/moe/token_dispatcher/ascend_tp.py new file mode 100644 index 000000000..d7c80e153 --- /dev/null +++ b/python/sglang/srt/layers/moe/token_dispatcher/ascend_tp.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +from typing import NamedTuple, Optional + +import torch + +from sglang.srt.hardware_backend.npu.moe.finalize_routing import ( + AllGatherFinalizeRoutingWrapper, + NPUFinalizeRouting, +) +from sglang.srt.hardware_backend.npu.moe.init_routing import ( + NPUMoEInitRouting_v2, +) +from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig +from sglang.srt.layers.moe.token_dispatcher.base import ( + BaseDispatcher, + CombineInputFormat, + DispatchOutputFormat, +) +from sglang.srt.layers.moe.topk import TopKOutput +from sglang.srt.layers.moe.utils import ( + DispatcherOutputDtype, + get_ascend_dispatcher_output_dtype, +) +from sglang.srt.runtime_context import get_parallel + + +class AscendTPDispatchOutput(NamedTuple): + hidden_states: torch.Tensor + hidden_states_scale: Optional[torch.Tensor] + topk_weights: torch.Tensor + topk_ids: torch.Tensor + expanded_row_idx: torch.Tensor + expert_tokens: torch.Tensor + group_list_type: int + + @property + def format(self) -> DispatchOutputFormat: + return DispatchOutputFormat.ASCEND_TP + + +class AscendTPCombineInput(NamedTuple): + hidden_states: torch.Tensor + + @property + def format(self) -> CombineInputFormat: + return CombineInputFormat.ASCEND_TP + + +class AscendTPDispatcher(BaseDispatcher): + def __init__(self, moe_runner_config: MoeRunnerConfig): + super().__init__() + self.num_experts = moe_runner_config.num_experts + self.top_k = moe_runner_config.top_k + self._dispatch_output: Optional[AscendTPDispatchOutput] = None + + self.quant_config: Optional[dict] = None + + # Initialise routing kernels with default (no quant config yet) + self.set_ascend_dispatcher_output_dtype() + + def set_quant_config(self, quant_config: dict) -> None: + self.quant_config = quant_config + self.set_ascend_dispatcher_output_dtype() + + # If the quantisation is GGUF and TP is active, wrap the finalizer + # with an all‑gather so that the dispatcher stays completely clean. + if ( + isinstance(self.quant_config, dict) + and self.quant_config.get("quant_type") == "gguf" + and get_parallel().tp_size > 1 + ): + self.finalize = AllGatherFinalizeRoutingWrapper(self.finalize, dim=-1) + + def set_ascend_dispatcher_output_dtype(self) -> None: + """Choose init & finalize routing kernels based on quant config.""" + self.ascend_dispatcher_output_dtype = get_ascend_dispatcher_output_dtype(self) + + if self.ascend_dispatcher_output_dtype == DispatcherOutputDtype.BF16: + self.init = NPUMoEInitRouting_v2(quant_mode=-1) + self.finalize = NPUFinalizeRouting(drop_pad_mode=2) + self.group_list_type = 1 + elif self.ascend_dispatcher_output_dtype == DispatcherOutputDtype.INT8: + self.init = NPUMoEInitRouting_v2(quant_mode=1) + self.finalize = NPUFinalizeRouting(drop_pad_mode=2) + self.group_list_type = 1 + else: + raise ValueError( + f"Unsupported ascend_dispatcher_output_dtype: {self.ascend_dispatcher_output_dtype}" + ) + + def dispatch( + self, hidden_states: torch.Tensor, topk_output: TopKOutput + ) -> AscendTPDispatchOutput: + topk_weights, topk_ids, _ = topk_output + topk_weights = topk_weights.to(hidden_states.dtype) + topk_ids = topk_ids.to(torch.int32) + + ( + permuted_hidden_states, + expanded_row_idx, + expert_tokens, + hidden_states_scale, + ) = self.init._init_routing( + hidden_states, + topk_ids, + self.num_experts, + self.top_k, + ) + + self._dispatch_output = AscendTPDispatchOutput( + hidden_states=permuted_hidden_states, + hidden_states_scale=hidden_states_scale, + topk_weights=topk_weights, + topk_ids=topk_ids, + expanded_row_idx=expanded_row_idx, + expert_tokens=expert_tokens, + group_list_type=self.group_list_type, + ) + return self._dispatch_output + + def combine(self, combine_input: AscendTPCombineInput) -> torch.Tensor: + if self._dispatch_output is None: + raise RuntimeError("combine() called before dispatch()") + + dispatch_out = self._dispatch_output + + # The finalizer (possibly wrapped with TP all‑gather) does all the work. + final_hidden_states = self.finalize._finalize_routing( + combine_input.hidden_states, + topk_weights=dispatch_out.topk_weights, + expanded_row_idx=dispatch_out.expanded_row_idx, + topk_ids=dispatch_out.topk_ids, + ) + + self._dispatch_output = None + return final_hidden_states diff --git a/python/sglang/srt/layers/moe/token_dispatcher/base.py b/python/sglang/srt/layers/moe/token_dispatcher/base.py index 3e62d9566..1ff2beb5b 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/base.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/base.py @@ -21,6 +21,8 @@ import torch if TYPE_CHECKING: from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs from sglang.srt.layers.moe.token_dispatcher import ( + AscendTPCombineInput, + AscendTPDispatchOutput, DeepEPLLCombineInput, DeepEPLLDispatchOutput, DeepEPNormalCombineInput, @@ -133,6 +135,12 @@ class DispatchOutputChecker: ) -> TypeGuard[StandardDispatchOutput]: return dispatch_output.format.is_standard() + @staticmethod + def format_is_ascend_tp( + dispatch_output: DispatchOutput, + ) -> TypeGuard[AscendTPDispatchOutput]: + return dispatch_output.format.is_ascend_tp() + @staticmethod def format_is_deepep_normal( dispatch_output: DispatchOutput, @@ -164,10 +172,14 @@ class DispatchOutputFormat(Enum): DEEPEP_NORMAL = "deepep_normal" DEEPEP_LL = "deepep_ll" FLASHINFER = "flashinfer" + ASCEND_TP = "ascend_tp" def is_standard(self) -> bool: return self == DispatchOutputFormat.STANDARD + def is_ascend_tp(self) -> bool: + return self == DispatchOutputFormat.ASCEND_TP + def is_deepep_normal(self) -> bool: return self == DispatchOutputFormat.DEEPEP_NORMAL @@ -204,6 +216,12 @@ class CombineInputChecker: ) -> TypeGuard[StandardCombineInput]: return combine_input.format == CombineInputFormat.STANDARD + @staticmethod + def format_is_ascend_tp( + combine_input: CombineInput, + ) -> TypeGuard[AscendTPCombineInput]: + return combine_input.format == CombineInputFormat.ASCEND_TP + @staticmethod def format_is_deepep_normal( combine_input: CombineInput, @@ -237,6 +255,7 @@ class CombineInputFormat(Enum): DEEPEP_NORMAL = "deepep_normal" DEEPEP_LL = "deepep_ll" FLASHINFER = "flashinfer" + ASCEND_TP = "ascend_tp" @runtime_checkable diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 40d4786b2..9ab8b166c 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -22,7 +22,7 @@ from sglang.srt.layers.moe.token_dispatcher.base import ( from sglang.srt.layers.moe.topk import TopKOutput from sglang.srt.layers.moe.utils import ( DeepEPMode, - DeepEPOutputDtype, + DispatcherOutputDtype, get_deepep_config, get_deepep_output_dtype, is_tbo_enabled, @@ -422,22 +422,22 @@ class _DeepEPDispatcherImplBase: # Configuration mapping for each dtype config_map = { - DeepEPOutputDtype.BF16: { + DispatcherOutputDtype.BF16: { "use_fp8": False, "use_nvfp4": False, }, - DeepEPOutputDtype.FP8: { + DispatcherOutputDtype.FP8: { "use_fp8": True, "use_nvfp4": False, }, # Needed for Ascend A2/A3 NPU case, # despite the use_fp8 flag, # quantization will be performed in int8 - DeepEPOutputDtype.INT8: { + DispatcherOutputDtype.INT8: { "use_fp8": True, "use_nvfp4": False, }, - DeepEPOutputDtype.NVFP4: { + DispatcherOutputDtype.NVFP4: { "use_fp8": False, "use_nvfp4": True, }, @@ -458,23 +458,23 @@ class _DeepEPDispatcherImplBase: def _validate_and_adjust_dtype(self) -> None: """Validate dtype against hardware and adjust if necessary.""" if _is_npu: - if self.deepep_output_dtype == DeepEPOutputDtype.FP8: + if self.deepep_output_dtype == DispatcherOutputDtype.FP8: logger.warning_once( "Ascend A2/A3 NPU does not support fp8 " "deepep_dispatcher_output_dtype, switching to int8..." ) - self.deepep_output_dtype = DeepEPOutputDtype.INT8 - elif self.deepep_output_dtype == DeepEPOutputDtype.NVFP4: + self.deepep_output_dtype = DispatcherOutputDtype.INT8 + elif self.deepep_output_dtype == DispatcherOutputDtype.NVFP4: raise RuntimeError( "Ascend A2/A3 NPU does not support nvfp4 deepep_dispatcher_output_dtype." ) else: - if self.deepep_output_dtype == DeepEPOutputDtype.INT8: + if self.deepep_output_dtype == DispatcherOutputDtype.INT8: logger.warning_once( "GPU does not support int8 " "deepep_dispatcher_output_dtype, switching to fp8..." ) - self.deepep_output_dtype = DeepEPOutputDtype.FP8 + self.deepep_output_dtype = DispatcherOutputDtype.FP8 # NVFP4 is supported on GPU, no adjustment needed def _update_int8_quant_env(self) -> None: diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 779ac855c..aaf1f1d09 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -33,6 +33,7 @@ class MoeA2ABackend(Enum): NIXL = "nixl" MORI = "mori" ASCEND_FUSEEP = "ascend_fuseep" + ASCEND_TP = "ascend_tp" FLASHINFER = "flashinfer" MEGAMOE = "megamoe" CUSTOMIZED = "customized" @@ -64,6 +65,9 @@ class MoeA2ABackend(Enum): def is_ascend_fuseep(self): return self == MoeA2ABackend.ASCEND_FUSEEP + def is_ascend_tp(self): + return self == MoeA2ABackend.ASCEND_TP + def is_mori(self): return self == MoeA2ABackend.MORI @@ -89,6 +93,7 @@ class MoeRunnerBackend(Enum): DEEP_GEMM = "deep_gemm" TRITON = "triton" TRITON_KERNELS = "triton_kernel" + ASCEND = "ascend" FLASHINFER_TRTLLM = "flashinfer_trtllm" EXPERIMENTAL_SGL_TRTLLM = "experimental_sgl_trtllm" FLASHINFER_TRTLLM_ROUTED = "flashinfer_trtllm_routed" @@ -109,6 +114,9 @@ class MoeRunnerBackend(Enum): def is_triton(self): return self == MoeRunnerBackend.TRITON + def is_ascend(self): + return self == MoeRunnerBackend.ASCEND + def is_triton_kernels(self): return self == MoeRunnerBackend.TRITON_KERNELS @@ -179,7 +187,7 @@ class DeepEPMode(Enum): return self == DeepEPMode.AUTO -class DeepEPOutputDtype(Enum): +class DispatcherOutputDtype(Enum): """ Describes the dispatch output data type for DeepEP. @@ -195,7 +203,7 @@ class DeepEPOutputDtype(Enum): NVFP4 = "nvfp4" -def get_deepep_output_dtype(self) -> DeepEPOutputDtype: +def get_deepep_output_dtype(self) -> DispatcherOutputDtype: """ Automatically choose the dispatch output dtype for DeepEP. @@ -212,7 +220,7 @@ def get_deepep_output_dtype(self) -> DeepEPOutputDtype: # 0. Parse server argument. server_args = get_server_args() if server_args and server_args.deepep_dispatcher_output_dtype != "auto": - return DeepEPOutputDtype(server_args.deepep_dispatcher_output_dtype) + return DispatcherOutputDtype(server_args.deepep_dispatcher_output_dtype) # 1. Parse deprecated environment variables. if envs.SGLANG_DEEPEP_BF16_DISPATCH.get(): @@ -221,18 +229,18 @@ def get_deepep_output_dtype(self) -> DeepEPOutputDtype: "and will be removed in future releases. Please use a new " "`--deepep-dispatcher-output-dtype bf16` argument instead." ) - return DeepEPOutputDtype.BF16 + return DispatcherOutputDtype.BF16 # 2. NVFP4 is detected inside dispatch_a / _dispatch_core via quant_config; no need to infer here. if self.quant_config is not None: input_global_scale = self.quant_config.get("input_global_scale", None) if input_global_scale is not None: - return DeepEPOutputDtype.NVFP4 + return DispatcherOutputDtype.NVFP4 # 3. Parse quant config to determine the output dtype of dispatcher dispatcher_output_dtype = self.quant_config.get("dispatcher_output_dtype", None) if dispatcher_output_dtype is not None: - return DeepEPOutputDtype(dispatcher_output_dtype) + return DispatcherOutputDtype(dispatcher_output_dtype) # 4. flashinfer_cutedsl / cutlass / humming expects BF16 dispatch if ( @@ -240,14 +248,31 @@ def get_deepep_output_dtype(self) -> DeepEPOutputDtype: or get_moe_runner_backend().is_cutlass() or get_moe_runner_backend().is_humming() ): - return DeepEPOutputDtype.BF16 + return DispatcherOutputDtype.BF16 # 5. Default on NPU → BF16 if _is_npu: - return DeepEPOutputDtype.BF16 + return DispatcherOutputDtype.BF16 # 6. Default → FP8 - return DeepEPOutputDtype.FP8 + return DispatcherOutputDtype.FP8 + + +def get_ascend_dispatcher_output_dtype(dispatcher): + """ + Automatically choose the dispatch output dtype for Ascend. + """ + + # 1. Parse quant config to determine the output dtype of dispatcher + if dispatcher.quant_config is not None: + dispatcher_output_dtype = dispatcher.quant_config.get( + "dispatcher_output_dtype", None + ) + if dispatcher_output_dtype is not None: + return DispatcherOutputDtype(dispatcher_output_dtype) + + # 2. Ascend dispatch defaults to BF16 + return DispatcherOutputDtype.BF16 def initialize_moe_config(server_args: ServerArgs): diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py index 9e8467f39..92bb0cdf5 100644 --- a/python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py @@ -6,17 +6,16 @@ from typing import TYPE_CHECKING import torch from sglang.srt.layers.linear import set_weight_attrs -from sglang.srt.layers.moe import ( - MoeRunner, - MoeRunnerBackend, - MoeRunnerConfig, - get_moe_runner_backend, -) +from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig +from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend from .awq_scheme import AWQMoESchemeBase if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) from sglang.srt.layers.quantization.awq.awq import AWQConfig, AWQMarlinConfig __all__ = ["AWQMoEScheme", "AWQAscendMoEScheme"] @@ -151,6 +150,43 @@ class AWQAscendMoEScheme(AWQMoEScheme): return AWQAscendMoEKernel(quant_config) def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + self, + layer: torch.nn.Module, + moe_runner_config: MoeRunnerConfig, + **extra_weight_attrs, ): + from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUWNA16Int4MoEMethod, + ) + self.moe_runner_config = moe_runner_config + layer.w13_kernel = NPUWNA16Int4MoEMethod() + layer.w2_kernel = NPUWNA16Int4MoEMethod() + moe_runner_config.layer = layer + backend = get_moe_runner_backend() + if backend.is_auto(): + backend = MoeRunnerBackend.ASCEND + self.runner = MoeRunner(backend, moe_runner_config) + + def apply_weights( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + from sglang.srt.layers.moe.moe_runner.ascend import ( + AscendQuantInfo, + ) + + quant_info = AscendQuantInfo( + w13_weight=layer.w13_qweight, + w2_weight=layer.w2_qweight, + w13_weight_scale=layer.w13_scales, + w2_weight_scale=layer.w2_scales, + w13_weight_offset=layer.w13_qzeros, + w2_weight_offset=layer.w2_qzeros, + w13_weight_bias=getattr(layer, "w13_weight_bias", None), + w2_weight_bias=getattr(layer, "w2_weight_bias", None), + w13_scale_bias=getattr(layer, "w13_scale_bias", None), + w2_scale_bias=getattr(layer, "w2_scale_bias", None), + ) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/base_scheme.py b/python/sglang/srt/layers/quantization/base_scheme.py index ee55caa3c..4155ad07f 100644 --- a/python/sglang/srt/layers/quantization/base_scheme.py +++ b/python/sglang/srt/layers/quantization/base_scheme.py @@ -1,15 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Optional +from typing import Optional import torch -from sglang.srt.layers.moe import MoeRunnerConfig - -if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput - __all__ = ["BaseLinearScheme", "BaseMoEScheme"] @@ -66,12 +61,6 @@ class BaseMoEScheme(ABC): """ raise NotImplementedError - @abstractmethod - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig - ): - raise NotImplementedError - @abstractmethod def process_weights_after_loading(self, layer: torch.nn.Module): """ @@ -79,21 +68,3 @@ class BaseMoEScheme(ABC): needs to occur. """ raise NotImplementedError - - @abstractmethod - def apply_weights( - self, - layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", - ): - """ - Run the forward pass for the particular scheme. This is where - scheme-specific dequant/quant steps/kernels should be applied. - - :param layer: torch.nn.Module with the registered weights and - other parameters relevant to the particular scheme. - :param x: input to the layer - :param bias: bias parameter - - """ - raise NotImplementedError diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int8_moe.py index b45b63fc4..45a25c6f0 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int8_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int8_moe.py @@ -5,10 +5,11 @@ from typing import TYPE_CHECKING import torch -from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( - NPUW4A8Int8DynamicMoEMethod, +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUW4A8Int8MoEMethod, ) -from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig +from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend from sglang.srt.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsMoEScheme, ) @@ -38,7 +39,14 @@ class NPUCompressedTensorsW4A8Int8DynamicMoE(CompressedTensorsMoEScheme): .get("group_1", {}) .get("activation_use_clip", False) ) - self.kernel = NPUW4A8Int8DynamicMoEMethod() + self.w13_kernel = NPUW4A8Int8MoEMethod( + is_per_channel_weight=self.is_per_channel_weight, + activation_use_clip=self.activation_use_clip, + ) + self.w2_kernel = NPUW4A8Int8MoEMethod( + is_per_channel_weight=self.is_per_channel_weight, + activation_use_clip=self.activation_use_clip, + ) def create_weights( self, @@ -257,37 +265,38 @@ class NPUCompressedTensorsW4A8Int8DynamicMoE(CompressedTensorsMoEScheme): set_weight_attrs(w2_scale_bias, extra_weight_attrs) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - self.kernel.process_weights_after_loading( - layer, self.is_per_channel_weight, self.activation_use_clip - ) + self.w13_kernel.process_weights_after_loading(layer, "w13") + self.w2_kernel.process_weights_after_loading(layer, "w2") def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): + layer.w13_kernel = self.w13_kernel + layer.w2_kernel = self.w2_kernel + moe_runner_config.layer = layer self.moe_runner_config = moe_runner_config + backend = get_moe_runner_backend() + if backend.is_auto(): + backend = MoeRunnerBackend.ASCEND + self.runner = MoeRunner(backend, moe_runner_config) def apply_weights( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: + from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo - return self.kernel.apply(layer, dispatch_output) - - def apply_weights_with_router_logits( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - return self.kernel.apply_without_routing_weights( - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, + quant_info = AscendQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_scale=layer.w2_weight_scale, + w13_weight_offset=layer.w13_weight_offset, + w2_weight_offset=layer.w2_weight_offset, + w13_scale_bias=layer.w13_scale_bias, + w2_scale_bias=layer.w2_scale_bias, + w13_weight_bias=getattr(layer, "w13_weight_bias", None), + w2_weight_bias=getattr(layer, "w2_weight_bias", None), ) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8_moe.py index 9df834726..1669899ee 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8_moe.py @@ -6,10 +6,11 @@ from typing import TYPE_CHECKING import torch from compressed_tensors.quantization import QuantizationStrategy -from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( - NPUW8A8Int8DynamicMoEMethod, +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUW8A8Int8MoEMethod, ) -from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig +from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend from sglang.srt.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsMoEScheme, ) @@ -31,7 +32,8 @@ class NPUCompressedTensorsW8A8Int8DynamicMoE(CompressedTensorsMoEScheme): def __init__(self, weight_quant, input_quant): self.weight_quant = weight_quant self.input_quant = input_quant - self.kernel = NPUW8A8Int8DynamicMoEMethod() + self.w13_kernel = NPUW8A8Int8MoEMethod() + self.w2_kernel = NPUW8A8Int8MoEMethod() self.static_input_scales = not self.input_quant.dynamic per_channel = ( @@ -118,37 +120,38 @@ class NPUCompressedTensorsW8A8Int8DynamicMoE(CompressedTensorsMoEScheme): layer.w2_input_scale = None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - self.kernel.process_weights_after_loading(layer) + self.w13_kernel.process_weights_after_loading(layer, "w13") + self.w2_kernel.process_weights_after_loading(layer, "w2") def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): + layer.w13_kernel = self.w13_kernel + layer.w2_kernel = self.w2_kernel + moe_runner_config.layer = layer self.moe_runner_config = moe_runner_config + backend = get_moe_runner_backend() + if backend.is_auto(): + backend = MoeRunnerBackend.ASCEND + self.runner = MoeRunner(backend, moe_runner_config) def apply_weights( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: + from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo - return self.kernel.apply(layer, dispatch_output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - # NPU MoE bypasses MoeRunner: expose the kernel's existing - # apply_without_routing_weights directly through the scheme. - return self.kernel.apply_without_routing_weights( - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, + quant_info = AscendQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_scale=layer.w2_weight_scale, + w13_weight_offset=layer.w13_weight_offset, + w2_weight_offset=layer.w2_weight_offset, + w13_weight_bias=getattr(layer, "w13_weight_bias", None), + w2_weight_bias=getattr(layer, "w2_weight_bias", None), + w13_scale_bias=getattr(layer, "w13_scale_bias", None), + w2_scale_bias=getattr(layer, "w2_scale_bias", None), ) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py index 236420847..cbdfe1144 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py @@ -11,10 +11,11 @@ from compressed_tensors import CompressionFormat from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import ( gptq_marlin_moe_repack, ) -from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( - NPUW4A16Int4DynamicMoEMethod, +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUWNA16Int4MoEMethod, ) -from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig +from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend from sglang.srt.layers.quantization.compressed_tensors.schemes import ( WNA16_SUPPORTED_BITS, CompressedTensorsMoEScheme, @@ -38,7 +39,6 @@ if TYPE_CHECKING: CompressedTensorsConfig, ) - __all__ = [ "CompressedTensorsWNA16MoE", "CompressedTensorsWNA16TritonMoE", @@ -578,7 +578,8 @@ class NPUCompressedTensorsW4A16Int4DynamicMoE(CompressedTensorsMoEScheme): else: self.group_size = 128 - self.kernel = NPUW4A16Int4DynamicMoEMethod() + self.w13_kernel = NPUWNA16Int4MoEMethod() + self.w2_kernel = NPUWNA16Int4MoEMethod() # TODO: See if we can merge this method's logic # with CompressedTensorsWNA16MoE. Need more models and tests. @@ -693,35 +694,38 @@ class NPUCompressedTensorsW4A16Int4DynamicMoE(CompressedTensorsMoEScheme): set_weight_attrs(w2_weight_shape, extra_weight_attrs) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - self.kernel.process_weights_after_loading(layer) + self.w13_kernel.process_weights_after_loading(layer, "w13") + self.w2_kernel.process_weights_after_loading(layer, "w2") def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): + layer.w13_kernel = self.w13_kernel + layer.w2_kernel = self.w2_kernel + moe_runner_config.layer = layer self.moe_runner_config = moe_runner_config + backend = get_moe_runner_backend() + if backend.is_auto(): + backend = MoeRunnerBackend.ASCEND + self.runner = MoeRunner(backend, moe_runner_config) def apply_weights( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: + from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo - return self.kernel.apply(layer, dispatch_output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - return self.kernel.apply_without_routing_weights( - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, + quant_info = AscendQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_scale=layer.w2_weight_scale, + w13_weight_offset=layer.w13_weight_offset, + w2_weight_offset=layer.w2_weight_offset, + w13_weight_bias=getattr(layer, "w13_weight_bias", None), + w2_weight_bias=getattr(layer, "w2_weight_bias", None), + w13_scale_bias=getattr(layer, "w13_scale_bias", None), + w2_scale_bias=getattr(layer, "w2_scale_bias", None), ) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/gguf.py b/python/sglang/srt/layers/quantization/gguf.py index 6b65e1ef7..e5fbc24ee 100644 --- a/python/sglang/srt/layers/quantization/gguf.py +++ b/python/sglang/srt/layers/quantization/gguf.py @@ -12,8 +12,13 @@ import torch from gguf import GGMLQuantizationType as WeightType from torch.nn.parameter import Parameter, UninitializedParameter +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUUnquantMoEMethod, +) +from sglang.srt.hardware_backend.npu.utils import npu_format_cast from sglang.srt.layers.linear import LinearBase -from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig +from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, LinearMethodBase, @@ -780,6 +785,8 @@ class GGUFMoEAscendMethod(FusedMoEMethodBase): def __init__(self, quant_config: GGUFConfig): self.quant_config = quant_config + self.w13_kernel = NPUUnquantMoEMethod() + self.w2_kernel = NPUUnquantMoEMethod() def create_weights( self, @@ -870,16 +877,18 @@ class GGUFMoEAscendMethod(FusedMoEMethodBase): torch.from_numpy(dequant_np) .to(dtype=self.params_dtype, device=w13_qweight.device) .reshape(rows, cols) - .transpose(-1, -2) .contiguous() ) w13_dequant_list.append(dequant) w13_full = torch.stack(w13_dequant_list, dim=0) - - layer.register_buffer("w13_dequant", w13_full, persistent=False) + layer.register_buffer( + "w13_dequant", npu_format_cast(w13_full), persistent=False + ) else: - layer.register_buffer("w13_dequant", w13_qweight.data, persistent=False) + layer.register_buffer( + "w13_dequant", npu_format_cast(w13_qweight.data), persistent=False + ) # Pre-dequantize w2 weights (down projection) w2_qweight = layer.w2_qweight @@ -901,137 +910,57 @@ class GGUFMoEAscendMethod(FusedMoEMethodBase): torch.from_numpy(dequant_np) .to(dtype=self.params_dtype, device=w2_qweight.device) .reshape(rows, cols) - .transpose(-1, -2) .contiguous() ) w2_dequant_list.append(dequant) w2_full = torch.stack(w2_dequant_list, dim=0) - layer.register_buffer("w2_dequant", w2_full, persistent=False) + layer.register_buffer( + "w2_dequant", npu_format_cast(w2_full), persistent=False + ) else: - layer.register_buffer("w2_dequant", w2_qweight.data, persistent=False) + layer.register_buffer( + "w2_dequant", npu_format_cast(w2_qweight.data), persistent=False + ) if hasattr(layer, "w2_qweight"): del layer.w2_qweight if hasattr(layer, "w13_qweight"): del layer.w13_qweight + if hasattr(layer, "dispatcher"): + layer.dispatcher.set_quant_config({"quant_type": "gguf"}) + def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): + layer.w13_kernel = self.w13_kernel + layer.w2_kernel = self.w2_kernel + moe_runner_config.layer = layer + moe_runner_config.use_tp_all_gather_activation = True self.moe_runner_config = moe_runner_config + backend = get_moe_runner_backend() + if backend.is_auto(): + backend = MoeRunnerBackend.ASCEND + self.runner = MoeRunner(backend, moe_runner_config) def apply( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: - """Apply MoE forward pass on NPU using npu_grouped_matmul for maximum performance.""" - from sglang.srt.distributed.communication_op import ( - tensor_model_parallel_all_gather, + from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo + + quant_info = AscendQuantInfo( + w13_weight=layer.w13_dequant, + w2_weight=layer.w2_dequant, + w13_weight_bias=getattr(layer, "w13_weight_bias", None), + w2_weight_bias=getattr(layer, "w2_weight_bias", None), + w13_scale_bias=getattr(layer, "w13_scale_bias", None), + w2_scale_bias=getattr(layer, "w2_scale_bias", None), ) - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - topk_weights, topk_ids, _ = topk_output - - # Check if pre-dequantized weights are available - use_pre_dequant = hasattr(layer, "w13_dequant") and hasattr(layer, "w2_dequant") - - if not use_pre_dequant: - raise RuntimeError( - "GGUF MoE on NPU requires pre-dequantization (FusedMoE fix). Please report if this occurs." - ) - - w13 = layer.w13_dequant - w2 = layer.w2_dequant - - num_experts = w13.shape[0] - - tp_size = getattr(layer, "moe_tp_size", 1) - - original_dtype = x.dtype - num_tokens = x.shape[0] - top_k = topk_ids.shape[1] - - # Ensure correct dtypes for NPU ops - topk_ids = topk_ids.to(torch.int32) - topk_weights = topk_weights.to(x.dtype) - - # MoE routing initialization - reorder tokens by expert - row_idx_len = num_tokens * top_k - row_idx = ( - torch.arange(0, row_idx_len, dtype=torch.int32, device=x.device) - .view(top_k, -1) - .permute(1, 0) - .contiguous() - ) - - sorted_hidden_states, expanded_row_idx, expanded_expert_idx = ( - torch.ops.npu.npu_moe_init_routing( - x, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens - ) - ) - - # Compute tokens per expert - expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens( - expanded_expert_idx, num_experts - ) - expert_tokens = expert_tokens.to(torch.int64) - - w13_gmm = w13 # No transpose needed - - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[sorted_hidden_states], - weight=[w13_gmm], - split_item=2, - group_list_type=0, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - - # Activation (SwiGLU) - hidden_states = torch.ops.npu.npu_swiglu(hidden_states) - - # TP all-gather for intermediate dimension if needed - if tp_size > 1: - hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=-1) - - w2_gmm = w2 - - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[w2_gmm], - split_item=2, - group_list_type=0, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - - # Finalize routing - reorder back and apply weights - final_hidden_states = torch.ops.npu.npu_moe_finalize_routing( - hidden_states, - skip1=None, - skip2=None, - bias=None, - scales=topk_weights, - expanded_src_to_dst_row=expanded_row_idx, - export_for_source_row=topk_ids, - ) - - if tp_size > 1: - final_hidden_states = tensor_model_parallel_all_gather( - final_hidden_states, dim=-1 - ) - - # Ensure output matches input dtype - final_hidden_states = final_hidden_states.to(dtype=original_dtype) - - return StandardCombineInput(hidden_states=final_hidden_states) + return self.runner.run(dispatch_output, quant_info) class GGUFEmbeddingAscendMethod(GGUFLinearAscendMethod): diff --git a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_moe.py b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_moe.py index c83bdba19..d3481cb5e 100644 --- a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_moe.py +++ b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_moe.py @@ -5,13 +5,20 @@ from typing import TYPE_CHECKING import torch +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUWNA16Int4MoEMethod, +) from sglang.srt.layers.linear import set_weight_attrs -from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig +from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend from .gptq_scheme import GPTQMoESchemeBase if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) from sglang.srt.layers.quantization.gptq.gptq import GPTQConfig, GPTQMarlinConfig __all__ = ["GPTQMoEAscendScheme", "GPTQMarlinMoEScheme"] @@ -122,9 +129,19 @@ class GPTQMoEAscendScheme(GPTQMoESchemeBase): set_weight_attrs(w2_qzeros, extra_weight_attrs) def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + self, + layer: torch.nn.Module, + moe_runner_config: MoeRunnerConfig, + **extra_weight_attrs, ): - self.kernel.create_moe_runner(layer, moe_runner_config) + self.moe_runner_config = moe_runner_config + layer.w13_kernel = NPUWNA16Int4MoEMethod() + layer.w2_kernel = NPUWNA16Int4MoEMethod() + moe_runner_config.layer = layer + backend = get_moe_runner_backend() + if backend.is_auto(): + backend = MoeRunnerBackend.ASCEND + self.runner = MoeRunner(backend, moe_runner_config) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: self.kernel.process_weights_after_loading(layer) @@ -133,8 +150,22 @@ class GPTQMoEAscendScheme(GPTQMoESchemeBase): self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, - ): - return self.kernel.apply(layer, dispatch_output) + ) -> CombineInput: + from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo + + quant_info = AscendQuantInfo( + w13_weight=layer.w13_qweight, + w2_weight=layer.w2_qweight, + w13_weight_scale=layer.w13_scales, + w2_weight_scale=layer.w2_scales, + w13_weight_offset=layer.w13_qzeros, + w2_weight_offset=layer.w2_qzeros, + w13_weight_bias=getattr(layer, "w13_weight_bias", None), + w2_weight_bias=getattr(layer, "w2_weight_bias", None), + w13_scale_bias=getattr(layer, "w13_scale_bias", None), + w2_scale_bias=getattr(layer, "w2_scale_bias", None), + ) + return self.runner.run(dispatch_output, quant_info) class GPTQMarlinMoEScheme(GPTQMoESchemeBase): diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index b69936985..7af3f9171 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -9,6 +9,8 @@ import torch from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( _NPULinearMethodBase, ) +from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig +from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_runner_backend from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, QuantizationConfig, @@ -34,7 +36,6 @@ if TYPE_CHECKING: from sglang.srt.layers.quantization.base_config import QuantizeMethodBase from sglang.srt.layers.quantization.modelslim.schemes import ( ModelSlimLinearScheme, - ModelSlimMoEScheme, ) logger = logging.getLogger(__name__) @@ -182,7 +183,14 @@ class ModelSlimConfig(QuantizationConfig): return UnquantizedLinearMethod() return ModelSlimLinearMethod(self) elif isinstance(layer, FusedMoE): - layer.scheme = self.get_moe_scheme(layer, prefix) + moe_schemes = self.get_moe_scheme(layer, prefix) + if moe_schemes is None: + raise ValueError(f"No ModelSlim MoE scheme found for layer {prefix}") + layer.w13_scheme, layer.w2_scheme = moe_schemes + layer.w13_kernel, layer.w2_kernel = ( + layer.w13_scheme.kernel, + layer.w2_scheme.kernel, + ) return ModelSlimFusedMoEMethod(self) return None @@ -219,29 +227,69 @@ class ModelSlimConfig(QuantizationConfig): self, layer: torch.nn.Module, prefix: str, - ) -> Optional[ModelSlimMoEScheme]: + ): moe_quant_schemes = [ ("W4A4_DYNAMIC", ModelSlimW4A4Int4MoE), ("W4A8_DYNAMIC", ModelSlimW4A8Int8MoE), ("W8A8_DYNAMIC", ModelSlimW8A8Int8MoE), ] - - moe_weight_suffixes = [".0.gate_proj.weight", ".0.w2.weight"] - quant_schemes = [ - self.quant_description.get(prefix + suffix, "") - for suffix in moe_weight_suffixes + w13_keys = [ + prefix + ".0.gate_proj.weight", + prefix + ".0.up_proj.weight", ] + w2_key = prefix + ".0.down_proj.weight" + w13_entries = { + key: self.quant_description[key] + for key in w13_keys + if key in self.quant_description + } + if not w13_entries or w2_key not in self.quant_description: + missing_groups = [] + if not w13_entries: + missing_groups.append(f"W13 ({', '.join(w13_keys)})") + if w2_key not in self.quant_description: + missing_groups.append(f"W2 ({w2_key})") + raise ValueError( + f"Missing ModelSlim MoE quantization description for layer {prefix}: " + + ", ".join(missing_groups) + ) - for scheme_name, scheme_class in moe_quant_schemes: - if any(s == scheme_name for s in quant_schemes): - logger.info_once(f"Using {scheme_class.__name__}") - return scheme_class(self) + w13_names = list(w13_entries.values()) + w2_name = self.quant_description[w2_key] - logger.warning( - f"Unsupported FusedMoe modelslim scheme: " - f"{quant_schemes} in layer: {prefix}" - ) - return None + # For w13, gate_proj and up_proj must agree on the scheme + unique_w13 = set(w13_names) + if len(unique_w13) > 1: + raise ValueError( + f"Mismatched ModelSlim quantization for W13 in layer {prefix}: " + f"{w13_entries}" + ) + w13_scheme_name = w13_names[0] + + # Map scheme names to classes + scheme_map = dict( + moe_quant_schemes + ) # dict: "W4A4_DYNAMIC" -> ModelSlimW4A4Int4MoE, etc. + + # Instantiate the schemes + def instantiate(name, weight_group): + cls = scheme_map.get(name) + if cls is None: + logger.warning(f"Unsupported scheme '{name}' for layer {prefix}") + return None + return cls(self, weight_group) + + w13_scheme = instantiate(w13_scheme_name, weight_group="w13") + w2_scheme = instantiate(w2_name, weight_group="w2") + if w13_scheme is None or w2_scheme is None: + raise ValueError( + f"Unsupported ModelSlim MoE schemes for layer {prefix}: " + f"gate/up={w13_names}, down_proj='{w2_name}'" + ) + logger.info_once(f"Using {type(w13_scheme).__name__} for gate_up_proj") + logger.info_once(f"Using {type(w2_scheme).__name__} for down_proj") + + return w13_scheme, w2_scheme def is_layer_skipped( self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({}) @@ -332,12 +380,19 @@ class ModelSlimLinearMethod(_NPULinearMethodBase): class ModelSlimFusedMoEMethod(FusedMoEMethodBase): + """ + Fused MoE method for ModelSlim quantization on Ascend NPU. + + Delegates routing, activation, and finalization to the modular NPU MoE + components introduced in the hardware backend refactoring. + """ def __init__(self, quantization_config: ModelSlimConfig): self.quantization_config = quantization_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - layer.scheme.process_weights_after_loading(layer) + layer.w13_scheme.process_weights_after_loading(layer) + layer.w2_scheme.process_weights_after_loading(layer) def create_weights( self, @@ -353,50 +408,53 @@ class ModelSlimFusedMoEMethod(FusedMoEMethodBase): the necessary parameters for the layer. See FusedMoEMethodBase for param details """ - layer.scheme.create_weights( + layer.w13_scheme.create_weights( layer=layer, num_experts=num_experts, hidden_size=hidden_size, intermediate_size_per_partition=intermediate_size_per_partition, - params_dtype=params_dtype, + weight_prefix="w13", + **extra_weight_attrs, + ) + layer.w2_scheme.create_weights( + layer=layer, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size_per_partition, + weight_prefix="w2", **extra_weight_attrs, ) def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): - return layer.scheme.create_moe_runner(layer, moe_runner_config) + moe_runner_config.layer = layer + self.moe_runner_config = moe_runner_config + backend = get_moe_runner_backend() + if backend.is_auto(): + backend = MoeRunnerBackend.ASCEND + self.runner = MoeRunner(backend, moe_runner_config) + # ------------------------------------------------------------------ + # Main apply() + # ------------------------------------------------------------------ def apply( self, - layer: torch.nn.Module, + layer, dispatch_output: StandardDispatchOutput, ) -> CombineInput: - """ - Use the output of create_weights and the ModelSlimMoEScheme - associated with the layer to apply the forward pass with the - layer input. See FusedMoEMethodBase for param details + from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo - """ - scheme = layer.scheme - if scheme is None: - raise ValueError("A scheme must be defined for each layer") - return scheme.apply_weights(layer, dispatch_output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - return layer.scheme.apply_without_routing_weights( - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, + quant_info = AscendQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_scale=layer.w2_weight_scale, + w13_weight_offset=layer.w13_weight_offset, + w2_weight_offset=layer.w2_weight_offset, + w13_scale_bias=getattr(layer, "w13_scale_bias", None), + w2_scale_bias=getattr(layer, "w2_scale_bias", None), + w13_weight_bias=getattr(layer, "w13_weight_bias", None), + w2_weight_bias=getattr(layer, "w2_weight_bias", None), ) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_scheme.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_scheme.py index c95a7637a..3d7a28741 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_scheme.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_scheme.py @@ -3,16 +3,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import abstractmethod -from typing import TYPE_CHECKING, Optional +from typing import Optional import torch -from sglang.srt.layers.moe import MoeRunnerConfig from sglang.srt.layers.quantization.base_scheme import BaseLinearScheme, BaseMoEScheme -if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput - __all__ = ["ModelSlimLinearScheme", "ModelSlimMoEScheme"] @@ -76,26 +72,3 @@ class ModelSlimMoEScheme(BaseMoEScheme): needs to occur. """ raise NotImplementedError - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig" - ): - raise NotImplementedError - - @abstractmethod - def apply_weights( - self, - layer, - dispatch_output: "StandardDispatchOutput", - ): - """ - Run the forward pass for the particular scheme. This is where - scheme-specific dequant/quant steps/kernels should be applied. - - :param layer: torch.nn.Module with the registered weights and - other parameters relevant to the particular scheme. - :param x: input to the layer - :param bias: bias parameter - - """ - raise NotImplementedError diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py index 10c9ff863..c55d71668 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py @@ -6,6 +6,7 @@ from typing import Any, Dict, List, Optional import torch +from sglang.srt.environ import envs from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( NPU_W4A4DynamicLinearMethod, ) @@ -29,6 +30,8 @@ class ModelSlimW4A4Int4(ModelSlimLinearScheme): def get_weight( input_size: int, output_size: int, params_dtype: torch.dtype ) -> Dict[str, Any]: + if envs.SGLANG_NPU_W4A4_NEW_PACKING.get(): + output_size = output_size // 2 params_dict = {"weight": torch.empty(output_size, input_size, dtype=torch.int8)} return params_dict @@ -54,15 +57,25 @@ class ModelSlimW4A4Int4(ModelSlimLinearScheme): ) -> None: output_size_per_partition = sum(output_partition_sizes) weight_loader = extra_weight_attrs.get("weight_loader") - + if envs.SGLANG_NPU_W4A4_NEW_PACKING.get(): + weight_output_size_per_partition = output_size_per_partition // 2 + else: + weight_output_size_per_partition = output_size_per_partition weight_dict = { "weight": torch.empty( - output_size_per_partition, input_size_per_partition, dtype=torch.int8 + weight_output_size_per_partition, + input_size_per_partition, + dtype=torch.int8, ) } for weight_name, weight_param in weight_dict.items(): param = torch.nn.Parameter(weight_param, requires_grad=False) - set_weight_attrs(param, {"input_dim": 1, "output_dim": 0}) + if envs.SGLANG_NPU_W4A4_NEW_PACKING.get(): + set_weight_attrs( + param, {"input_dim": 1, "output_dim": 0, "pack_factor": 2} + ) + else: + set_weight_attrs(param, {"input_dim": 1, "output_dim": 0}) layer.register_parameter(weight_name, param) set_weight_attrs(param, extra_weight_attrs) diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py index 95ce416d3..cd1d722a0 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py @@ -1,23 +1,17 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, Dict +from typing import Any, Dict import torch -from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( - NPUW4A4Int4DynamicMoEMethod, +from sglang.srt.environ import envs +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUW4A4Int4MoEMethod, ) from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimMoEScheme from sglang.srt.utils import set_weight_attrs -if TYPE_CHECKING: - from sglang.srt.layers.moe import MoeRunnerConfig - from sglang.srt.layers.moe.token_dispatcher import ( - CombineInput, - StandardDispatchOutput, - ) - logger = logging.getLogger(__name__) __all__ = [ @@ -26,14 +20,30 @@ __all__ = [ class ModelSlimW4A4Int4MoE(ModelSlimMoEScheme): + """ + W4A4 integer MoE scheme that creates weights for either the + w13 (gate+up) or w2 (down) projection group. + + Two instances of this class are used per MoE layer: + - weight_prefix="w13" → handles the fused gate_proj + up_proj weights + - weight_prefix="w2" → handles the down_proj weights + """ def __init__( self, quant_config: Dict[str, Any], - prefix: str = None, - ): + weight_prefix: str, # "w13" or "w2" + group_size: int = 0, + ) -> None: self.quant_config = quant_config - self.kernel = NPUW4A4Int4DynamicMoEMethod() + self.kernel = NPUW4A4Int4MoEMethod() + self.weight_prefix = weight_prefix + self.group_size = group_size + self.is_per_channel_weight = group_size == 0 + if weight_prefix not in ("w13", "w2"): + raise ValueError( + f"weight_prefix must be 'w13' or 'w2', got '{weight_prefix}'" + ) def create_weights( self, @@ -41,103 +51,87 @@ class ModelSlimW4A4Int4MoE(ModelSlimMoEScheme): num_experts: int, hidden_size: int, intermediate_size_per_partition: int, - params_dtype: torch.dtype, **extra_weight_attrs, ) -> None: + """ + Create and register weight, scale, and offset parameters for the layer. + Shape depends on the W4A4 packing environment flag and whether the weight + prefix is "w13" or "w2". + """ from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported - self.num_experts = num_experts extra_weight_attrs.update( {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} ) + # --- compute shapes based on the packing path and prefix --- + if envs.SGLANG_NPU_W4A4_NEW_PACKING.get(): + if self.weight_prefix == "w13": + out_features = intermediate_size_per_partition + in_features = hidden_size + else: # w2 + out_features = hidden_size // 2 + in_features = intermediate_size_per_partition - # weight - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size, - dtype=torch.int8, - ), + weight_shape = (num_experts, out_features, in_features) + scale_shape = (num_experts, 2 * out_features, 1) + else: + if self.weight_prefix == "w13": + a_dim = 2 * intermediate_size_per_partition + b_dim = hidden_size + else: # w2 + a_dim = hidden_size + b_dim = intermediate_size_per_partition + + weight_shape = (num_experts, a_dim, b_dim) + scale_shape = (num_experts, a_dim, 1) + + offset_shape = scale_shape # offset always matches scale + + self._create_weight_params( + layer, + self.weight_prefix, + weight_shape, + scale_shape, + offset_shape, + extra_weight_attrs, + ) + + @staticmethod + def _create_weight_params( + layer: torch.nn.Module, + prefix: str, + weight_shape: tuple, + scale_shape: tuple, + offset_shape: tuple, + extra_weight_attrs: dict, + ) -> None: + """Helper that registers weight, scale, and offset as parameters.""" + # Weight + weight = torch.nn.Parameter( + torch.empty(weight_shape, dtype=torch.int8), requires_grad=False, ) - layer.register_parameter("w13_weight", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition, - dtype=torch.int8, - ), + layer.register_parameter(f"{prefix}_weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + + # Scale + scale = torch.nn.Parameter( + torch.empty(scale_shape, dtype=torch.float32), requires_grad=False, ) - layer.register_parameter("w2_weight", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - # scale - w13_weight_scale = torch.nn.Parameter( - torch.empty( - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 - ), + layer.register_parameter(f"{prefix}_weight_scale", scale) + set_weight_attrs(scale, extra_weight_attrs) + + # Offset + offset = torch.nn.Parameter( + torch.empty(offset_shape, dtype=torch.float32), requires_grad=False, ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - w2_weight_scale = torch.nn.Parameter( - torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - # offset - w13_weight_offset = torch.nn.Parameter( - torch.empty( - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_offset", w13_weight_offset) - set_weight_attrs(w13_weight_offset, extra_weight_attrs) - w2_weight_offset = torch.nn.Parameter( - torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_weight_offset", w2_weight_offset) - set_weight_attrs(w2_weight_offset, extra_weight_attrs) + layer.register_parameter(f"{prefix}_weight_offset", offset) + set_weight_attrs(offset, extra_weight_attrs) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - self.kernel.process_weights_after_loading(layer) - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig - ): - self.moe_runner_config = moe_runner_config - - def apply_weights( - self, - layer, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - return self.kernel.apply(layer, dispatch_output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - logger.warning_once( - "Warning: Performance may be reduced, because DeepEP Dispatcher does not support 4-bit quantization, " - "switching to the bf16 dispatcher, quantization will be performed separately..." - ) - return self.kernel.apply_without_routing_weights( - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ) + """ + Delegate weight processing to the NPU kernel for the fixed weight group. + """ + self.kernel.process_weights_after_loading(layer, self.weight_prefix) diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_int8_moe.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_int8_moe.py index 2c292d40c..ca2eef848 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_int8_moe.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_int8_moe.py @@ -1,43 +1,55 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, Dict +from typing import Any, Dict import torch -from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( - NPUW4A8Int8DynamicMoEMethod, +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUW4A8Int8MoEMethod, ) from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimMoEScheme from sglang.srt.utils import set_weight_attrs -if TYPE_CHECKING: - from sglang.srt.layers.moe import MoeRunnerConfig - from sglang.srt.layers.moe.token_dispatcher import ( - CombineInput, - StandardDispatchOutput, - ) - logger = logging.getLogger(__name__) -__all__ = [ - "ModelSlimW4A8Int8MoE", -] +__all__ = ["ModelSlimW4A8Int8MoE"] class ModelSlimW4A8Int8MoE(ModelSlimMoEScheme): + """ + W4A8 MoE scheme for a single weight group (w13 or w2). + + Two instances of this class are created per MoE layer: + - weight_prefix="w13" → handles gate + up projections + - weight_prefix="w2" → handles down projection + + Configuration flags (``is_per_channel_weight``, ``activation_use_clip``) + are passed to the underlying NPU kernel. + """ def __init__( self, quant_config: Dict[str, Any], - prefix: str = None, - ): + weight_prefix: str, + group_size: int = 0, + tp_size: int = 1, + activation_use_clip: bool = False, + ) -> None: + if weight_prefix not in ("w13", "w2"): + raise ValueError( + f"weight_prefix must be 'w13' or 'w2', got '{weight_prefix}'" + ) self.quant_config = quant_config - self.group_size = 0 - self.is_per_channel_weight = self.group_size == 0 - self.tp_size = 1 - self.activation_use_clip = False - self.kernel = NPUW4A8Int8DynamicMoEMethod() + self.weight_prefix = weight_prefix + self.group_size = group_size + self.tp_size = tp_size + self.is_per_channel_weight = group_size == 0 + self.activation_use_clip = activation_use_clip + self.kernel = NPUW4A8Int8MoEMethod( + is_per_channel_weight=self.is_per_channel_weight, + activation_use_clip=self.activation_use_clip, + ) def create_weights( self, @@ -45,173 +57,89 @@ class ModelSlimW4A8Int8MoE(ModelSlimMoEScheme): num_experts: int, hidden_size: int, intermediate_size_per_partition: int, - params_dtype: torch.dtype, **extra_weight_attrs, ) -> None: from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported - self.is_per_channel_weight = self.group_size == 0 self.num_experts = num_experts extra_weight_attrs.update( {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} ) - # >> weight - w13_output_size = intermediate_size_per_partition - w2_output_size = hidden_size // 2 - w13_weight = torch.nn.Parameter( - torch.empty(num_experts, w13_output_size, hidden_size, dtype=torch.int8), - requires_grad=False, - ) - layer.register_parameter("w13_weight", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - w2_output_size, - intermediate_size_per_partition, - dtype=torch.int8, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) + # Determine dimensions based on weight group + if self.weight_prefix == "w13": + out_features = intermediate_size_per_partition + in_features = hidden_size + bias_last_dim = 1 + else: # w2 + out_features = hidden_size // 2 + in_features = intermediate_size_per_partition + bias_last_dim = 16 // self.tp_size - # >> scale - w13_weight_scale = torch.nn.Parameter( - torch.empty( - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 - ), + prefix = self.weight_prefix + + # ---- weight ---- + weight = torch.nn.Parameter( + torch.empty(num_experts, out_features, in_features, dtype=torch.int8), requires_grad=False, ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) + layer.register_parameter(f"{prefix}_weight", weight) + set_weight_attrs(weight, extra_weight_attrs) - w2_weight_scale = torch.nn.Parameter( - torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), + # ---- scale ---- + + scale = torch.nn.Parameter( + torch.empty(num_experts, 2 * out_features, 1, dtype=torch.float32), requires_grad=False, ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) + layer.register_parameter(f"{prefix}_weight_scale", scale) + set_weight_attrs(scale, extra_weight_attrs) - # >> offset - w13_weight_offset = torch.nn.Parameter( - torch.empty( - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 - ), + # ---- offset ---- + offset = torch.nn.Parameter( + torch.empty(num_experts, 2 * out_features, 1, dtype=torch.float32), requires_grad=False, ) - layer.register_parameter("w13_weight_offset", w13_weight_offset) - set_weight_attrs(w13_weight_offset, extra_weight_attrs) + layer.register_parameter(f"{prefix}_weight_offset", offset) + set_weight_attrs(offset, extra_weight_attrs) - w2_weight_offset = torch.nn.Parameter( - torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_weight_offset", w2_weight_offset) - set_weight_attrs(w2_weight_offset, extra_weight_attrs) - - # >>> special param for w4a8 + # ---- per‑group second scale/offset (when not per‑channel) ---- if not self.is_per_channel_weight: - w13_weight_scale_second = torch.nn.Parameter( + scale_second = torch.nn.Parameter( torch.empty( num_experts, - 2 * intermediate_size_per_partition, - hidden_size // self.group_size, + 2 * out_features, + in_features // self.group_size, dtype=torch.float32, ), requires_grad=False, ) - layer.register_parameter("w13_weight_scale_second", w13_weight_scale_second) - set_weight_attrs(w13_weight_scale_second, extra_weight_attrs) - w13_weight_offset_second = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size // self.group_size, - dtype=torch.float32, - ), - requires_grad=False, - ) - layer.register_parameter( - "w13_weight_offset_second", w13_weight_offset_second - ) - set_weight_attrs(w13_weight_offset_second, extra_weight_attrs) + layer.register_parameter(f"{prefix}_weight_scale_second", scale_second) + set_weight_attrs(scale_second, extra_weight_attrs) - w2_weight_scale_second = torch.nn.Parameter( + offset_second = torch.nn.Parameter( torch.empty( num_experts, - hidden_size, - intermediate_size_per_partition // self.group_size, + 2 * out_features, + in_features // self.group_size, dtype=torch.float32, ), requires_grad=False, ) - layer.register_parameter("w2_weight_scale_second", w2_weight_scale_second) - set_weight_attrs(w2_weight_scale_second, extra_weight_attrs) + layer.register_parameter(f"{prefix}_weight_offset_second", offset_second) + set_weight_attrs(offset_second, extra_weight_attrs) - w2_weight_offset_second = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition // self.group_size, - dtype=torch.float32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_offset_second", w2_weight_offset_second) - set_weight_attrs(w2_weight_offset_second, extra_weight_attrs) - - w13_scale_bias = torch.nn.Parameter( + # ---- bias for scale (activation clip path) ---- + # This parameter is always created; the kernel uses it only when activation_use_clip is True. + scale_bias = torch.nn.Parameter( torch.empty( - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + num_experts, 2 * out_features, bias_last_dim, dtype=torch.float32 ), requires_grad=False, ) - layer.register_parameter("w13_scale_bias", w13_scale_bias) - set_weight_attrs(w13_scale_bias, extra_weight_attrs) - - w2_scale_bias = torch.nn.Parameter( - torch.empty( - num_experts, hidden_size, 16 // self.tp_size, dtype=torch.float32 - ), - requires_grad=False, - ) - layer.register_parameter("w2_scale_bias", w2_scale_bias) - set_weight_attrs(w2_scale_bias, extra_weight_attrs) + layer.register_parameter(f"{prefix}_scale_bias", scale_bias) + set_weight_attrs(scale_bias, extra_weight_attrs) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - self.kernel.process_weights_after_loading( - layer, self.is_per_channel_weight, self.activation_use_clip - ) - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig - ): - self.moe_runner_config = moe_runner_config - - def apply_weights( - self, - layer, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - # FIXME W4A8 without EP can give 0 accuracy - return self.kernel.apply(layer, dispatch_output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - return self.kernel.apply_without_routing_weights( - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ) + """Delegate weight processing to the kernel for the assigned weight group.""" + self.kernel.process_weights_after_loading(layer, self.weight_prefix) diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py index 01b705fef..22c3623aa 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py @@ -1,23 +1,16 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, Dict +from typing import Any, Dict import torch -from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( - NPUW8A8Int8DynamicMoEMethod, +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUW8A8Int8MoEMethod, ) from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimMoEScheme from sglang.srt.utils import set_weight_attrs -if TYPE_CHECKING: - from sglang.srt.layers.moe import MoeRunnerConfig - from sglang.srt.layers.moe.token_dispatcher import ( - CombineInput, - StandardDispatchOutput, - ) - logger = logging.getLogger(__name__) __all__ = [ @@ -26,14 +19,27 @@ __all__ = [ class ModelSlimW8A8Int8MoE(ModelSlimMoEScheme): + """ + W8A8 integer MoE scheme that creates weights for either the + w13 (gate+up) or w2 (down) projection group. + + Two instances of this class are used per MoE layer: + - weight_prefix="w13" → handles the fused gate_proj + up_proj weights + - weight_prefix="w2" → handles the down_proj weights + """ def __init__( self, quant_config: Dict[str, Any], - prefix: str = None, - ): + weight_prefix: str, # "w13" or "w2" + ) -> None: self.quant_config = quant_config - self.kernel = NPUW8A8Int8DynamicMoEMethod() + self.kernel = NPUW8A8Int8MoEMethod() + self.weight_prefix = weight_prefix + if weight_prefix not in ("w13", "w2"): + raise ValueError( + f"weight_prefix must be 'w13' or 'w2', got '{weight_prefix}'" + ) def create_weights( self, @@ -41,7 +47,6 @@ class ModelSlimW8A8Int8MoE(ModelSlimMoEScheme): num_experts: int, hidden_size: int, intermediate_size_per_partition: int, - params_dtype: torch.dtype, **extra_weight_attrs, ) -> None: from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported @@ -51,89 +56,45 @@ class ModelSlimW8A8Int8MoE(ModelSlimMoEScheme): {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} ) - # weight - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size, - dtype=torch.int8, - ), + # Determine shape based on weight group + if self.weight_prefix == "w13": + a_dim = 2 * intermediate_size_per_partition + b_dim = hidden_size + else: # w2 + a_dim = hidden_size + b_dim = intermediate_size_per_partition + + prefix = self.weight_prefix + + # Create and register weight + weight_name = f"{prefix}_weight" + weight = torch.nn.Parameter( + torch.empty(num_experts, a_dim, b_dim, dtype=torch.int8), requires_grad=False, ) - layer.register_parameter("w13_weight", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition, - dtype=torch.int8, - ), + layer.register_parameter(weight_name, weight) + set_weight_attrs(weight, extra_weight_attrs) + + # Create and register scale + scale_name = f"{prefix}_weight_scale" + scale = torch.nn.Parameter( + torch.empty(num_experts, a_dim, 1, dtype=torch.float32), requires_grad=False, ) - layer.register_parameter("w2_weight", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - # scale - w13_weight_scale = torch.nn.Parameter( - torch.empty( - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 - ), + layer.register_parameter(scale_name, scale) + set_weight_attrs(scale, extra_weight_attrs) + + # Create and register offset + offset_name = f"{prefix}_weight_offset" + offset = torch.nn.Parameter( + torch.empty(num_experts, a_dim, 1, dtype=torch.float32), requires_grad=False, ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - w2_weight_scale = torch.nn.Parameter( - torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - # offset - w13_weight_offset = torch.nn.Parameter( - torch.empty( - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_offset", w13_weight_offset) - set_weight_attrs(w13_weight_offset, extra_weight_attrs) - w2_weight_offset = torch.nn.Parameter( - torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_weight_offset", w2_weight_offset) - set_weight_attrs(w2_weight_offset, extra_weight_attrs) + layer.register_parameter(offset_name, offset) + set_weight_attrs(offset, extra_weight_attrs) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - self.kernel.process_weights_after_loading(layer) - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig - ): - self.moe_runner_config = moe_runner_config - - def apply_weights( - self, - layer, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - return self.kernel.apply(layer, dispatch_output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - return self.kernel.apply_without_routing_weights( - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ) + """ + Delegate weight processing to the NPU kernel for the fixed weight group. + """ + self.kernel.process_weights_after_loading(layer, self.weight_prefix) diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 10da8b562..2f6406d2c 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -49,6 +49,9 @@ if TYPE_CHECKING: ) from sglang.srt.server_args import ServerArgs +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUUnquantMoEMethod, +) _is_cpu_amx_available = cpu_has_amx_support() _is_hip = is_hip() @@ -60,9 +63,6 @@ if _use_aiter: from aiter.ops.shuffle import shuffle_weight from aiter.tuned_gemm import tgemm -if _is_npu: - from sglang.srt.hardware_backend.npu.utils import npu_format_cast - class Bf16GemmBackend(Enum): AUTO = "auto" @@ -403,9 +403,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): layer.num_local_experts, *new_shape_w2 ) if _is_npu: - for weight_name in ["w13_weight", "w2_weight"]: - weight = getattr(layer, weight_name) - weight.data = npu_format_cast(weight) + layer.w13_kernel.process_weights_after_loading(layer, "w13") + layer.w2_kernel.process_weights_after_loading(layer, "w2") + if hasattr(layer, "dispatcher"): + layer.dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"}) return @@ -474,6 +475,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): backend = MoeRunnerBackend.DEEP_GEMM elif self.use_triton_kernels: backend = MoeRunnerBackend.TRITON_KERNELS + elif _is_npu: + layer.w13_kernel = NPUUnquantMoEMethod() + layer.w2_kernel = NPUUnquantMoEMethod() + moe_runner_config.layer = layer + backend = MoeRunnerBackend.ASCEND else: backend = MoeRunnerBackend.TRITON self.runner = MoeRunner(backend, moe_runner_config) @@ -721,151 +727,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): dispatch_output: DispatchOutput, ) -> CombineInput: - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker - - if DispatchOutputChecker.format_is_deepep(dispatch_output): - return self._forward_npu_deepep(layer, dispatch_output) - - # x.shape = [B*S, H] - x = dispatch_output.hidden_states - # topk_weights.shape = [B*S, K]; topk_ids.shape = [B*S, K] - topk_weights, topk_ids, _ = dispatch_output.topk_output - - original_dtype = x.dtype - num_tokens = x.shape[0] - topk_weights = topk_weights.to(x.dtype) - topk_ids = topk_ids.to(torch.int32) - num_experts = layer.num_experts - top_k = layer.top_k or topk_ids.shape[1] # in case layer.top_k is not set - - hidden_states, expanded_row_idx, expert_tokens, _ = ( - torch.ops.npu.npu_moe_init_routing_v2( - x, - topk_ids, - active_num=num_tokens * top_k, - expert_num=num_experts, - expert_tokens_num_type=1, - expert_tokens_num_flag=True, - active_expert_range=[0, num_experts], - quant_mode=-1, - ) - ) - expert_tokens = expert_tokens.to(torch.int64) - w13_bias = [layer.w13_weight_bias] if self.with_bias else None - w2_bias = [layer.w2_weight_bias] if self.with_bias else None - - # gmm1: gate_up_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w13_weight.transpose(1, 2)], - bias=w13_bias, - split_item=2, - group_list_type=1, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - - # act_fn: - if self.moe_runner_config.activation == "npu_swiglu_oai": - from sgl_kernel_npu.activation.swiglu_oai import swiglu_oai_triton - - # `hidden_states` is the gmm1 output of shape [num_tokens, 2 * inter]. - # Pass the gate_up dim from the activation itself instead of letting - # swiglu_oai() derive it from layer.w13_weight.shape[2]: w13_weight is - # now stored un-transposed (transposed on the fly for the grouped - # matmuls above), so shape[2] is `hidden`, not the gate_up dim, which - # makes the kernel's view(-1, dim) reshape fail. - hidden_states = swiglu_oai_triton( - hidden_states, - hidden_states.shape[-1], - self.moe_runner_config.gemm1_alpha, - self.moe_runner_config.gemm1_clamp_limit, - ) - elif self.moe_runner_config.activation == "silu": - if self.moe_runner_config.gemm1_clamp_limit is not None: - from sgl_kernel_npu.activation.swiglu_quant import swiglu_quant - - hidden_states, _ = swiglu_quant( - hidden_states, - group_list=expert_tokens, - group_list_type=1, - need_quant=False, - do_limit=True, - limit=self.moe_runner_config.gemm1_clamp_limit, - ) - else: - hidden_states = torch.ops.npu.npu_swiglu(hidden_states) - else: - from sglang.srt.layers.activation import GeluAndMul - - hidden_states = GeluAndMul()(hidden_states) - - # gmm2: down_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w2_weight.transpose(1, 2)], - bias=w2_bias, - split_item=2, - group_list_type=1, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - - final_hidden_states = torch.ops.npu.npu_moe_finalize_routing( - hidden_states, - skip1=None, - skip2=None, - bias=None, - scales=topk_weights, - expanded_src_to_dst_row=expanded_row_idx, - export_for_source_row=topk_ids, - drop_pad_mode=2, - ) - - return StandardCombineInput(hidden_states=final_hidden_states) - - def _forward_npu_deepep( - self, - layer: torch.nn.Module, - dispatch_output: DispatchOutput, - ) -> CombineInput: - from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( - npu_fused_moe_without_routing_weights_bf16, - ) - from sglang.srt.layers.moe.token_dispatcher import ( - DeepEPLLCombineInput, - DeepEPNormalCombineInput, - ) - from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker - - # NOTE: Ascend's Dispatch & Combine does not support FP16 - output_dtype = torch.bfloat16 - group_list_type = 1 - - if DispatchOutputChecker.format_is_deepep_normal(dispatch_output): - hidden_states, _, _, _, num_recv_tokens_per_expert = dispatch_output - group_list = torch.tensor( - num_recv_tokens_per_expert, - dtype=torch.int64, - device=hidden_states.device, - ) - combine_cls = DeepEPNormalCombineInput - else: - hidden_states, _, _, _, group_list, _ = dispatch_output - group_list = group_list.to(torch.int64) - combine_cls = DeepEPLLCombineInput - - hidden_states = npu_fused_moe_without_routing_weights_bf16( - layer, hidden_states, group_list_type, group_list, output_dtype - ) - return combine_cls( - hidden_states=hidden_states, - topk_ids=dispatch_output.topk_ids, - topk_weights=dispatch_output.topk_weights, - ) + return self.runner.run(dispatch_output, layer) def forward_tpu(self, *args, **kwargs) -> CombineInput: raise NotImplementedError("The TPU backend currently does not support MoE.") diff --git a/python/sglang/srt/models/afmoe.py b/python/sglang/srt/models/afmoe.py index 7b281a32c..72bc43a6a 100644 --- a/python/sglang/srt/models/afmoe.py +++ b/python/sglang/srt/models/afmoe.py @@ -62,7 +62,7 @@ from sglang.srt.utils import add_prefix, is_npu _is_npu = is_npu() if _is_npu: - from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( + from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( fused_moe_npu as fused_moe, ) diff --git a/python/sglang/srt/models/dbrx.py b/python/sglang/srt/models/dbrx.py index 6707fe7c2..d0425e603 100644 --- a/python/sglang/srt/models/dbrx.py +++ b/python/sglang/srt/models/dbrx.py @@ -27,7 +27,7 @@ from sglang.srt.configs import DbrxConfig from sglang.srt.distributed import ( tensor_model_parallel_all_reduce, ) -from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( fused_moe_npu, ) from sglang.srt.layers.linear import ( diff --git a/python/sglang/srt/models/deepseek.py b/python/sglang/srt/models/deepseek.py index 133c0e6f9..cce2d5aee 100644 --- a/python/sglang/srt/models/deepseek.py +++ b/python/sglang/srt/models/deepseek.py @@ -59,7 +59,7 @@ if _is_cpu and _is_cpu_amx_available: import sgl_kernel # noqa: F401 if _is_npu: - from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( + from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( fused_moe_npu as fused_moe, ) else: diff --git a/python/sglang/srt/models/xverse_moe.py b/python/sglang/srt/models/xverse_moe.py index 943395e1c..964a6ba4c 100644 --- a/python/sglang/srt/models/xverse_moe.py +++ b/python/sglang/srt/models/xverse_moe.py @@ -22,7 +22,7 @@ from transformers import PretrainedConfig from sglang.srt.distributed import ( tensor_model_parallel_all_reduce, ) -from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( +from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( fused_moe_npu, ) from sglang.srt.layers.activation import SiluAndMul diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6d3cec64c..ea83e2041 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -274,6 +274,7 @@ MOE_A2A_BACKEND_CHOICES = [ "ascend_fuseep", "flashinfer", "megamoe", + "ascend_tp", ] MXFP8_MOE_RUNNER_BACKEND_CHOICES = [ @@ -1901,6 +1902,7 @@ class ServerArgs: "ascend_fuseep", "flashinfer", "megamoe", + "ascend_tp", ], Arg( help="Choose the backend for MoE A2A.", @@ -1924,6 +1926,10 @@ class ServerArgs: Literal["auto", "normal", "low_latency"], "Select the mode when enable DeepEP or MoriEP MoE, could be `normal`, `low_latency` or `auto`. Default is `auto`, which means `low_latency` for decode batch and `normal` for prefill batch.", ] = "auto" + fuseep_mode: A[ + Literal[1, 2], + "Select the mode when enable Ascend FuseEP MoE, 1 -> dispatch_gmm_combine_decode is executed；2 -> dispatch_ffn_combine is executed (support hybrid deployment when 2).", + ] = 2 deepep_dispatcher_output_dtype: A[ Literal["auto", "bf16", "fp8", "int8", "nvfp4"], "Select DeepEP dispatcher output dtype", @@ -5558,20 +5564,17 @@ class ServerArgs: f"Nixl MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]." ) - if a2a_backend == "ascend_fuseep": + if ( + self.moe_a2a_backend == "none" and is_npu() + ) or self.moe_a2a_backend == "ascend_tp": + # FIXME (OrangeRedeng): for some reasons if pass "ascend_tp" accuracy drops to zero + self.moe_a2a_backend = "none" + + if self.moe_a2a_backend == "ascend_fuseep": logger.warning( f"Ascend fused EP MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]." ) - fuse_mode = envs.SGLANG_NPU_FUSED_MOE_MODE.get() - if fuse_mode not in [1, 2]: - raise ValueError( - f"Wrong value of {fuse_mode=}, the NPU only support 1 or 2." - ) - elif fuse_mode == 2: - assert ( - resolved_view(self).quantization == "modelslim" - ), "When fuse_mode is set to 2, the NPU supports only ModelSlim quantization." - if a2a_backend == "flashinfer": + if self.moe_a2a_backend == "flashinfer": assert ( resolved_view(self).enable_dp_attention and self.dp_size == self.tp_size ), "Flashinfer MoE A2A is only supported with dp_size == tp_size and --enable-dp-attention" @@ -7684,6 +7687,32 @@ def get_global_server_args() -> ServerArgs: return get_context().server_args +def _has_cli_arg(argv: List[str], flag: str) -> bool: + return any(arg == flag or arg.startswith(f"{flag}=") for arg in argv) + + +def _apply_fuseep_mode_env_compat( + raw_args: argparse.Namespace, argv: List[str] +) -> None: + if not envs.SGLANG_NPU_FUSED_MOE_MODE.is_set() or _has_cli_arg( + argv, "--fuseep-mode" + ): + return + + fuseep_mode = envs.SGLANG_NPU_FUSED_MOE_MODE.get() + if fuseep_mode not in (1, 2): + raise ValueError( + f"Wrong value of SGLANG_NPU_FUSED_MOE_MODE={fuseep_mode}, " + "the NPU only supports 1 or 2." + ) + + logger.warning( + "The env variable SGLANG_NPU_FUSED_MOE_MODE is deprecated and will be " + "removed in a future release. Please use --fuseep-mode instead." + ) + raw_args.fuseep_mode = fuseep_mode + + def prepare_server_args(argv: List[str]) -> ServerArgs: """ Prepare the server arguments from the command line arguments. @@ -7718,6 +7747,8 @@ def prepare_server_args(argv: List[str]) -> ServerArgs: force=True, ) + _apply_fuseep_mode_env_compat(raw_args, argv) + return ServerArgs.from_cli_args(raw_args) diff --git a/python/sglang/test/ascend/test_ascend_utils.py b/python/sglang/test/ascend/test_ascend_utils.py index 3e6a9a194..8a70c8776 100644 --- a/python/sglang/test/ascend/test_ascend_utils.py +++ b/python/sglang/test/ascend/test_ascend_utils.py @@ -77,6 +77,9 @@ DOTS_OCR_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "rednote-hilab/dots.ocr" ECO_TECH_QWEN3_32B_W4A4_LAOS_WEIGHTS_PATH = os.path.join( MODEL_WEIGHTS_DIR, "Eco-Tech/Qwen3-32B-w4a4-LAOS" ) +ECO_TECH_QWEN3_30B_A3B_W4A4_LAOS_WEIGHTS_PATH = os.path.join( + MODEL_WEIGHTS_DIR, "Eco-Tech/Qwen3-30B-A3B-w4a4-LAOS" +) ERNIE_4_5_21B_A3B_PT_WEIGHTS_PATH = os.path.join( MODEL_WEIGHTS_DIR, "baidu/ERNIE-4.5-21B-A3B-PT" ) @@ -301,9 +304,6 @@ DEEPSEEK_R1_0528_W4A8_PER_CHANNEL_WEIGHTS_PATH = os.path.join( DEEPSEEK_R1_0528_W8A8_WEIGHTS_PATH = os.path.join( MODEL_WEIGHTS_DIR, "vllm-ascend/DeepSeek-R1-0528-W8A8" ) -QWEN3_30B_MODELSLIM_INT4_WEIGHTS_PATH = os.path.join( - MODEL_WEIGHTS_DIR, "Eco-Tech/Qwen3-30B-A3B-w4a4-LAOS" -) QWEN3_5_397B_W4A8_MODEL_PATH = os.path.join( MODEL_WEIGHTS_DIR, "Eco-Tech/Qwen3.5-397B-A17B-w4a8-mtp" ) diff --git a/test/registered/ascend/llm_models/test_npu_qwen3_30b_fuseep.py b/test/registered/ascend/llm_models/test_npu_qwen3_30b_fuseep.py new file mode 100644 index 000000000..693e2b3fb --- /dev/null +++ b/test/registered/ascend/llm_models/test_npu_qwen3_30b_fuseep.py @@ -0,0 +1,45 @@ +import os +import unittest + +from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin +from sglang.test.ascend.test_ascend_utils import QWEN3_30B_A3B_W8A8_WEIGHTS_PATH +from sglang.test.ci.ci_register import register_npu_ci +from sglang.test.test_utils import CustomTestCase + +register_npu_ci(est_time=400, suite="per-commit-4-npu-a3") + + +class TestQwen330Bw8a8FuseEP(GSM8KAscendMixin, CustomTestCase): + """Testcase: Verify that the inference accuracy of the Qwen/Qwen3-30B-A3B-w8a8 model on the GSM8K dataset is no less than 0.90. + + [Test Category] Model + [Test Target] Qwen/Qwen3-30B-A3B-w8a8 + """ + + model = QWEN3_30B_A3B_W8A8_WEIGHTS_PATH + accuracy = 0.90 + other_args = [ + "--trust-remote-code", + "--mem-fraction-static", + 0.8, + "--attention-backend", + "ascend", + "--cuda-graph-max-bs-decode", + 128, + "--tp-size", + 4, + "--moe-a2a-backend", + "ascend_fuseep", + "--fuseep-mode", + 2, + ] + + env = { + **os.environ, + "SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES": "100", + "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "188416", + } + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/ascend/llm_models/test_npu_qwen3_30b_w4a4.py b/test/registered/ascend/llm_models/test_npu_qwen3_30b_w4a4.py index 1b4dae5d7..3911fedd2 100644 --- a/test/registered/ascend/llm_models/test_npu_qwen3_30b_w4a4.py +++ b/test/registered/ascend/llm_models/test_npu_qwen3_30b_w4a4.py @@ -1,28 +1,28 @@ import unittest from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin -from sglang.test.ascend.test_ascend_utils import QWEN3_30B_MODELSLIM_INT4_WEIGHTS_PATH +from sglang.test.ascend.test_ascend_utils import ( + ECO_TECH_QWEN3_30B_A3B_W4A4_LAOS_WEIGHTS_PATH, +) from sglang.test.ci.ci_register import register_npu_ci from sglang.test.test_utils import CustomTestCase register_npu_ci(est_time=400, suite="per-commit-2-npu-a2") -class TestQwen317BGPTQInt8(GSM8KAscendMixin, CustomTestCase): +class TestQwen330Bw4a4(GSM8KAscendMixin, CustomTestCase): """Testcase: Verify that the inference accuracy of the Eco-Tech/Qwen3-30B-A3B-w4a4-LAOS model on the GSM8K dataset is no less than 0.85. [Test Category] Model - [Test Target] Qwen/Qwen3-1.7B-GPTQ-Int8 + [Test Target] Eco-Tech/Qwen3-30B-A3B-w4a4-LAOS """ - model = QWEN3_30B_MODELSLIM_INT4_WEIGHTS_PATH + model = ECO_TECH_QWEN3_30B_A3B_W4A4_LAOS_WEIGHTS_PATH accuracy = 0.85 other_args = [ "--trust-remote-code", "--mem-fraction-static", 0.8, - "--max-running-requests", - 32, "--attention-backend", "ascend", "--disable-cuda-graph", diff --git a/test/registered/ascend/performance/minimax_m2_5/test_npu_minimax_m2_5_w8a8_8p_in3k5_out1k5_50ms_gpqa.py b/test/registered/ascend/performance/minimax_m2_5/test_npu_minimax_m2_5_w8a8_8p_in3k5_out1k5_50ms_gpqa.py index 8d686340f..6a10b1992 100644 --- a/test/registered/ascend/performance/minimax_m2_5/test_npu_minimax_m2_5_w8a8_8p_in3k5_out1k5_50ms_gpqa.py +++ b/test/registered/ascend/performance/minimax_m2_5/test_npu_minimax_m2_5_w8a8_8p_in3k5_out1k5_50ms_gpqa.py @@ -30,7 +30,6 @@ MINIMAX_M2_5_HIGH_THROUGHPUT_ENVS = { "ASCEND_USE_FIA": "1", "SGLANG_SET_CPU_AFFINITY": "1", "SGLANG_ENABLE_OVERLAP_PLAN_STREAM": "1", - "SGLANG_NPU_FUSED_MOE_MODE": "2", "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "204800", "PYTHONPATH": f"{MINIMAX_M2_5_EAGLE3_MODEL_PATH}:{os.environ.get('PYTHONPATH', '')}", "SGLANG_EXTERNAL_MODEL_PACKAGE": "custom_eagle3", @@ -68,8 +67,8 @@ MINIMAX_M2_5_HIGH_THROUGHPUT_OTHER_ARGS = [ 20, "--moe-a2a-backend", "ascend_fuseep", - "--deepep-mode", - "auto", + "--fuseep-mode", + 2, "--quantization", "modelslim", "--speculative-algorithm",